diff options
Diffstat (limited to 'kpov_judge/create_disk_images.py')
-rwxr-xr-x | kpov_judge/create_disk_images.py | 245 |
1 files changed, 124 insertions, 121 deletions
diff --git a/kpov_judge/create_disk_images.py b/kpov_judge/create_disk_images.py index ce79f85..381ba25 100755 --- a/kpov_judge/create_disk_images.py +++ b/kpov_judge/create_disk_images.py @@ -23,27 +23,30 @@ def get_prepare_disks(db, class_id, task_id): 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): +def create_snapshot(class_id, task_id, student_id, disk_name, fmt='vmdk', 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) + task_id, disk_name, snap_hash, fmt) backing = [] - template = disk_name + '.' + settings.STUDENT_DISK_FORMAT + template = disk_name + '.' + fmt 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]) + # ensure task dir exists + os.makedirs(task_path, exist_ok=True) - else: - # create task dir - os.makedirs(task_path, exist_ok=True) + if fmt == 'vdi': + # don’t use backing files, just copy the template + os.chdir(task_path) + if settings.STUDENT_DISK_COW: + subprocess.call(['cp', '--reflink=always', os.path.join(settings.DISK_TEMPLATE_PATH, template), snap]) + else: + subprocess.call(['cp', os.path.join(settings.DISK_TEMPLATE_PATH, template), snap]) + elif fmt == 'vmdk': # qemu-img create stores backing-file path as given, so link all # backing images to task directory where target image will be # generated @@ -60,11 +63,93 @@ def create_snapshot(class_id, task_id, student_id, disk_name, overwrite=True): # make overlay image os.chdir(task_path) subprocess.call(['qemu-img', 'create', - '-f', settings.STUDENT_DISK_FORMAT, + '-f', fmt, '-b', template, snap]) return task_dir, snap, backing +def prepare_task_disks(class_id, task_id, student_id, fmt, computers): + disks = collections.defaultdict(dict) + templates = collections.defaultdict(dict) + for computer in computers: + lock_fp.write('creating computer ' + computer['name'] + '\n') + if not computer['disks']: + continue + + manual_disks = [] + try_automount = False + + g = guestfs.GuestFS() + for disk in computer['disks']: + lock_fp.write("register " + disk['name'] + '\n') + task_dir, snap, backing = create_snapshot(class_id, task_id, student_id, disk['name'], fmt=fmt) + 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( + 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)) + + # add disk or update existing record with new format + disks[computer['name']][disk['name']] = [snap] + 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') + for mountpoint, device in sorted(mps): + if mountpoint not in mounted: + try: + g.mount(device, mountpoint, ) + lock_fp.write( 'mounted ' + device + ' on ' + mountpoint + '\n') + 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: + g.mount_options(opts, device, mountpoint) + else: + g.mount(device, mountpoint) + 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 + } + if 'TASK_URL' in vars(settings): + global_params['task_url'] = settings.TASK_URL + '/' + class_id + '/' + + 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) + 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() + + return disks + if __name__ == '__main__': if len(sys.argv) != 1: print("Usage: {0}") @@ -77,119 +162,37 @@ if __name__ == '__main__': all_computers[(computer['class_id'], computer['task_id'], computer['student_id'])] += [computer] 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: + if db.student_computers.find_one({'class_id': class_id, 'task_id': task_id, 'student_id': student_id}) 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') - try: - 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 = {} - all_disks = collections.defaultdict(list) - parts = {} - for computer in computers: - lock_fp.write('creating computer ' + computer['name'] + '\n') - if len(computer['disks']) == 0: + with open(lock_file, 'w') as lock_fp: + try: + fcntl.lockf(lock_fp, fcntl.LOCK_EX | fcntl.LOCK_NB) + except IOError: continue - manual_disks = [] - try_automount = False - - g = guestfs.GuestFS() - for disk in computer['disks']: - lock_fp.write("register " + disk['name'] + '\n') - 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( - 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']] += [{ - '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') - for mountpoint, device in sorted(mps): - if mountpoint not in mounted: - try: - g.mount(device, mountpoint, ) - lock_fp.write( 'mounted ' + device + ' on ' + mountpoint + '\n') - 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: - g.mount_options(opts, device, mountpoint) - else: - g.mount(device, mountpoint) - 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 - } - 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 computer in computers: - comp_name = computer['name'] - disks = all_disks[comp_name] - lock_fp.write('urls: '+ str(disks) + '\n') - l = db.student_computers.update({ - 'disk_urls': {'$exists': False}, - 'student_id': student_id, - 'task_id': task_id, - 'class_id': class_id, - 'name': comp_name}, - {'$set': { 'disk_urls': disks } - }) - - os.unlink(lock_file) - lock_fp.close() + print("Creating {}/{} for {}".format(class_id, task_id, student_id)) + all_disks = collections.defaultdict(dict) + for fmt in settings.STUDENT_DISK_FORMATS: + for computer, disks in prepare_task_disks(class_id, task_id, student_id, fmt, computers).items(): + for disk, urls in disks.items(): + d = all_disks[computer].setdefault(disk, {'formats': []}) + d['formats'] += [fmt] + d[fmt] = urls + + lock_fp.write("saving URLs\n") + for computer in computers: + comp_name = computer['name'] + disks = all_disks[comp_name] + lock_fp.write('urls: '+ str(disks) + '\n') + db.student_computers.update({ + 'disk_urls': {'$exists': False}, + 'student_id': student_id, + 'task_id': task_id, + 'class_id': class_id, + 'name': comp_name}, + {'$set': { 'disk_urls': disks }}) + + os.unlink(lock_file) |