| import contextlib |
| import io |
| import json |
| import os |
| from pathlib import Path |
| import tempfile |
| import time |
| import unittest |
| from unittest.mock import patch |
|
|
| from checkpoint_archive import Archive, atomic_json, sha256_file, stable_stat |
| import cleanup_run1_run40 as cleanup |
|
|
|
|
| class CleanupGuardTest(unittest.TestCase): |
| def setUp(self): |
| self.tmp = tempfile.TemporaryDirectory() |
| self.base = Path(self.tmp.name).resolve() |
| self.root = self.base / 'checkpoints' |
| self.migration = self.root / 'migration' |
| self.migration.mkdir(parents=True) |
| self.patch = patch.multiple(cleanup, BASE=self.base, ROOT=self.root, MIGRATION=self.migration) |
| self.patch.start() |
| self.source = self.base / 'source' |
| self.source.mkdir() |
| self.weight = self.source / 'svae.pt' |
| self.weight.write_bytes(b'complete final model') |
| old = time.time() - 7200 |
| os.utime(self.weight, (old, old)) |
| self.cache = self.source / 'features_rank0_part00000.pt' |
| self.cache.write_bytes(b'regenerable feature cache') |
| os.utime(self.cache, (old, old)) |
| self.keep = self.source / 'config.yaml' |
| self.keep.write_text('preserve: true') |
| self.archive = Archive(self.root) |
| |
| for i in range(53): |
| self.archive.add_model(key=f'svae/run{i:02d}', name=f'run{i:02d}_test', source=self.source, files=['svae.pt']) |
| with contextlib.redirect_stderr(io.StringIO()): |
| self.archive.verify(reconstruct=True, report=self.migration / 'verification.json', workers=4) |
| atomic_json(self.migration / 'plan.json', {'fixture': True}) |
| rows = [{'path': str(p), 'category': cat, 'archive_id': 'svae/run00', 'stat': stable_stat(p)} |
| for p, cat in [(self.weight, 'archived_svae_original'), (self.cache, 'regenerable_feature_shard')]] |
| atomic_json(self.migration / 'cleanup-plan.json', {'approved_plan_sha256': sha256_file(self.migration / 'plan.json'), |
| 'files': rows, 'remove_if_empty': [], 'preserved': ['configuration']}) |
| atomic_json(self.migration / 'load-checks.json', {'status': 'passed', 'eval_dry_run_exit_code': 0, |
| 'cpu_load_successes': ['synthetic-fixture'], 'catalog_sha256': sha256_file(self.root / 'catalog.json')}) |
|
|
| def tearDown(self): |
| self.patch.stop() |
| self.tmp.cleanup() |
|
|
| def test_only_verified_manifest_files_are_deleted(self): |
| with contextlib.redirect_stdout(io.StringIO()): |
| cleanup.apply_plan() |
| self.assertFalse(self.weight.exists()) |
| self.assertFalse(self.cache.exists()) |
| self.assertEqual(self.keep.read_text(), 'preserve: true') |
| self.assertEqual((self.root / 'svae/run00_test/private/svae.pt').read_bytes(), b'complete final model') |
| report = json.loads((self.migration / 'cleanup-report.json').read_text()) |
| self.assertEqual(report['removed_files'], 2) |
|
|
| def test_failed_validation_prevents_all_deletion(self): |
| atomic_json(self.migration / 'load-checks.json', {'status': 'failed'}) |
| with self.assertRaises(ValueError): |
| cleanup.apply_plan() |
| self.assertTrue(self.weight.exists()) |
| self.assertTrue(self.cache.exists()) |
|
|
| def test_changed_source_prevents_all_deletion(self): |
| self.cache.write_bytes(b'new training in progress') |
| with self.assertRaises(ValueError): |
| cleanup.apply_plan() |
| self.assertTrue(self.weight.exists()) |
| self.assertTrue(self.cache.exists()) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
