File size: 1,980 Bytes
795f737 | 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 | from pathlib import Path
import tempfile
import unittest
from baim.registry import Registry,file_hash
class RegistryTests(unittest.TestCase):
def test_gates_integrity_and_rollback(self):
with tempfile.TemporaryDirectory() as temp:
root = Path(temp)
source = root/'checkpoint'
source.mkdir()
(source/'model.safetensors').write_bytes(b'unit-test-fixture-not-a-model')
registry = Registry(root/'registry')
try:
registry.register('v000',source,{})
with self.assertRaisesRegex(ValueError,'MISSING_EVIDENCE'):
registry.promote('v000')
self.assertIsNone(registry.champion)
metrics = dict(suite_sha256='a'*64,artifact_sha256=file_hash(source/'model.safetensors'),
task_success_rate=.9,out_of_domain_success=.9,recovery_rate=.9,unsafe_action_rate=0,
premature_finish_rate=0,loop_rate=0,p95_action_ms=1000,peak_service_rss_bytes=1024**3,
target_hardware='linux-epyc9354p-2vcpu',evaluator_version='unit-test-only')
registry.register('v001',source,metrics,parent='v000')
registry.promote('v001')
self.assertEqual(registry.champion,'v001')
registry.register('v002',source,{**metrics,'p95_action_ms':900},parent='v001')
registry.promote('v002')
registry.rollback('v001')
self.assertEqual(registry.champion,'v001')
(root/'registry/v002/model.safetensors').write_bytes(b'changed')
with self.assertRaisesRegex(ValueError,'integrity'):
registry.promote('v002')
self.assertEqual(registry.champion,'v001')
with self.assertRaisesRegex(ValueError,'immutable'):
registry.register('v001',source,metrics)
finally:
registry.close()
|