import hashlib import json import math from pathlib import Path import subprocess import tempfile import unittest import torch from scripts import prepare_expansion_backup as backup from scripts.publish_hf_snapshot import prepare as publication_prepare, verify_stage from selection import validation_selection, SELECTION_METRIC class ExpansionBackupTests(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) self.root = Path(self.temp.name) self.source = self.root / 'source' self.source.mkdir() (self.source / 'scripts').mkdir() (self.source / 'training_model.py').write_text('# unchanged implementation\n') (self.source / 'scripts/prepare_expanded_data.py').write_text('# frozen data builder\n') (self.source / 'HF_MODEL_CARD.md').write_text('---\nlicense: unknown\n---\nFixture model card\n') for command in (['init', '-q'], ['add', '.'], ['-c', 'user.name=Fixture', '-c', 'user.email=fixture@example.invalid', 'commit', '-qm', 'fixture']): subprocess.run(['git', *command], cwd=self.source, check=True, capture_output=True) self.revision = subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=self.source, text=True).strip() source_sha = {name: backup.digest(self.source / name) for name in ('training_model.py', 'scripts/prepare_expanded_data.py')} self.old, self.new = self.root / 'v1', self.root / 'v2' for directory in (self.old, self.new): directory.mkdir() for split in backup.SPLITS: (directory / (split + '.jsonl')).write_bytes(b'opaque frozen payload\n') (self.new / 'train.jsonl').write_bytes(b'opaque frozen payload\nexpanded train\n') hashes = lambda directory: {s: backup.digest(directory / (s + '.jsonl')) for s in backup.SPLITS} base = {'split_sha256': hashes(self.old)} backup.write_json(self.old / 'manifest.json', base) (self.new / 'diagnostics').mkdir() (self.new / 'diagnostics/new_sources.jsonl').write_text('opaque diagnostic payload\n') (self.new / 'raw/piqa').mkdir(parents=True) (self.new / 'raw/piqa/README.md').write_text('Fixture upstream license evidence\n') self.data = {'split_sha256': hashes(self.new), 'base_dataset': {'path': str(self.old), 'manifest_sha256': backup.digest(self.old / 'manifest.json'), 'split_sha256': base['split_sha256']}, 'source_script_sha256': source_sha['scripts/prepare_expanded_data.py'], 'diagnostics': {'path': 'diagnostics/new_sources.jsonl', 'selection_eligible': False, 'sha256': backup.digest(self.new / 'diagnostics/new_sources.jsonl')}, 'sources': {'piqa': {'repo': 'fixture/upstream', 'revision': 'pinned', 'license': 'AFL-3.0', 'license_reference': 'README.md', 'files': {'README.md': {'local_path': 'raw/piqa/README.md', 'sha256': backup.digest(self.new / 'raw/piqa/README.md')}}}}} backup.write_json(self.new / 'manifest.json', self.data) proof = self.new.with_name(self.new.name + '-proof') proof.mkdir() for name in ('tokenization-proof.json', 'data_filter.json', 'postbuild-audit.json', 'prepare_expanded_data.py', 'tokenize_only.py'): (proof / name).write_bytes(b'opaque provenance fixture\n') provenance = {'model_id': 'fixture', 'revision': 'pin'} signature = lambda data: hashlib.sha256(json.dumps({'data': data['split_sha256'], 'model': provenance, 'implementation': source_sha['training_model.py'], 'max_tokens': 512}, sort_keys=True).encode()).hexdigest() self.rows = [] for family in ('arc', 'banking', 'boolq', 'snli'): for index in range(128): logits = [float(index % 3), 0.0] logsum = math.log(sum(math.exp(v) for v in logits)) logs = [v - logsum for v in logits] self.rows.append({'id': f'{family}-{index}', 'group': f'{family}-{index // 2}', 'family': family, 'choices': ['a', 'b'], 'target': index % 2, 'logits': logits, 'log_probabilities': logs, 'probabilities': [math.exp(v) for v in logs]}) scores = validation_selection(self.rows) self.parent = self.root / 'parent' self.parent.mkdir() parent_config = {'dataset': str(self.old), 'validation_per_family': 128, 'max_tokens': 512} saved = {'format': 'opensysone-adapter-v1', 'step': 20, 'config': parent_config, 'trainable_state': {'head': torch.ones(2)}, 'source_commit': self.revision, 'data_signature': signature(base), 'model_provenance': provenance, 'selection_metric': SELECTION_METRIC, 'best_validation_selection_score': scores['score'], 'best_validation_macro_nll': scores['raw_macro_nll']} torch.save(saved, self.parent / 'best.pt') self.run = self.root / 'run' self.pilot = self.run / 'artifacts' self.pilot.mkdir(parents=True) (self.run / 'exit_code').write_text('0\n') initial = {'kind': 'warm_start', 'restores_optimizer': False, 'restores_rng': False, 'parent_checkpoint': str(self.parent / 'best.pt'), 'parent_checkpoint_sha256': backup.digest(self.parent / 'best.pt'), 'parent_step': 20, 'parent_source_commit': self.revision, 'data_transition': {'kind': 'training_split_only', 'parent_dataset': str(self.old), 'dataset': str(self.new), 'parent_data_signature': signature(base), 'data_signature': signature(self.data), 'parent_manifest_sha256': backup.digest(self.old / 'manifest.json'), 'manifest_sha256': backup.digest(self.new / 'manifest.json'), 'protected_split_sha256': {s: base['split_sha256'][s] for s in backup.SPLITS[1:]}}} config = {**parent_config, 'dataset': str(self.new)} self.best = {**saved, 'step': 0, 'config': config, 'initialization': initial, 'data_signature': signature(self.data)} torch.save(self.best, self.pilot / 'best.pt') latest = {**self.best, 'step': 8, 'optimizer': {'state': {0: {'step': torch.tensor(8.), 'exp_avg': torch.zeros(2)}}}, 'random_state': (), 'torch_rng': torch.zeros(2, dtype=torch.uint8), 'cuda_rng': []} torch.save(latest, self.pilot / 'checkpoint.pt') self.manifest = {'git_commit': self.revision, 'source_sha256': source_sha, 'config': config, 'data_signature': signature(self.data), 'model_provenance': provenance} backup.write_json(self.pilot / 'manifest.json', self.manifest) backup.write_json(self.parent / 'manifest.json', {**self.manifest, 'config': parent_config, 'data_signature': signature(base)}) self.summary = {'status': 'completed_step_target', 'final_correctness_status': 'passed', 'completed_steps': 8, 'best_sha256': backup.digest(self.pilot / 'best.pt'), 'checkpoint_sha256': backup.digest(self.pilot / 'checkpoint.pt'), **{k: self.best[k] for k in ('best_validation_macro_nll', 'best_validation_selection_score', 'selection_metric')}} backup.write_json(self.pilot / 'summary.json', self.summary) checks = {name + '_probability_max_abs': 0.0 for name in ('branch_chunks_1', 'branch_chunks_2', 'branch_chunks_4', 'question_isolation', 'candidate_permutation', 'repeat')} checks['tolerance_probability_abs'] = 1e-5 for directory in (self.parent, self.pilot): for name, value in [('correctness_initial.json', checks), ('correctness_final.json', checks), ('data_filter.json', {}), ('best_validation_selection.json', scores)]: backup.write_json(directory / name, value) for name in ('parent-snapshot.json', 'source-data-manifest.json'): backup.write_json(self.parent / name, {}) backup.write_json(self.pilot / 'best_validation_predictions.json', self.rows) backup.write_json(self.parent / 'validation_step_000020_predictions.json', self.rows) reference = self.root / 'reference.json' backup.write_json(reference, self.rows) self.plan = self.root / 'plan.json' backup.write_json(self.plan, {'reference_predictions': str(reference), 'reference_sha256': backup.digest(reference)}) self.output = self.root / 'exports/snapshot' def prepare(self): return backup.prepare(self.run, self.output, self.plan, source_root=self.source) def test_complete_backup_preserves_bytes_and_passes_existing_publisher(self): result = self.prepare() self.assertEqual(len(result['artifacts']), 2) self.assertFalse(result['gpu_initialized_during_backup']) self.assertEqual((self.output / 'data/v2/test.jsonl').read_bytes(), (self.new / 'test.jsonl').read_bytes()) self.assertTrue((self.output / 'data/v2/diagnostics/new_sources.jsonl').is_file()) self.assertEqual((self.output / 'data/v2/proof/tokenization-proof.json').read_bytes(), b'opaque provenance fixture\n') self.assertTrue((self.output / 'artifacts/warm-start-parent/best.pt').is_file()) self.assertEqual(backup.read_json(self.output / 'fleet/reference-validation-predictions.json'), self.rows) stage = publication_prepare(self.output, 'andyshu/opensysone', self.source, self.root / 'publication') _, _, proof = verify_stage(stage) self.assertEqual(proof['training_source_commits'], [self.revision]) self.assertFalse(torch.cuda.is_initialized()) def test_incomplete_pilot_rejected_before_any_stage(self): self.summary['status'] = 'interrupted' backup.write_json(self.pilot / 'summary.json', self.summary) with self.assertRaisesRegex(ValueError, 'not complete'): self.prepare() self.assertFalse(self.output.parent.exists()) def test_source_revision_and_parent_corruption_rejected(self): self.manifest['source_sha256']['training_model.py'] = '0' * 64 backup.write_json(self.pilot / 'manifest.json', self.manifest) with self.assertRaisesRegex(ValueError, 'source does not match'): self.prepare() self.manifest['source_sha256']['training_model.py'] = backup.digest(self.source / 'training_model.py') backup.write_json(self.pilot / 'manifest.json', self.manifest) (self.parent / 'best.pt').write_bytes(b'corrupt') with self.assertRaisesRegex(ValueError, 'parent checkpoint checksum'): self.prepare() def test_protected_data_and_diagnostics_mismatch_rejected(self): (self.new / 'test.jsonl').write_bytes(b'changed reserved data') self.data['split_sha256']['test'] = backup.digest(self.new / 'test.jsonl') backup.write_json(self.new / 'manifest.json', self.data) with self.assertRaisesRegex(ValueError, 'protected split'): self.prepare() def test_corrupted_completed_stage_rejected_by_publisher(self): self.prepare() path = self.output / 'data/v2/train.jsonl' path.chmod(0o644) path.write_bytes(b'changed') with self.assertRaisesRegex(ValueError, 'checksum'): backup.verify_artifacts(self.output, 'andyshu/opensysone') if __name__ == '__main__': unittest.main()