source/tests/test_expansion_backup.py

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()