Buckets:
| import json | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| import torch | |
| from rl.runtime import ( | |
| latest_complete_checkpoint, | |
| load_rl_checkpoint, | |
| save_reference_state, | |
| save_rl_checkpoint, | |
| verify_rl_checkpoint, | |
| ) | |
| from shared.reference import FrozenReference | |
| from shared.artifacts import rng_state | |
| class _Saver(torch.nn.Linear): | |
| def save_pretrained(self, path, **_kwargs): | |
| torch.save(self.state_dict(), Path(path) / "adapter.pt") | |
| (Path(path) / "manifest.json").write_text("adapter-owned\n") | |
| class _Processor: | |
| def save_pretrained(self, path): | |
| (Path(path) / "processor.json").write_text("{}\n") | |
| class _Config: | |
| def fingerprint(self): | |
| return "rl-config" | |
| def as_dict(self): | |
| return { | |
| "run": {"name": "fixture"}, | |
| "checkpoint": {"resume_from": "/prior/checkpoint"}, | |
| } | |
| class RLRuntimeTest(unittest.TestCase): | |
| def test_checkpoint_restores_the_calling_ranks_rng(self): | |
| actor = _Saver(1, 1, bias=False) | |
| reference = FrozenReference(actor) | |
| optimizer = torch.optim.SGD(actor.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0) | |
| states, expected = [], [] | |
| for seed in (42, 71): | |
| torch.manual_seed(seed) | |
| states.append(rng_state()) | |
| expected.append(torch.rand(4)) | |
| with tempfile.TemporaryDirectory() as directory: | |
| checkpoint = save_rl_checkpoint( | |
| checkpoints_dir=Path(directory) / "checkpoints", actor=actor, | |
| processor=_Processor(), optimizer=optimizer, scheduler=scheduler, | |
| trainer_state={"attempted_updates": 0, "successful_updates": 0, "query_update": 0}, | |
| config=_Config(), provenance_fingerprints={}, reference=reference, | |
| world_size=2, rank_rng_states=states, | |
| ) | |
| for rank in (0, 1): | |
| load_rl_checkpoint(checkpoint, optimizer=optimizer, scheduler=scheduler, | |
| config=_Config(), provenance_fingerprints={}, reference=reference, | |
| world_size=2, rank=rank) | |
| torch.testing.assert_close(torch.rand(4), expected[rank], rtol=0, atol=0) | |
| def test_atomic_checkpoint_restores_frozen_run_start_reference_and_cursor(self): | |
| actor = _Saver(1, 1, bias=False) | |
| actor.weight.data.fill_(1.0) | |
| reference = FrozenReference(actor) | |
| optimizer = torch.optim.SGD(actor.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR( | |
| optimizer, lambda _step: 1.0 | |
| ) | |
| state = { | |
| "attempted_updates": 2, | |
| "successful_updates": 2, | |
| "query_update": 2, | |
| } | |
| provenance = {"source": "source", "model": "model"} | |
| with tempfile.TemporaryDirectory() as directory: | |
| run = Path(directory) | |
| save_reference_state(run, reference) | |
| artifact = run / "rollouts/attempt-000002" | |
| artifact.mkdir(parents=True) | |
| (artifact / "update.json").write_text("{}\n") | |
| (artifact / "metrics.json").write_text("{}\n") | |
| actor.weight.data.fill_(5.0) | |
| checkpoint = save_rl_checkpoint( | |
| checkpoints_dir=run / "checkpoints", | |
| actor=actor, | |
| processor=_Processor(), | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| trainer_state=state, | |
| config=_Config(), | |
| provenance_fingerprints=provenance, | |
| reference=reference, | |
| artifact_roots=(artifact,), | |
| ) | |
| self.assertTrue((checkpoint / "COMPLETE").is_file()) | |
| self.assertTrue((checkpoint / "manifest.json").is_file()) | |
| manifest = verify_rl_checkpoint(checkpoint) | |
| self.assertIn("adapter/manifest.json", manifest["files"]) | |
| self.assertEqual(latest_complete_checkpoint(run / "checkpoints"), checkpoint) | |
| payload = torch.load( | |
| checkpoint / "training_state.pt", weights_only=False | |
| ) | |
| self.assertNotIn("trajectories", payload) | |
| self.assertEqual( | |
| payload["reference_descriptor"], reference.descriptor() | |
| ) | |
| self.assertEqual( | |
| payload["committed_artifacts"], | |
| ["rollouts/attempt-000002"], | |
| ) | |
| self.assertEqual(payload["resume_parent"], "/prior/checkpoint") | |
| self.assertIn("resume_from: null", (checkpoint / "config.yaml").read_text()) | |
| marker = json.loads((artifact / "COMMITTED.json").read_text()) | |
| self.assertEqual( | |
| marker["checkpoint"], | |
| "checkpoints/step-000002-attempt-000002", | |
| ) | |
| resumed_reference = FrozenReference(actor) | |
| restored = load_rl_checkpoint( | |
| checkpoint, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| config=_Config(), | |
| provenance_fingerprints=provenance, | |
| reference=resumed_reference, | |
| ) | |
| self.assertEqual(restored, state) | |
| with resumed_reference.apply(actor): | |
| self.assertEqual(float(actor.weight), 1.0) | |
| self.assertEqual(float(actor.weight), 5.0) | |
| (checkpoint / "rng.pt").write_bytes(b"tampered") | |
| with self.assertRaisesRegex(ValueError, "checkpoint file hash"): | |
| verify_rl_checkpoint(checkpoint) | |
| with self.assertRaisesRegex(ValueError, "refusing implicit rollback"): | |
| latest_complete_checkpoint(run / "checkpoints") | |
| retry_artifact = run / "rollouts/attempt-000002-retry-0001" | |
| retry_artifact.mkdir() | |
| (retry_artifact / "update.json").write_text("{}\n") | |
| (retry_artifact / "metrics.json").write_text("{}\n") | |
| retry = save_rl_checkpoint( | |
| checkpoints_dir=run / "checkpoints", | |
| actor=actor, | |
| processor=_Processor(), | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| trainer_state=state, | |
| config=_Config(), | |
| provenance_fingerprints=provenance, | |
| reference=reference, | |
| artifact_roots=(retry_artifact,), | |
| ) | |
| self.assertEqual( | |
| retry.name, | |
| "step-000002-attempt-000002-retry-0001", | |
| ) | |
| self.assertEqual(latest_complete_checkpoint(run / "checkpoints"), retry) | |
| def test_latest_checkpoint_ignores_malformed_complete_directories(self): | |
| with tempfile.TemporaryDirectory() as directory: | |
| root = Path(directory) | |
| malformed = root / "step-not-a-checkpoint" | |
| malformed.mkdir() | |
| (malformed / "COMPLETE").write_text("complete\n") | |
| with self.assertRaises(FileNotFoundError): | |
| latest_complete_checkpoint(root) | |
| malformed_json = root / "step-000000-attempt-000001" | |
| malformed_json.mkdir() | |
| (malformed_json / "COMPLETE").write_text("[]\n") | |
| (malformed_json / "manifest.json").write_text("[]\n") | |
| valid = root / "step-000001-attempt-000003" | |
| valid.mkdir() | |
| (valid / "COMPLETE").write_text("complete\n") | |
| with self.assertRaisesRegex(ValueError, "refusing implicit rollback"): | |
| latest_complete_checkpoint(root) | |
| if __name__ == "__main__": | |
| unittest.main() | |
Xet Storage Details
- Size:
- 7.6 kB
- Xet hash:
- facfb16ed9b180d4d9aa9d4f785b12f019067a42746ea8251685ff119b026d02
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.