Buckets:
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| import torch | |
| from opd.runtime import ( | |
| latest_complete_checkpoint, | |
| load_opd_checkpoint, | |
| save_opd_checkpoint, | |
| ) | |
| from opd.config import OPDConfig | |
| from opd.trainer import ( | |
| _evaluation_due, | |
| _rollout_sampling_seed, | |
| _scheduler, | |
| _unique_artifact_path, | |
| ) | |
| class _Saver(torch.nn.Linear): | |
| def save_pretrained(self, path, **_kwargs): | |
| torch.save(self.state_dict(), Path(path) / "adapter.pt") | |
| class _Processor: | |
| def save_pretrained(self, path): | |
| (Path(path) / "processor.json").write_text("{}\n") | |
| class _Config: | |
| def fingerprint(self): | |
| return "config-fingerprint" | |
| def as_dict(self): | |
| return {"run": {"name": "fixture"}} | |
| class OPDRuntimeTest(unittest.TestCase): | |
| def test_rollout_sampling_seed_is_stateless_and_phase_scoped(self): | |
| config = OPDConfig() | |
| first = _rollout_sampling_seed( | |
| config, phase="validation", update=0, sample_id="a" | |
| ) | |
| self.assertEqual( | |
| first, | |
| _rollout_sampling_seed( | |
| config, phase="validation", update=0, sample_id="a" | |
| ), | |
| ) | |
| self.assertNotEqual( | |
| first, | |
| _rollout_sampling_seed(config, phase="train", update=0, sample_id="a"), | |
| ) | |
| self.assertNotEqual( | |
| first, | |
| _rollout_sampling_seed( | |
| config, phase="validation", update=1, sample_id="a" | |
| ), | |
| ) | |
| def test_constant_scheduler_only_uses_the_fixed_warmup(self): | |
| config = OPDConfig() | |
| config.scheduler.kind = "constant" | |
| config.scheduler.warmup_updates = 2 | |
| config.trainer.max_updates = 100 | |
| model = torch.nn.Linear(1, 1) | |
| optimizer = torch.optim.SGD(model.parameters(), lr=1.0) | |
| scheduler, warmup = _scheduler(optimizer, config) | |
| scale = scheduler.lr_lambdas[0] | |
| self.assertEqual(warmup, 2) | |
| self.assertEqual([scale(step) for step in (0, 1, 2, 50, 100)], [0.5, 0.5, 1.0, 1.0, 1.0]) | |
| def test_final_validation_is_due_even_off_the_periodic_cadence(self): | |
| config = OPDConfig() | |
| config.trainer.max_updates = 7 | |
| config.evaluation.every_updates = 5 | |
| self.assertFalse(_evaluation_due(config, 6)) | |
| self.assertTrue(_evaluation_due(config, 7)) | |
| config.evaluation.at_end = False | |
| self.assertFalse(_evaluation_due(config, 7)) | |
| def test_retry_artifact_path_preserves_orphaned_attempt_output(self): | |
| with tempfile.TemporaryDirectory() as directory: | |
| original = Path(directory) / "validation_step_000007.json" | |
| original.write_text("{}\n") | |
| retry = _unique_artifact_path(original) | |
| self.assertEqual(retry.name, "validation_step_000007-retry-0001.json") | |
| self.assertTrue(original.is_file()) | |
| def test_latest_checkpoint_ignores_incomplete_staging_and_steps(self): | |
| with tempfile.TemporaryDirectory() as directory: | |
| root = Path(directory) | |
| (root / "step-000003").mkdir() | |
| complete = root / "step-000002" | |
| complete.mkdir() | |
| (complete / "COMPLETE").write_text("complete\n") | |
| (root / ".tmp-step-000004-dead").mkdir() | |
| self.assertEqual(latest_complete_checkpoint(root), complete) | |
| def test_atomic_checkpoint_roundtrip_excludes_rollout_and_failure_queue(self): | |
| model = _Saver(1, 1, bias=False) | |
| optimizer = torch.optim.SGD(model.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _step: 1.0) | |
| state = {"attempted_updates": 3, "successful_updates": 3, "query_update": 3} | |
| ema = {"schema": "shared_ema_teacher_v1", "update_count": 3} | |
| provenance = {"source": "source", "cache": "cache", "model": "model", "verifier": "verifier"} | |
| with tempfile.TemporaryDirectory() as directory: | |
| checkpoint = save_opd_checkpoint( | |
| checkpoints_dir=Path(directory), | |
| model=model, | |
| processor=_Processor(), | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| trainer_state=state, | |
| config=_Config(), | |
| provenance_fingerprints=provenance, | |
| teacher_state=ema, | |
| ) | |
| payload = torch.load(checkpoint / "training_state.pt", weights_only=False) | |
| self.assertEqual(payload["schema"], "opd_checkpoint_v1") | |
| self.assertNotIn("rollouts", payload) | |
| self.assertNotIn("failure_queue", payload) | |
| self.assertTrue((checkpoint / "COMPLETE").exists()) | |
| loaded_teacher = [] | |
| restored = load_opd_checkpoint( | |
| checkpoint, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| config=_Config(), | |
| provenance_fingerprints=provenance, | |
| teacher_state_loader=loaded_teacher.append, | |
| ) | |
| self.assertEqual(restored, state) | |
| self.assertEqual(loaded_teacher, [ema]) | |
| def test_distributed_checkpoint_restores_each_rank_rng(self): | |
| model = _Saver(1, 1, bias=False) | |
| optimizer = torch.optim.SGD(model.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _step: 1.0) | |
| state = {"attempted_updates": 1, "successful_updates": 1, "query_update": 1} | |
| ema = {"schema": "shared_ema_teacher_v1", "update_count": 1} | |
| provenance = {"source": "s", "cache": "c", "model": "m"} | |
| rng_states = [] | |
| for seed in (7, 9): | |
| torch.manual_seed(seed) | |
| from shared.artifacts import rng_state | |
| rng_states.append(rng_state()) | |
| with tempfile.TemporaryDirectory() as directory: | |
| checkpoint = save_opd_checkpoint( | |
| checkpoints_dir=Path(directory), | |
| model=model, | |
| processor=_Processor(), | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| trainer_state=state, | |
| config=_Config(), | |
| provenance_fingerprints=provenance, | |
| teacher_state=ema, | |
| world_size=2, | |
| rank_rng_states=rng_states, | |
| ) | |
| payload = torch.load(checkpoint / "training_state.pt", weights_only=False) | |
| self.assertEqual(payload["schema"], "opd_checkpoint_v2") | |
| self.assertTrue((checkpoint / "rng-rank-00.pt").is_file()) | |
| self.assertTrue((checkpoint / "rng-rank-01.pt").is_file()) | |
| loaded_teacher = [] | |
| restored = load_opd_checkpoint( | |
| checkpoint, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| config=_Config(), | |
| provenance_fingerprints=provenance, | |
| teacher_state_loader=loaded_teacher.append, | |
| rank=1, | |
| world_size=2, | |
| ) | |
| self.assertEqual(restored, state) | |
| self.assertEqual(torch.get_rng_state().tolist(), rng_states[1]["torch_cpu"].tolist()) | |
| if __name__ == "__main__": | |
| unittest.main() | |
Xet Storage Details
- Size:
- 7.22 kB
- Xet hash:
- 247b9713fc3221035708cdcd7a86da04a87a1c9c4fb9376ebbd940cfecc678d4
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.