#!/usr/bin/env python3 """Continue SQ3.25 Resurface from its exact v2 FP32/Adam/scaler checkpoint. Only the final 3072-additional/4608-cumulative update export is a candidate. The one-update --smoke run is discarded and the formal run reloads the parent. """ from __future__ import annotations import argparse import copy import hashlib import json import math import os from pathlib import Path import sys import time import traceback import torch ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from mamba2_recall import resurface_data as data, resurface_native as native from mamba2_recall import resurface_loss as objective, runtime from mamba2_recall.state_quant import StateQuant from mamba2_recall.state_training import StateQuantTraining from train_quant_first import load_training_inputs, optimizer_for, write_json from evaluate_quant_first import load_candidate as load_parent_candidate PARENT_STEPS = 1536 STEPS = 3072 PROTOCOL_SHA = '4c2c47aa00936ded52cf7b337126f1cce9556da7021e0a75c6e9df83f4949330' PARENT_PROTOCOL_SHA = '24466642ce87c75fc2136a42ed14e69733c462b507a0be82c062b4da7f836bcb' PARENT_REPORT_SHA = 'e5a77d86cf2fb0e2389247e3cb325f74e89957861a6043e92a891d6d402ae359' PARENT_CHECKPOINT_SHA = 'bc548dd427d114098048fa1863f8e602c095dc2d9fde56348795628ae8e2c78f' PARENT_ADAPTER_SHA = '7144b88265a797ef935b1f94845395abec5532d8dc5b2a4a8fa0a8e07ce205f0' CALIBRATION_SHA = 'c366cd577967e64635f1dd960237dc1c0024d7413685ce0b4c4e40da869a7023' FORMAT = 'MAMBA2_SQ_MORE_RESURFACE_TRAIN_V1' CHECKPOINT_FORMAT = 'MAMBA2_SQ_MORE_RESURFACE_CHECKPOINT_V1' INITIAL_ADAPTER = 'exact parent FP32 masters/Adam/GradScaler checkpoint continuation' EXPECTED_SCALER = {'scale': 16.0, 'growth_factor': 2.0, 'backoff_factor': 0.5, 'growth_interval': 2000, '_growth_tracker': 1536} CACHE_BYTES = 28499968 def lr_factor(r): if not 0 <= r < STEPS: raise ValueError('Additional successful update index outside frozen tail') return .01 + .09 * (1 + math.cos(math.pi * r / (STEPS - 1))) / 2 def schedule(): return sum((torch.randperm(PARENT_STEPS, generator=torch.Generator(device='cpu').manual_seed(seed)).tolist() for seed in (2026092804, 2026092805)), []) def code_hashes(): paths = sorted((ROOT/'mamba2_recall').glob('*.py')) paths += [Path(__file__), ROOT/'scripts'/'train_quant_first.py', ROOT/'scripts'/'evaluate_quant_first.py', ROOT/'scripts'/'run_statequant.py', ROOT/'docs'/'QUANT_FIRST_PROTOCOL.md', ROOT/'docs'/'RESURFACE_MORE_PROTOCOL.md'] return {str(p.relative_to(ROOT)): data.sha_file(p) for p in paths} def equal_tree(actual, expected, where='state'): """Verify values and dtypes recursively; CPU/GPU residency may differ.""" if isinstance(expected, torch.Tensor): if (not isinstance(actual, torch.Tensor) or actual.dtype != expected.dtype or actual.shape != expected.shape or not torch.equal(actual.detach().cpu(), expected.detach().cpu())): raise RuntimeError(f'Exact tensor restoration failed: {where}') elif isinstance(expected, dict): if not isinstance(actual, dict) or set(actual) != set(expected): raise RuntimeError(f'Exact dictionary restoration failed: {where}') for key in expected: equal_tree(actual[key], expected[key], f'{where}.{key}') elif isinstance(expected, (tuple, list)): if type(actual) is not type(expected) or len(actual) != len(expected): raise RuntimeError(f'Exact sequence restoration failed: {where}') for i, (a, b) in enumerate(zip(actual, expected)): equal_tree(a, b, f'{where}[{i}]') elif type(actual) is not type(expected) or actual != expected: raise RuntimeError(f'Exact scalar restoration failed: {where}') return True def optimizer_names(bank): return [[n for n in bank.masters if n.endswith(('.V_read', '.g_read'))], [n for n in bank.masters if n.endswith(('.router_w', '.router_b'))]] def check_optimizer(optimizer_state, masters, steps, factor): groups, states = optimizer_state['param_groups'], optimizer_state['state'] names = [[n for n in masters if n.endswith(('.V_read', '.g_read'))], [n for n in masters if n.endswith(('.router_w', '.router_b'))]] if len(groups) != 2 or len(states) != 224 or any(len(ns) != 112 for ns in names): raise ValueError('Wrong resumed Adam parameter inventory') all_ids = [] for group, ns, base_lr in zip(groups, names, (1e-4, 3e-4)): if (len(group['params']) != len(ns) or group['base_lr'] != base_lr or not math.isclose(group['lr'], base_lr * factor, rel_tol=1e-14) or group['betas'] != (.9, .999) or group['eps'] != 1e-8 or group['weight_decay'] != 0. or group['amsgrad'] is not False or group['maximize'] is not False or group['differentiable'] is not False): raise ValueError('Resumed Adam parameter groups differ from frozen recipe') for index, name in zip(group['params'], ns): all_ids.append(index) value = states[index] if set(value) != {'step', 'exp_avg', 'exp_avg_sq'}: raise ValueError('Unexpected Adam parameter state') if value['step'].numel() != 1 or value['step'].item() != steps: raise ValueError('Adam successful step count differs') for field in ('exp_avg', 'exp_avg_sq'): tensor = value[field] if (tensor.dtype != torch.float32 or tensor.shape != masters[name].shape or not bool(torch.isfinite(tensor).all())): raise ValueError('Invalid Adam moments: ' + name) if not bool((value['exp_avg_sq'] >= 0).all()): raise ValueError('Adam variance is negative') if len(set(all_ids)) != 224 or set(all_ids) != set(states): raise ValueError('Adam parameters have missing/duplicate states') return names def load_parent(args, tokenizer, old_binding): if data.sha_file(ROOT/'docs'/'RESURFACE_MORE_PROTOCOL.md') != PROTOCOL_SHA: raise ValueError('Continuation protocol changed') if data.sha_file(ROOT/'docs'/'QUANT_FIRST_PROTOCOL.md') != PARENT_PROTOCOL_SHA: raise ValueError('Parent protocol changed') if (data.sha_file(args.parent_training_report) != PARENT_REPORT_SHA or data.sha_file(args.parent_checkpoint) != PARENT_CHECKPOINT_SHA or data.sha_file(args.calibration) != CALIBRATION_SHA): raise ValueError('Parent report/checkpoint/calibration identity changed') parent_args = argparse.Namespace(calibration=args.calibration, training_report=args.parent_training_report) _, _, parent, adapter_path = load_parent_candidate(parent_args, tokenizer) if parent['binding'] != old_binding or data.sha_file(adapter_path) != PARENT_ADAPTER_SHA: raise ValueError('Parent training binding/adapter differs') for relative, expected in parent['code_sha256'].items(): if data.sha_file(ROOT/relative) != expected: raise ValueError('Parent execution code changed: ' + relative) if parent['checkpoints'][-1]['sha256'] != PARENT_CHECKPOINT_SHA: raise ValueError('Parent report does not identify this final checkpoint') checkpoint = torch.load(args.parent_checkpoint, map_location='cpu', weights_only=True) if (checkpoint['format'] != 'MAMBA2_SQ_FIRST_RESURFACE_CHECKPOINT_V1' or checkpoint['binding'] != parent['binding'] or checkpoint['successful_updates'] != PARENT_STEPS or checkpoint['attempts'] != parent['attempts'] or checkpoint['scaler'] != EXPECTED_SCALER): raise ValueError('Parent checkpoint provenance/update/scaler differs') adapter = native.read_fp16(adapter_path, expected_binding=parent['binding']) if set(checkpoint['masters']) != set(adapter['tensors']) or len(checkpoint['masters']) != 224: raise ValueError('Parent FP32 master inventory differs from FP16 export') for name, value in checkpoint['masters'].items(): if (value.dtype != torch.float32 or not bool(torch.isfinite(value).all()) or not torch.equal(value.half(), adapter['tensors'][name])): raise ValueError('Parent FP32 master cast differs from export: ' + name) check_optimizer(checkpoint['optimizer'], checkpoint['masters'], PARENT_STEPS, .1) return parent, checkpoint, adapter_path def pair_for(r, ordered, examples, windows, prose_order): row = examples[ordered[r]] full = torch.tensor(row['full_ids'], dtype=torch.long, device='cuda')[None] answer_mask = torch.zeros_like(full[:, :-1], dtype=torch.bool) answer_mask[:, row['ce_hidden_positions']] = True if int(answer_mask.sum()) != row['answer_target_count']: raise ValueError('Numeric answer suffix mask differs') j = PARENT_STEPS + r prose_index, prose_start = prose_order[j % 448], 512 * ((j // 448) % 4) prose = windows[prose_index, prose_start:prose_start+512].to('cuda')[None] if tuple(prose.shape) != (1, 512): raise ValueError('Prose segment geometry differs') return {'id': row['id'], 'schedule_entry': ordered[r], 'global_prose_index': j, 'prose_window': prose_index, 'prose_start': prose_start, 'mk_ids': full[:, :-1], 'mk_targets': full[:, 1:], 'answer_mask': answer_mask, 'prose_ids': prose[:, :-1], 'prose_targets': prose[:, 1:]} def expected_scaler_after(before, overflow): after = dict(before) if overflow: # Preserve v2 manual update(new_scale=scale/2): growth tracker is unchanged. after['scale'] *= .5 else: after['_growth_tracker'] += 1 if after['_growth_tracker'] == after['growth_interval']: after['scale'] *= after['growth_factor'] after['_growth_tracker'] = 0 return after def attempt(bank, teacher, teacher_execution, pair, optimizer, scaler, r): optimizer.zero_grad(set_to_none=True) factor = lr_factor(r) for group in optimizer.param_groups: group['lr'] = group['base_lr'] * factor before = dict(scaler.state_dict()) hidden, gates = bank.forward_hidden(pair['mk_ids'], use_checkpoint=True) mk = objective.staged_loss_backward(hidden, pair['mk_targets'], bank.model.lm_head, gates=gates, scaler=scaler, answer_mask=pair['answer_mask'], ce_weight=1., kl_weight=0., closure_weight=0., chunk_tokens=64) del hidden, gates with torch.no_grad(): teacher_hidden = teacher_execution.backbone(pair['prose_ids'], reset=True) hidden, gates = bank.forward_hidden(pair['prose_ids'], use_checkpoint=True) prose = objective.staged_loss_backward(hidden, pair['prose_targets'], bank.model.lm_head, gates=gates, scaler=scaler, teacher_hidden=teacher_hidden, teacher_head=teacher.lm_head, ce_weight=.5, kl_weight=.5, closure_weight=3., chunk_tokens=64) del hidden, gates, teacher_hidden scaler.unscale_(optimizer) gradients = [p.grad for p in bank.parameters()] if any(g is None for g in gradients): raise RuntimeError('Missing adapter gradient') finite = all(bool(torch.isfinite(g).all()) for g in gradients) overflow = (not finite or not mk['scaled_hidden_gradient_finite'] or not prose['scaled_hidden_gradient_finite']) if overflow: scaler.update(new_scale=scaler.get_scale()*.5) magnitude = None else: magnitude = float(torch.nn.utils.clip_grad_norm_(bank.parameters(), 1., error_if_nonfinite=True)) scaler.step(optimizer) scaler.update() if any(not bool(torch.isfinite(p).all()) or not bool(torch.isfinite(p.half()).all()) for p in bank.parameters()): raise FloatingPointError('Nonfinite updated adapter master or FP16 cast') after = dict(scaler.state_dict()) equal_tree(after, expected_scaler_after(before, overflow), 'scaler transition') bank.assert_base_frozen() return {'mk_ce': mk['ce'], 'prose_ce': prose['ce'], 'prose_kl': prose['teacher_to_student_kl'], 'prose_closure': prose['closure'], 'overflow': overflow, 'all_adapter_gradients_finite': finite, 'mk_scaled_hidden_gradient_finite': mk['scaled_hidden_gradient_finite'], 'prose_scaled_hidden_gradient_finite': prose['scaled_hidden_gradient_finite'], 'gradient_norm_before_clip': magnitude, 'loss_scale_before': before['scale'], 'loss_scale': after['scale'], 'scaler_before': before, 'scaler_after': after, 'lr_factor': factor, 'learning_rates': [g['lr'] for g in optimizer.param_groups]} def save_checkpoint(path, bank, optimizer, scaler, binding, success, attempts): if path.exists(): raise FileExistsError(path) temporary = path.with_suffix('.pt.tmp') if temporary.exists(): raise FileExistsError(temporary) masters, optimizer_state = bank.state_dict(), optimizer.state_dict() mapping = check_optimizer(optimizer_state, masters, PARENT_STEPS+success, lr_factor(success-1)) with temporary.open('wb') as stream: torch.save({'format': CHECKPOINT_FORMAT, 'binding': binding, 'additional_successful_updates': success, 'cumulative_successful_updates': PARENT_STEPS+success, 'successful_updates': PARENT_STEPS+success, 'additional_attempts': attempts, 'attempts': attempts, 'masters': masters, 'optimizer': optimizer_state, 'optimizer_parameter_names': mapping, 'scaler': scaler.state_dict()}, stream) stream.flush(); os.fsync(stream.fileno()) temporary.replace(path) return {'path': str(path), 'file': path.name, 'bytes': path.stat().st_size, 'sha256': data.sha_file(path), 'additional_successful_updates': success, 'cumulative_successful_updates': PARENT_STEPS+success} def validate_smoke(path, binding, hashes, output_dir): if path is None: raise ValueError('Formal continuation requires --smoke-report from a discarded update') smoke = json.loads(path.read_text()) if (path.parent.resolve() == output_dir.resolve() or smoke.get('format') != FORMAT or smoke.get('mode') != 'smoke' or smoke.get('complete') is not True or smoke.get('additional_successful_updates') != 1 or smoke.get('cumulative_successful_updates') != PARENT_STEPS+1 or smoke.get('binding') != binding or smoke.get('code_sha256') != hashes or smoke.get('checkpoints') != [] or 'adapter' in smoke or smoke.get('parent_resume_check', {}).get('packed_training_forward_bitwise_equal') is not True or smoke.get('frozen_base_check', {}).get('identity_version_gradients_unchanged') is not True or smoke.get('frozen_state_calibration_check') is not True or not 1 <= smoke.get('additional_attempts', 0) <= 9): raise ValueError('Discarded resumed smoke is incomplete or differs from formal inputs/code') return {'path': str(path), 'sha256': data.sha_file(path), 'discarded': True, 'additional_successful_updates': 1, 'parent_reloaded_for_formal': True} def main(): p = argparse.ArgumentParser(description=__doc__) p.add_argument('--source-dir', type=Path, required=True) p.add_argument('--calibration', type=Path, required=True) p.add_argument('--data-root', type=Path, required=True) p.add_argument('--train-manifest-sha256', required=True) p.add_argument('--prose-manifest', type=Path, required=True) p.add_argument('--prose-tokens', type=Path, required=True) p.add_argument('--parent-training-report', type=Path, required=True) p.add_argument('--parent-checkpoint', type=Path, required=True) p.add_argument('--out-dir', type=Path, required=True) p.add_argument('--smoke-report', type=Path) p.add_argument('--smoke', action='store_true', help='Discard one resumed update; no export/checkpoint') args = p.parse_args() if args.smoke and args.smoke_report is not None: p.error('--smoke-report is only used for the formal continuation') if args.out_dir.exists(): raise FileExistsError('New output directory required') args.out_dir.mkdir(parents=True) torch.set_num_threads(8) torch.manual_seed(2026092804); torch.cuda.manual_seed_all(2026092804) torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False torch.set_float32_matmul_precision('highest') hashes = code_hashes() report = {'format': FORMAT, 'complete': False, 'mode': 'smoke' if args.smoke else 'formal', 'started_unix': time.time(), 'source_sha256': data.sha_file(__file__), 'code_sha256': hashes, 'continuation_protocol_sha256': PROTOCOL_SHA, 'parent_successful_updates': PARENT_STEPS, 'planned_additional_successful_updates': STEPS, 'planned_cumulative_successful_updates': PARENT_STEPS+STEPS, 'history': [], 'checkpoints': []} bank = student_execution = teacher_execution = None try: tokenizer, examples, windows, prose_order, old_binding, permutations = load_training_inputs(args) parent, checkpoint, parent_adapter = load_parent(args, tokenizer, old_binding) binding = dict(old_binding, continuation_protocol_sha256=PROTOCOL_SHA, parent_training_report_sha256=PARENT_REPORT_SHA, parent_checkpoint_sha256=PARENT_CHECKPOINT_SHA, parent_adapter_sha256=PARENT_ADAPTER_SHA, parent_successful_updates=PARENT_STEPS, additional_successful_updates=STEPS, successful_updates=PARENT_STEPS+STEPS, initial_adapter=INITIAL_ADAPTER) report['binding'] = binding report['parent'] = {'training_report': str(args.parent_training_report), 'training_report_sha256': PARENT_REPORT_SHA, 'checkpoint': str(args.parent_checkpoint), 'checkpoint_sha256': PARENT_CHECKPOINT_SHA, 'adapter': str(parent_adapter), 'adapter_sha256': PARENT_ADAPTER_SHA, 'attempts': parent['attempts']} if not args.smoke: report['discarded_smoke'] = validate_smoke(args.smoke_report, binding, hashes, args.out_dir) write_json(args.out_dir/'report.json', report) print('[setup] loading frozen source student and independent S16 teacher', flush=True) student = runtime.load_source_model(args.source_dir) teacher = runtime.load_source_model(args.source_dir) probe_ids = windows[0, :128].cuda()[None] with torch.no_grad(), native.install_fp16(student, parent_adapter, expected_binding=parent['binding']): with StateQuant(student, 'sq3p25', permutations) as packed: parent_hidden = packed.backbone(probe_ids, reset=True) parent_cache = packed.cache_breakdown() if parent_cache['total_bytes'] != CACHE_BYTES: raise RuntimeError('Parent packed cache changed') bank = native.ResurfaceNative(student, 'soft') bank.load_state_dict(checkpoint['masters']) optimizer = optimizer_for(bank) optimizer.load_state_dict(copy.deepcopy(checkpoint['optimizer'])) scaler = torch.amp.GradScaler('cuda', init_scale=1024., growth_factor=2., backoff_factor=.5, growth_interval=2000) scaler.load_state_dict(copy.deepcopy(checkpoint['scaler'])) equal_tree(bank.state_dict(), checkpoint['masters'], 'masters') equal_tree(optimizer.state_dict(), checkpoint['optimizer'], 'optimizer') equal_tree(scaler.state_dict(), checkpoint['scaler'], 'scaler') report['optimizer_parameter_names'] = optimizer_names(bank) student_execution = StateQuantTraining(student, permutations).install() teacher_execution = StateQuant(teacher, 's16').install() with torch.no_grad(): restored_hidden, _ = bank.forward_hidden(probe_ids, use_checkpoint=False) if not torch.equal(parent_hidden, restored_hidden): raise RuntimeError('Restored training forward differs from parent packed FP16 inference') report['parent_resume_check'] = {'masters_exact': True, 'optimizer_exact': True, 'scaler_exact': True, 'master_fp16_cast_equals_parent_export': True, 'optimizer_steps': PARENT_STEPS, 'optimizer_states': 224, 'master_tensors': 224, 'scaler': scaler.state_dict(), 'packed_training_forward_bitwise_equal': True, 'probe_tokens': 128, 'probe_token_sha256': runtime.token_digest(probe_ids.cpu().numpy()), 'cache': parent_cache} del parent_hidden, restored_hidden, probe_ids, checkpoint print('[setup] exact parent masters/Adam/scaler restored; 128-token packed forward matches', flush=True) ordered = schedule() report['schedule'] = {'numeric_seeds': [2026092804, 2026092805], 'numeric_order_sha256_int64le': runtime.token_digest(torch.tensor(ordered, dtype=torch.int64).numpy()), 'prose_global_start': PARENT_STEPS, 'additional_updates': STEPS} successful = attempts = overflows = 0 limit = 1 if args.smoke else STEPS while successful < limit: pair = pair_for(successful, ordered, examples, windows, prose_order) r = successful start = time.monotonic() result = attempt(bank, teacher, teacher_execution, pair, optimizer, scaler, r) student_execution.assert_frozen() attempts += 1 overflows += int(result['overflow']) successful += int(not result['overflow']) row = {'attempt': attempts, 'additional_update_index': r, 'additional_successful_updates': successful, 'cumulative_successful_updates': PARENT_STEPS+successful, 'successful_updates': PARENT_STEPS+successful, 'case_id': pair['id'], 'schedule_entry': pair['schedule_entry'], 'global_prose_index': pair['global_prose_index'], 'prose_window': pair['prose_window'], 'prose_start': pair['prose_start'], 'answer_targets': int(pair['answer_mask'].sum()), 'seconds': time.monotonic()-start, **result} report['history'].append(row) report.update(additional_successful_updates=successful, cumulative_successful_updates=PARENT_STEPS+successful, successful_updates=PARENT_STEPS+successful, additional_attempts=attempts, attempts=attempts, additional_overflows=overflows) if overflows > 8: raise RuntimeError('Additional overflow retry budget exceeded') if not args.smoke and successful and successful % 768 == 0 and not result['overflow']: report['checkpoints'].append(save_checkpoint(args.out_dir/f'checkpoint_{successful:04d}.pt', bank, optimizer, scaler, binding, successful, attempts)) if attempts % 16 == 0 or successful == limit or result['overflow']: report['gpu_memory'] = runtime.gpu_memory_receipt() write_json(args.out_dir/'report.json', report) print(json.dumps({'additional_updates': successful, 'cumulative_updates': PARENT_STEPS+successful, 'attempts': attempts, 'mk_ce': result['mk_ce'], 'prose_ce': result['prose_ce'], 'overflow': result['overflow'], 'loss_scale': result['loss_scale']}), flush=True) report['frozen_base_check'] = bank.assert_base_frozen() report['frozen_state_calibration_check'] = student_execution.assert_frozen() report['teacher_base_parameters_frozen'] = all(not p.requires_grad and p.grad is None for p in teacher.parameters()) report['final_scaler'] = scaler.state_dict() check_optimizer(optimizer.state_dict(), bank.state_dict(), PARENT_STEPS+successful, lr_factor(successful-1)) if not args.smoke: adapter_path = args.out_dir/'adapter_fp16.pt' report['adapter'] = bank.export_fp16(adapter_path, binding=binding) actual_adapter = native.read_fp16(adapter_path, expected_binding=binding) final_checkpoint = torch.load(args.out_dir/'checkpoint_3072.pt', map_location='cpu', weights_only=True) equal_tree(final_checkpoint['masters'], bank.state_dict(), 'final checkpoint masters') equal_tree(final_checkpoint['optimizer'], optimizer.state_dict(), 'final checkpoint optimizer') equal_tree(final_checkpoint['scaler'], scaler.state_dict(), 'final checkpoint scaler') for name, value in final_checkpoint['masters'].items(): if not torch.equal(value.half(), actual_adapter['tensors'][name]): raise RuntimeError('Final checkpoint master cast differs from exported FP16 tensor') report['final_checkpoint_export_check'] = {'all_master_casts_equal_export': True, 'master_tensors': 224, 'optimizer_exact': True, 'scaler_exact': True, 'final_checkpoint_sha256': report['checkpoints'][-1]['sha256']} export_probe = windows[0, :128].cuda()[None] with torch.no_grad(): training_hidden, _ = bank.forward_hidden(export_probe, use_checkpoint=False) student_execution.close(); bank.close() with torch.no_grad(), native.install_fp16(student, adapter_path, expected_binding=binding): with StateQuant(student, 'sq3p25', permutations) as deployed: deployed_hidden = deployed.backbone(export_probe, reset=True) deploy_cache = deployed.cache_breakdown() if not torch.equal(training_hidden, deployed_hidden) or deploy_cache != parent_cache: raise RuntimeError('Final packed training/deployment forward or cache differs') report['deployed_export_check'] = {'packed_training_forward_bitwise_equal': True, 'probe_tokens': 128, 'probe_token_sha256': runtime.token_digest(export_probe.cpu().numpy()), 'cache': deploy_cache, 'cache_unchanged_from_parent': True} report.update(complete=True, gpu_memory=runtime.gpu_memory_receipt()) except BaseException as error: report.update(error=repr(error), traceback=traceback.format_exc()) raise finally: for resource in (student_execution, teacher_execution, bank): if resource is not None: resource.close() report['finished_unix'] = time.time() write_json(args.out_dir/'report.json', report) if __name__ == '__main__': main()