diff options
Diffstat (limited to 'kpov_judge/create_disk_images.py')
| -rwxr-xr-x | kpov_judge/create_disk_images.py | 189 |
1 files changed, 102 insertions, 87 deletions
diff --git a/kpov_judge/create_disk_images.py b/kpov_judge/create_disk_images.py index 9073538..5e630be 100755 --- a/kpov_judge/create_disk_images.py +++ b/kpov_judge/create_disk_images.py @@ -1,62 +1,70 @@ -#!/usr/bin/env python +#!/usr/bin/env python3 -import pymongo -import sys +import hashlib +import collections +import fcntl +import glob import inspect -import kpov_util -import settings -import guestfs import os -import glob +import re import subprocess -import fcntl +import sys + +import guestfs +import pymongo + +import settings +import kpov_util from util import write_default_config def get_prepare_disks(db, class_id, task_id): - prepare_disks_source = db.prepare_disks.find_one({'class_id': class_id, 'task_id':task_id})['source'] + prepare_disks_source = db.prepare_disks.find_one({'class_id': class_id, 'task_id': task_id})['source'] d = {} exec(compile(prepare_disks_source, 'prepare_disks.py', 'exec'), globals(), d) return d['prepare_disks'] -def create_snapshot(class_id, task_id, student_id, disk_name, overwrite = True, cow = False): - print(os.path.join(settings.DISK_TEMPLATE_PATH, disk_name) + '.*') - template_paths = glob.glob(os.path.join(settings.DISK_TEMPLATE_PATH, disk_name) + '.*') - filtered_paths = list(filter((lambda x: os.path.splitext(x)[1] == '.' + settings.STUDENT_DISK_FORMAT), template_paths)) - if filtered_paths: - template_path = filtered_paths[0] - else: - template_path = template_paths[0] - if cow: - d = os.path.join(student_id, class_id, task_id, disk_name) + os.path.splitext(template_path)[1] - else: - d = os.path.join(student_id, class_id, task_id, disk_name) + '.qcow2' - try: - os.makedirs(os.path.join(settings.STUDENT_DISK_PATH, student_id, class_id, task_id)) - except: - pass - disk_file = os.path.join(settings.STUDENT_DISK_PATH, d) - if overwrite or not os.path.exists(disk_file): - if cow: - subprocess.call(['cp', '--reflink=always', template_path, disk_file]) +def create_snapshot(class_id, task_id, student_id, disk_name, overwrite=True): + # add a hash to filename to allow multiple students using the same directory + snap_hash = hashlib.sha1((disk_name+class_id+task_id+student_id).encode()).hexdigest()[:4] + snap = '{}-{}-{}.{}'.format( + task_id, disk_name, snap_hash, settings.STUDENT_DISK_FORMAT) + backing = [] + + template = disk_name + '.' + settings.STUDENT_DISK_FORMAT + task_dir = os.path.join(student_id, class_id, task_id) + task_path = os.path.join(settings.STUDENT_DISK_PATH, task_dir) + + if not os.path.exists(os.path.join(task_path)) or overwrite: + if settings.STUDENT_DISK_COW: + # don’t use backing files, just copy the template + # (requires a cow-capable filesystem) + subprocess.call(['cp', '--reflink=always', template, snap]) + else: - subprocess.call(['qemu-img', 'create', '-f', 'qcow2', '-o', 'backing_file=' + template_path, - disk_file]) - return d - -def publish_snapshot(d): - if os.path.splitext(d)[1][1:] != settings.STUDENT_DISK_FORMAT: - snap_name = os.path.splitext(d)[0] + '.' + settings.STUDENT_DISK_FORMAT - disk_file = os.path.join(settings.STUDENT_DISK_PATH, d) - snap_file = os.path.join(settings.STUDENT_DISK_PATH, snap_name) - subprocess.call(['qemu-img', 'convert', '-f', 'qcow2', '-O', settings.STUDENT_DISK_FORMAT, disk_file, - snap_file]) - url = settings.STUDENT_DISK_URL + snap_name - else: - url = settings.STUDENT_DISK_URL + d - return url - - - + # create task dir + os.makedirs(task_path, exist_ok=True) + + # qemu-img create stores backing-file path as given, so link all + # backing images to task directory where target image will be + # generated + os.chdir(settings.DISK_TEMPLATE_PATH) # qemu-img info is saner when called from image directory + output = subprocess.check_output( + ['qemu-img', 'info', '--backing-chain', template], universal_newlines=True) + for image in [template] + [m.group(1) for m in re.finditer(r'backing file: (.*)', output)]: + backing += [image] + dest = os.path.join(task_path, image) + if not os.path.exists(dest): + os.symlink(os.path.join(settings.DISK_TEMPLATE_PATH, image), dest) + # would be great if someone finds a way to avoid the stuff above + + # make overlay image + os.chdir(task_path) + subprocess.call(['qemu-img', 'create', + '-f', settings.STUDENT_DISK_FORMAT, + '-b', template, snap]) + + return task_dir, snap, backing + if __name__ == '__main__': if len(sys.argv) != 1: print("Usage: {0}") @@ -64,20 +72,18 @@ if __name__ == '__main__': db = pymongo.MongoClient(settings.DB_URI).get_default_database() - dev_prefix = settings.GUESTFS_DEV_PREFIX - l = db.student_computers.find({"disk_urls": {"$exists": False}}) - computers_by_class_task_student = dict() - for computer in l: - student_id, task_id, class_id = computer['student_id'], computer['task_id'], computer['class_id'] - if (class_id, task_id, student_id) not in computers_by_class_task_student: - computers_by_class_task_student[(class_id, task_id, student_id)] = list() - computers_by_class_task_student[(class_id, task_id, student_id)].append(computer) + all_computers = collections.defaultdict(list) + for computer in db.student_computers.find({"disk_urls": {"$exists": False}}): + all_computers[(computer['class_id'], computer['task_id'], computer['student_id'])] += [computer] - for (class_id, task_id, student_id), computers in list(computers_by_class_task_student.items()): - print("Creating {}/{} for {}".format(class_id, task_id, student_id)) - l = db.student_computers.find_one({'class_id': class_id, 'task_id': task_id, 'student_id':student_id, "disk_urls": {"$exists": False}}) + for (class_id, task_id, student_id), computers in all_computers.items(): + # TODO check why we iterate over student_computers twice + l = db.student_computers.find_one({'class_id': class_id, 'task_id': task_id, 'student_id': student_id}) if l is None: continue + + print("Creating {}/{} for {}".format(class_id, task_id, student_id)) + lock_file = os.path.join(settings.STUDENT_LOCKFILE_PATH, '{0}-{1}-{2}.lock'.format(student_id, class_id, task_id)) lock_fp = open(lock_file, 'w') @@ -85,48 +91,56 @@ if __name__ == '__main__': fcntl.lockf(lock_fp, fcntl.LOCK_EX | fcntl.LOCK_NB) except IOError: continue + task_params = db.task_params.find_one({'class_id': class_id, 'task_id': task_id, 'student_id': student_id})['params'] prepare_disks = get_prepare_disks(db, class_id, task_id) + # tule odpri, ustvari snapshote za vajo - templates = dict() - all_disks = dict() - parts = dict() + templates = {} + all_disks = collections.defaultdict(list) + parts = {} for computer in computers: lock_fp.write('creating computer ' + computer['name'] + '\n') - all_disks[computer['name']] = dict() - manual_disks = list() - this_computers_disks = set() - try_automount = False if len(computer['disks']) == 0: continue + + manual_disks = [] + try_automount = False + g = guestfs.GuestFS() for disk in computer['disks']: lock_fp.write("register " + disk['name'] + '\n') - snap = create_snapshot(class_id, task_id, student_id, disk['name'], - cow = settings.STUDENT_DISK_COW) - snap_file = os.path.join(settings.STUDENT_DISK_PATH, snap) + task_dir, snap, backing = create_snapshot(class_id, task_id, student_id, disk['name']) + snap_file = os.path.join(settings.STUDENT_DISK_PATH, task_dir, snap) if 'options' in disk: g.add_drive_opts(snap_file, **(disk['options'])) else: g.add_drive(snap_file) if 'parts' in disk: for p in disk['parts']: - lock_fp.write("part {}: {}\n".format(dev_prefix + p['dev'], p['path'])) - manual_disks.append((dev_prefix + p['dev'], - p['path'], p.get('options', None))) + lock_fp.write("part {}: {}\n".format( + settings.GUESTFS_DEV_PREFIX + p['dev'], p['path'])) + manual_disks.append( + (settings.GUESTFS_DEV_PREFIX + p['dev'], p['path'], p.get('options', None))) else: try_automount = True + templates[disk['name']] = g - lock_fp.write(" templates[{}] = {}\n".format( - disk['name'], disk)) - all_disks[computer['name']][disk['name']] = snap + lock_fp.write(" templates[{}] = {}\n".format(disk['name'], disk)) + + all_disks[computer['name']] += [{ + 'name': disk['name'], + 'file': snap, + 'backing': backing, + }] + g.launch() mounted = set() if try_automount: roots = g.inspect_os() for root in roots: mps = g.inspect_get_mountpoints(root) - lock_fp.write("detected:"+str(mps)+'\n') + lock_fp.write('detected: ' + str(mps) + '\n') for mountpoint, device in sorted(mps): if mountpoint not in mounted: try: @@ -135,6 +149,7 @@ if __name__ == '__main__': except RuntimeError as msg: lock_fp.write( "%s (ignored)\n" % msg) mounted.add(mountpoint) + for device, mountpoint, opts in manual_disks: try: if opts is not None: @@ -144,35 +159,35 @@ if __name__ == '__main__': lock_fp.write('manually mounted ' + device + " on " + mountpoint + '\n') except RuntimeError as msg: lock_fp.write( "%s (ignored)\n" % msg) + lock_fp.write("preparing disks\n") global_params = { 'task_name': task_id, 'class_id': class_id, - 'username': student_id} + 'username': student_id + } if 'TASK_URL' in vars(settings): global_params['task_url'] = settings.TASK_URL + '/' + class_id + '/' + prepare_disks(templates, task_params, global_params) + # pospravi za seboj. lock_fp.write("unmounting\n") for g in set(templates.values()): g.umount_all() g.close() + lock_fp.write("saving URLs\n") - for comp_name, d_dict in all_disks.items(): - disk_urls = list() - for d_name, d in d_dict.items(): - lock_fp.write('publishing '+ str(d) + '\n') - url = publish_snapshot(d) - lock_fp.write('published as '+ url + '\n') - disk_urls.append({'name': d_name, 'url': url}) - lock_fp.write('urls: '+ str(disk_urls) + '\n') + for comp_name, disks in all_disks.items(): + lock_fp.write('urls: '+ str(disks) + '\n') l = db.student_computers.update({ - "disk_urls": {"$exists": False}, + 'disk_urls': {'$exists': False}, 'student_id': student_id, 'task_id': task_id, 'class_id': class_id, 'name': comp_name}, - {'$set': { 'disk_urls': disk_urls }}) - # print "done for ", student_id, task_id + {'$set': { 'disk_urls': disks } + }) + os.unlink(lock_file) lock_fp.close() |