"""Snapshot integration checks with a fake scheduler; real lifecycle remains separate.""" from contextlib import contextmanager import hashlib import importlib.util import shutil from pathlib import Path from types import SimpleNamespace import uuid import pytest if Path(__file__).with_name("managed_trial.py").is_file(): spec = importlib.util.spec_from_file_location("staged_managed_trial", Path(__file__).with_name("managed_trial.py")) managed = importlib.util.module_from_spec(spec) spec.loader.exec_module(managed) else: from hessian_recipes import managed_trial as managed class Checkpoints: def __init__(self, root): self.root, self.reported, self.fail_store = root, {}, False @contextmanager def store_path(self, metadata): storage_id = str(uuid.uuid4()) path = self.root / storage_id path.mkdir(parents=True) yield path, storage_id if self.fail_store: raise RuntimeError("scheduler report failed") self.reported[storage_id] = metadata @contextmanager def restore_path(self, storage_id): assert storage_id in self.reported yield self.root / storage_id class Context: def __init__(self, storage, preempt_at=None): self.checkpoint, self.preempt_at, self.acknowledged = storage, preempt_at, False self.training, self.validation, self.progress, self.completed = [], [], [], [] self.train = SimpleNamespace(report_training_metrics=lambda **kw: self.training.append(kw), report_validation_metrics=lambda **kw: self.validation.append(kw)) self.preempt = SimpleNamespace(should_preempt=self.should_preempt, acknowledge_preemption_signal=self.acknowledge) self.searcher = SimpleNamespace(operations=lambda: iter([ SimpleNamespace(length=8, report_progress=self.progress.append, report_completed=self.completed.append)])) def should_preempt(self, *, auto_ack): assert auto_ack is False return bool(self.preempt_at and self.training and self.training[-1]["steps_completed"] >= self.preempt_at) def acknowledge(self): assert self.checkpoint.reported self.acknowledged = True def info(storage_id=None, *, trial=23, allocation="allocation-one"): return SimpleNamespace(task_type="TRIAL", latest_checkpoint=storage_id, allocation_id=allocation, trial=SimpleNamespace(trial_id=trial, experiment_id=17), container_addrs=["127.0.0.1"], slot_ids=[]) def adapter(tmp_path, context, identity=None, cluster_info=None, clock=lambda: 100): return managed.ManagedTrial(context, cluster_info or info(), identity=identity or {"source_sha256": "source", "config_sha256": "config"}, binding_path=tmp_path / "binding.json", max_wall_seconds=1000, clock=clock) def native_output(path): checkpoint = path / "checkpoint-2" checkpoint.mkdir(parents=True) for name in managed.NATIVE_STATE: (checkpoint / name).write_bytes(("native-" + name).encode()) for name in managed.ROOT_FILES: (path / name).write_bytes(("root-" + name).encode()) for name in managed.PREFIX_FILES: (path / name).write_bytes(b'{"committed":1}\n{"orphan":2}\n') tokenizer = path / "policy-tokenizer" tokenizer.mkdir() (tokenizer / "tokenizer.json").write_text("frozen tokenizer") seal = {"optimizer_step": 2, "files_sha256": {name: managed._digest(checkpoint / name) for name in managed.NATIVE_STATE}, **{key: managed._digest(path / name) for name, key in managed.ROOT_FILES.items()}, **{key: {"bytes": len(b'{"committed":1}\n'), "sha256": hashlib.sha256(b'{"committed":1}\n').hexdigest()} for name, key in managed.PREFIX_FILES.items()}} managed._write_json(checkpoint / managed.SEAL, seal) return checkpoint def test_cpu_pause_restore_finishes_same_reference_with_new_allocation(tmp_path): storage, clock = Checkpoints(tmp_path / "storage"), [100.] first = Context(storage, preempt_at=2) def sleeper(seconds): clock[0] += seconds result = managed.cpu_lifecycle(first, info(), tmp_path / "run", delay=3, clock=lambda: clock[0], sleeper=sleeper) assert result["status"] == "preempted" and result["steps_completed"] == 2 assert first.acknowledged and first.completed == [] original_deadline = result["deadline"] second = Context(storage) result = managed.cpu_lifecycle(second, info(result["scheduler_uuid"], allocation="allocation-two"), tmp_path / "run", delay=3, clock=lambda: clock[0], sleeper=sleeper) assert result["status"] == "completed" and result["restored"] assert result["state"]["allocations"] == ["allocation-one", "allocation-two"] assert result["deadline"] == original_deadline assert [row["steps_completed"] for row in second.training] == list(range(3, 9)) assert second.validation[0]["metrics"]["reference_matches"] == 1 reference = {"steps_completed": 0, "value": 17, "chain_sha256": hashlib.sha256(b"hessian-cpu").hexdigest()} for _ in range(8): reference = managed._advance(reference) assert all(result["state"][key] == value for key, value in reference.items()) def test_native_restore_uses_exact_prefix_and_preserves_partial_output(tmp_path): storage = Checkpoints(tmp_path / "storage") first, output = adapter(tmp_path, Context(storage)), tmp_path / "training" storage_id = first.publish_trainer(output, native_output(output)) original = (output / "rollout-evidence.jsonl").read_bytes() (output / "partial.txt").write_text("preserve") validated = [] def validate(staging, canonical, name): assert canonical == output and name == "checkpoint-2" assert (staging / name / managed.SEAL).is_file() assert (staging / "rollout-evidence.jsonl").read_bytes() == b'{"committed":1}\n' validated.append(True) second = adapter(tmp_path, Context(storage), cluster_info=info(storage_id, allocation="allocation-two")) assert second.restore_trainer(output, validate) == output / "checkpoint-2" assert validated assert (output / "rollout-evidence.jsonl").read_bytes() == b'{"committed":1}\n' quarantine = list((tmp_path / "training.managed-recovery").iterdir())[0] assert (quarantine / "partial.txt").read_text() == "preserve" assert (quarantine / "rollout-evidence.jsonl").read_bytes() == original def test_corrupt_scheduler_payload_cannot_replace_local_state(tmp_path): storage = Checkpoints(tmp_path / "storage") first, output = adapter(tmp_path, Context(storage)), tmp_path / "training" storage_id = first.publish_trainer(output, native_output(output)) (storage.root / storage_id / "payload/checkpoint-2/optimizer.pt").write_bytes(b"corrupt") second = adapter(tmp_path, Context(storage), cluster_info=info(storage_id)) with pytest.raises(ValueError, match="payload hashes"): second.restore_trainer(output, lambda *args: pytest.fail("validator must not run")) assert (output / "checkpoint-2/optimizer.pt").read_bytes() != b"corrupt" assert not (tmp_path / "training.managed-recovery").exists() def test_native_validator_failure_preserves_existing_output(tmp_path): storage = Checkpoints(tmp_path / "storage") first, output = adapter(tmp_path, Context(storage)), tmp_path / "training" storage_id = first.publish_trainer(output, native_output(output)) (output / "partial.txt").write_text("preserve") def reject(*args): raise ValueError("native source changed") second = adapter(tmp_path, Context(storage), cluster_info=info(storage_id)) with pytest.raises(ValueError, match="source changed"): second.restore_trainer(output, reject) assert (output / "partial.txt").read_text() == "preserve" def test_deadline_survives_restarts_before_first_snapshot(tmp_path): storage = Checkpoints(tmp_path / "storage") assert adapter(tmp_path, Context(storage)).binding["deadline"] == 1100 assert adapter(tmp_path, Context(storage), clock=lambda: 500).binding["deadline"] == 1100 with pytest.raises(TimeoutError): adapter(tmp_path, Context(storage), clock=lambda: 1100) def test_absent_local_binding_restores_original_scheduler_deadline(tmp_path): storage = Checkpoints(tmp_path / "storage") first, output = adapter(tmp_path, Context(storage)), tmp_path / "training" storage_id = first.publish_trainer(output, native_output(output)) (tmp_path / "binding.json").unlink() second = adapter(tmp_path, Context(storage), cluster_info=info(storage_id), clock=lambda: 500) assert second.binding is None second.restore_trainer(output, lambda *args: None) assert second.binding["deadline"] == 1100 @pytest.mark.parametrize("change", ["trial", "source", "output", "uuid"]) def test_snapshot_binding_rejects_changed_trial_source_output_or_uuid(tmp_path, change): storage = Checkpoints(tmp_path / "storage") first, output = adapter(tmp_path, Context(storage)), tmp_path / "training" storage_id = first.publish_trainer(output, native_output(output)) if change in {"trial", "source"}: with pytest.raises(ValueError, match="binding|identity"): adapter(tmp_path, Context(storage), identity={"source_sha256": "changed"} if change == "source" else None, cluster_info=info(storage_id, trial=24 if change == "trial" else 23)) return if change == "uuid": path = storage.root / storage_id / "snapshot.json" envelope = managed._read_json(path) envelope["scheduler_uuid"] = str(uuid.uuid4()) managed._write_json(path, envelope) second = adapter(tmp_path, Context(storage), cluster_info=info(storage_id)) with pytest.raises(ValueError, match="snapshot identity|another canonical"): second.restore_trainer(tmp_path / "other" if change == "output" else output, lambda *args: None) def test_preemption_ack_waits_for_successful_current_checkpoint_report(tmp_path): storage, output = Checkpoints(tmp_path / "storage"), tmp_path / "training" context = Context(storage) current = adapter(tmp_path, context) current.preempt_step = 2 with pytest.raises(ValueError, match="Publish"): current.acknowledge_preemption() checkpoint = native_output(output) storage.fail_store = True with pytest.raises(RuntimeError, match="scheduler report failed"): current.publish_trainer(output, checkpoint) with pytest.raises(ValueError, match="Publish"): current.acknowledge_preemption() assert not context.acknowledged storage.fail_store = False current.publish_trainer(output, checkpoint) current.acknowledge_preemption() assert context.acknowledged @pytest.mark.parametrize("change", ["prefix", "native", "symlink"]) def test_incomplete_or_changed_native_state_is_not_reported(tmp_path, change): storage = Checkpoints(tmp_path / "storage") current, output = adapter(tmp_path, Context(storage)), tmp_path / "training" checkpoint = native_output(output) if change == "prefix": (output / "rollout-evidence.jsonl").write_bytes(b"changed prefix") elif change == "native": (checkpoint / "optimizer.pt").unlink() else: (output / "policy-tokenizer/extra").symlink_to(checkpoint / "optimizer.pt") with pytest.raises(ValueError): current.publish_trainer(output, checkpoint) assert storage.reported == {} def test_unsafe_names_are_rejected(): for name in ["../escape", "/absolute", "a/../b", "a//b", "a\\b"]: with pytest.raises(ValueError): managed._relative(name) def test_native_managed_binding_is_snapshotted_and_must_match_this_trial(tmp_path): storage = Checkpoints(tmp_path / "storage") current, output = adapter(tmp_path, Context(storage)), tmp_path / "training" checkpoint = native_output(output) managed._write_json(output / "managed-binding.json", current.binding) seal = managed._read_json(checkpoint / managed.SEAL) seal["managed_binding_sha256"] = managed._digest(output / "managed-binding.json") managed._write_json(checkpoint / managed.SEAL, seal) storage_id = current.publish(output, checkpoint) payload = storage.root / storage_id / "payload" assert managed._read_json(payload / "managed-binding.json") == current.binding changed = dict(current.binding, trial_id=999) managed._write_json(output / "managed-binding.json", changed) with pytest.raises(ValueError, match="managed binding"): current.publish(output, checkpoint) def test_preempt_receipt_is_durable_before_synchronous_acknowledgement(tmp_path): storage = Checkpoints(tmp_path / "storage") context = Context(storage, preempt_at=2) output = tmp_path / "cpu" def immediate_termination(): receipt = managed._read_json(output / "trial-23/status.json") assert receipt["status"] == "preempted" assert receipt["steps_completed"] == 2 assert receipt["scheduler_uuid"] in storage.reported raise SystemExit("allocation terminated after acknowledgement") context.preempt.acknowledge_preemption_signal = immediate_termination with pytest.raises(SystemExit, match="allocation terminated"): managed.cpu_lifecycle(context, info(), output, delay=0, clock=lambda: 100, sleeper=lambda seconds: None) def test_restore_retains_earlier_dev_checkpoints_and_validates_all_before_swap(tmp_path): storage = Checkpoints(tmp_path / "storage") output = tmp_path / "training" checkpoint = native_output(output) later = output / "checkpoint-4" shutil.copytree(checkpoint, later) seal = managed._read_json(later / managed.SEAL) seal["optimizer_step"] = 4 managed._write_json(later / managed.SEAL, seal) current = adapter(tmp_path, Context(storage)) with pytest.raises(ValueError, match="native validator"): current.publish(output, later) published = [] current.native_validator = lambda root, path: published.append(path.name) storage_id = current.publish(output, later) assert published == ["checkpoint-2", "checkpoint-4"] restored = adapter(tmp_path, Context(storage), cluster_info=info(storage_id)) validated = [] result = restored.restore(output, lambda root, canonical, name: validated.append(name)) assert result == output / "checkpoint-4" assert validated == ["checkpoint-2", "checkpoint-4"] assert (output / "checkpoint-2" / managed.SEAL).is_file() assert (output / "checkpoint-4" / managed.SEAL).is_file() @pytest.mark.parametrize("step", [1, 2]) def test_checkpoint_steps_must_increase_before_a_second_core_uuid(tmp_path, step): storage = Checkpoints(tmp_path / "storage") current, output = adapter(tmp_path, Context(storage)), tmp_path / "training" identifier = current.publish(output, native_output(output)) with pytest.raises(ValueError, match="must advance"): current._publish({}, step=step, kind="cpu") assert list(storage.reported) == [identifier] assert len(list(storage.root.iterdir())) == 1 def test_restored_checkpoint_cannot_be_registered_again_at_the_same_step(tmp_path): storage = Checkpoints(tmp_path / "storage") first, output = adapter(tmp_path, Context(storage)), tmp_path / "training" identifier = first.publish(output, native_output(output)) resumed = adapter(tmp_path, Context(storage), cluster_info=info(identifier, allocation="allocation-two")) checkpoint = resumed.restore(output, lambda *args: None) with pytest.raises(ValueError, match="must advance"): resumed.publish(output, checkpoint) assert list(storage.reported) == [identifier] assert len(list(storage.root.iterdir())) == 1