File size: 6,585 Bytes
e7aa2d6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | import concurrent.futures
import json
from pathlib import Path
import struct
import tempfile
import unittest
from checkpoint_archive import Archive, checked_path, sha256_file
def source(root, run, reason=b'12345678', step=10):
folder = root / ('source-' + run)
folder.mkdir()
tensors = {
'action_backbone.test': {'dtype': 'F32', 'shape': [2], 'data_offsets': [0, 8]},
'video_backbone.reason1.test': {'dtype': 'BF16', 'shape': [4], 'data_offsets': [8, 16]},
'vlm_expert.vlm.visual.test': {'dtype': 'BF16', 'shape': [2], 'data_offsets': [16, 20]},
}
# Non-canonical whitespace/metadata must survive byte-identical restoration.
tensors['__metadata__'] = {'description': 'original header preserved'}
raw = json.dumps(tensors, indent=1).encode()
path = folder / f'checkpoint_step_{step}.safetensors'
path.write_bytes(struct.pack('<Q', len(raw)) + raw + b'abcdefgh' + reason + b'WXYZ')
(folder / 'config.yaml').write_text('model: {}\n')
(folder / 'normalization_stats.npy').write_bytes(b'stats-' + run.encode())
(folder / 'reason1').mkdir()
(folder / 'reason1/config.json').write_text('{}')
return path
class ArchiveTest(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.base = Path(self.temp.name)
self.a = Archive(self.base / 'archive')
self.a.init()
def tearDown(self):
self.temp.cleanup()
def add(self, run, **kwargs):
src = source(self.base, run, **kwargs)
return src, self.a.add_policy(run=run, name=run + '_test', stage='sft', source=src)
def test_lossless_relocation_and_real_copies(self):
src, e = self.add('run1')
self.a.verify(reconstruct=True)
relocated = self.base / 'moved'
self.a.root.rename(relocated)
new = Archive(relocated)
output = new.materialize('run1', self.base / 'restored')
self.assertEqual(src.read_bytes(), (output / src.name).read_bytes())
self.assertEqual((src.parent / 'config.yaml').read_bytes(), (output / 'config.yaml').read_bytes())
self.assertNotEqual(src.stat().st_ino, (output / src.name).stat().st_ino)
self.assertFalse(any(p.is_symlink() for p in relocated.rglob('*')))
def test_shared_exact_match_and_changed_component(self):
self.add('run1')
self.add('run2')
self.assertEqual(len(list((self.a.root / 'shared/reason1').glob('*/model.safetensors'))), 1)
self.add('run3', reason=b'changed!')
self.assertEqual(len(list((self.a.root / 'shared/reason1').glob('*/model.safetensors'))), 2)
self.a.verify(reconstruct=True)
def test_incremental_idempotence_and_conflict(self):
src, before = self.add('run40')
self.a.add_policy(run='run40', name='run40_test', stage='sft', source=src)
self.add('run41')
self.assertEqual(before, self.a.resolve('run40'))
changed = source(self.base, 'different', reason=b'new data')
with self.assertRaises(FileExistsError):
self.a.add_policy(run='run40', name='run40_test', stage='sft', source=changed)
(src.parent / 'config.yaml').write_text('changed: true')
with self.assertRaises(FileExistsError):
self.a.add_policy(run='run40', name='run40_test', stage='sft', source=src)
def test_corruption_never_publishes_materialized_output(self):
src, e = self.add('run1')
original = sha256_file(src)
record = self.a.index(e)['files'][0]
dest = self.a.root / record['path']
raw = bytearray(dest.read_bytes()); raw[-1] ^= 1; dest.write_bytes(raw)
with self.assertRaises(ValueError):
self.a.materialize('run1', self.base / 'restore')
self.assertFalse((self.base / 'restore').exists())
self.assertEqual(sha256_file(src), original)
def test_incomplete_cache_marker_cannot_hide_missing_weights(self):
src, _ = self.add('run1')
out = self.a.materialize('run1', self.base / 'restore')
marker = out / '.archive-materialized.json'
data = json.loads(marker.read_text())
data['files'] = []
marker.write_text(json.dumps(data))
(out / src.name).unlink()
with self.assertRaises(ValueError):
self.a.materialize('run1', out)
def test_added_metadata_is_not_silently_ignored_on_reimport(self):
src, _ = self.add('run1')
(src.parent / 'new-config.json').write_text('{}')
with self.assertRaises(FileExistsError):
self.a.add_policy(run='run1', name='run1_test', stage='sft', source=src)
def test_missing_dependency_never_enters_catalog(self):
src = source(self.base, 'run1')
with self.assertRaises(ValueError):
self.a.add_policy(run='run1', name='run1_test', stage='sft', source=src, dependencies={'svae': 'svae/run99'})
self.assertEqual(self.a.catalog()['entries'], {})
self.assertTrue(src.exists())
def test_portable_model_and_dependency(self):
src = self.base / 'svae'; src.mkdir(); (src / 'svae.pt').write_bytes(b'complete model')
self.a.add_model(key='svae/run1', name='run1_svae', source=src, files=['svae.pt'])
ckpt = source(self.base, 'run1')
self.a.add_policy(run='run1', name='run1_test', stage='sft', source=ckpt, dependencies={'svae': 'svae/run1'})
self.a.verify(reconstruct=True)
out = self.a.materialize('svae/run1', self.base / 'restored-model')
self.assertEqual((out / 'svae.pt').read_bytes(), b'complete model')
def test_path_and_symlink_guards(self):
with self.assertRaises(ValueError):
checked_path(self.a.root, '../other')
with self.assertRaises(ValueError):
checked_path(self.a.root, '/absolute')
(self.a.root / 'link').symlink_to(self.base)
with self.assertRaises(ValueError):
checked_path(self.a.root, 'link/anything')
def test_concurrent_append_preserves_both_catalog_entries(self):
sources = [source(self.base, 'run1'), source(self.base, 'run2')]
def add(pair):
i, src = pair
return Archive(self.a.root).add_policy(run=f'run{i}', name=f'run{i}_test', stage='sft', source=src)
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
list(pool.map(add, enumerate(sources, 1)))
self.assertEqual(set(self.a.catalog()['entries']), {'sft/run01', 'sft/run02'})
self.a.verify(reconstruct=True)
if __name__ == '__main__':
unittest.main()
|