"""Managed lifecycle hooks preserve all native checkpoint/audit boundaries.""" from __future__ import annotations import copy import json from pathlib import Path import random import pytest from hessian_recipes import gym_training_long as long from hessian_recipes.training import sha256 import test_gym_training_long as scripted setup_run = scripted.setup_run URL = scripted.URL class Reporter: def __init__(self, *, preempt_at=3): self.binding = {"format_version": "hessian-managed-snapshot-v1", "experiment_id": 11, "trial_id": 12, "identity": {"config": "fixture"}, "started_at": 100, "deadline": 200, "max_wall_seconds": 100} self.preempt_at = preempt_at self.events = [] self.fail_publish = False self.fail_report = False self.published_step = None def report(self, step, metrics): self.events.append(("report", step, dict(metrics))) if self.fail_report: raise RuntimeError("fixture reporter failed") def preempt_requested(self, step): self.events.append(("preempt_poll", step)) return step == self.preempt_at def publish(self, output, checkpoint): proof = long.validate_checkpoint(output, checkpoint) step = proof["optimizer_step"] assert proof["managed_binding_sha256"] == sha256(output / "managed-binding.json") self.events.append(("publish", step)) if self.fail_publish: raise RuntimeError("fixture publication failed") self.published_step = step def acknowledge_preemption(self): assert self.published_step == self.preempt_at self.events.append(("ack", self.preempt_at)) @pytest.fixture def managed(setup_run, monkeypatch): original = long._managed_callback monkeypatch.setattr(long, "_managed_callback", lambda context: original(context, callback_base=object)) return setup_run, Reporter() def test_offinterval_preemption_defers_publication_until_cleanup_and_external_ack(managed): setup, reporter = managed result = long.run(setup.config(steps=8), URL, checkpoint_interval=4, managed_reporter=reporter) assert result["status"] == "preempted_at_checkpoint" and result["optimizer_step"] == 3 checkpoint = Path(result["checkpoint"]) seal = long.validate_checkpoint(setup.root / "run", checkpoint, expected_step=3) assert seal["checkpoint_reason"] == "managed_preemption" assert seal["managed_binding_sha256"] == sha256(setup.root / "run/managed-binding.json") assert [event[1] for event in reporter.events if event[0] == "report"] == [1, 2, 3] assert all(environment._closed for environment in setup.state.environments) manifest = json.loads((setup.root / "run/training-manifest.json").read_text()) assert manifest["status"] == "preempted_at_checkpoint" and manifest["evidence_sha256"] assert not any(event[0] in {"publish", "ack"} for event in reporter.events) reporter.publish(setup.root / "run", checkpoint) # Outer controller after cleanup. reporter.acknowledge_preemption() assert reporter.events[-2:] == [("publish", 3), ("ack", 3)] def test_managed_resume_matches_uninterrupted_optimizer_rng_and_evidence(managed): setup, reporter = managed full = long.run(setup.config("full", steps=8), URL, checkpoint_interval=4) config = setup.config("managed", steps=8) first = long.run(config, URL, checkpoint_interval=4, managed_reporter=reporter) assert not any(event[0] == "publish" for event in reporter.events) reporter.publish(setup.root / "managed", Path(first["checkpoint"])) resumed = long.run(config, URL, checkpoint_interval=4, managed_reporter=reporter, resume_from_checkpoint=first["checkpoint"]) assert resumed["status"] == "completed" and resumed["optimizer_step"] == 8 assert (Path(full["checkpoint"]) / "adapter_model.safetensors").read_bytes() == ( Path(resumed["checkpoint"]) / "adapter_model.safetensors").read_bytes() assert (setup.root / "full/rollout-evidence.jsonl").read_bytes() == ( setup.root / "managed/rollout-evidence.jsonl").read_bytes() assert [event[1] for event in reporter.events if event[0] == "publish"] == [3, 4] reporter.publish(setup.root / "managed", Path(resumed["checkpoint"])) assert [event[1] for event in reporter.events if event[0] == "publish"] == [3, 4, 8] @pytest.mark.parametrize("mutation", ["no_reporter", "wrong_trial", "changed_binding", "symlink_binding"]) def test_managed_output_cannot_optimize_without_matching_immutable_reporter(managed, mutation): setup, reporter = managed config = setup.config(steps=8) first = long.run(config, URL, checkpoint_interval=4, managed_reporter=reporter) before = setup.state.heavy_calls binding = setup.root / "run/managed-binding.json" if mutation == "wrong_trial": reporter.binding = reporter.binding | {"trial_id": 99} elif mutation == "changed_binding": binding.write_text(json.dumps(reporter.binding | {"trial_id": 99})) elif mutation == "symlink_binding": stored = setup.root / "binding-copy.json" binding.rename(stored) binding.symlink_to(stored) with pytest.raises(ValueError, match="managed|Managed"): long.run(config, URL, checkpoint_interval=4, managed_reporter=None if mutation == "no_reporter" else reporter, resume_from_checkpoint=first["checkpoint"]) assert setup.state.heavy_calls == before @pytest.mark.parametrize("file", ["optimizer.pt", "rng_state.pth", "trainer_state.json"]) def test_managed_preemption_does_not_allow_corrupt_native_checkpoint(managed, file): setup, reporter = managed config = setup.config(steps=8) first = long.run(config, URL, checkpoint_interval=4, managed_reporter=reporter) (Path(first["checkpoint"]) / file).write_bytes(b"corrupt") before = setup.state.heavy_calls with pytest.raises(ValueError, match="hashes"): long.run(config, URL, checkpoint_interval=4, managed_reporter=reporter, resume_from_checkpoint=first["checkpoint"]) assert setup.state.heavy_calls == before def test_missing_reason_retains_strict_default_scheduled_boundary(managed): setup, reporter = managed result = long.run(setup.config(steps=8), URL, checkpoint_interval=4, managed_reporter=reporter) checkpoint = Path(result["checkpoint"]) seal_path = checkpoint / long.SEAL seal = json.loads(seal_path.read_text()) seal.pop("checkpoint_reason") seal_path.write_text(json.dumps(seal)) with pytest.raises(ValueError, match="trainer state"): long.validate_checkpoint(setup.root / "run", checkpoint) def test_unmanaged_reason_cannot_enable_offinterval_resume(managed): setup, reporter = managed result = long.run(setup.config(steps=8), URL, checkpoint_interval=4, managed_reporter=reporter) output = setup.root / "run" checkpoint = Path(result["checkpoint"]) plan_path = output / "run-plan.json" frozen = json.loads(plan_path.read_text()) frozen.pop("managed_training") plan_path.write_text(json.dumps(frozen)) (output / "managed-binding.json").unlink() seal_path = checkpoint / long.SEAL seal = json.loads(seal_path.read_text()) seal["run_plan_sha256"] = sha256(plan_path) seal.pop("managed_binding_sha256") seal_path.write_text(json.dumps(seal)) with pytest.raises(ValueError, match="boundary reason"): long.validate_checkpoint(output, checkpoint) @pytest.mark.parametrize("failure", ["partial_save", "publish", "cleanup", "report"]) def test_failed_native_save_publish_report_or_cleanup_cannot_ack(managed, failure): setup, reporter = managed if failure == "partial_save": setup.state.mode = "partial_save" elif failure == "publish": reporter.preempt_at = None # Fail an intermediate publication at step four. reporter.fail_publish = True elif failure == "cleanup": setup.state.cleanup_failure = True else: reporter.fail_report = True with pytest.raises((OSError, RuntimeError)): long.run(setup.config(steps=8), URL, checkpoint_interval=4, managed_reporter=reporter) assert not any(event[0] == "ack" for event in reporter.events) assert all(environment._closed for environment in setup.state.environments) manifest = json.loads((setup.root / "run/training-manifest.json").read_text()) assert manifest["status"] == "failed" if failure in {"partial_save", "report"}: assert not any(event[0] == "publish" for event in reporter.events) assert not (setup.root / "run/checkpoint-3" / long.SEAL).exists() def test_callback_changes_no_rng_and_only_forces_save_at_optimizer_boundary(): reporter = Reporter(preempt_at=3) context = type("Context", (), {"managed_reporter": reporter, "report": {"training": {"max_steps": 8}}, "preemption_step": None})() callback = long._managed_callback(context, callback_base=object) state = type("State", (), {"global_step": 3})() control = type("Control", (), {"should_save": False, "should_training_stop": False})() before = copy.deepcopy(random.getstate()) assert callback.on_step_end(None, state, control) is control assert callback.on_log(None, state, control, logs={"loss": 0.1}) is control assert random.getstate() == before assert context.preemption_step == 3 and control.should_save and control.should_training_stop assert not any(event[0] == "publish" or event[0] == "ack" for event in reporter.events) def test_existing_scheduled_seal_without_optional_fields_still_validates(setup_run): setup = setup_run result = long.run(setup.config(steps=8), URL, checkpoint_interval=4, stop_after_step=4) checkpoint = Path(result["checkpoint"]) seal_path = checkpoint / long.SEAL seal = json.loads(seal_path.read_text()) seal.pop("checkpoint_reason") assert "managed_binding_sha256" not in seal seal_path.write_text(json.dumps(seal)) assert long.validate_checkpoint(setup.root / "run", checkpoint, expected_step=4) def test_real_upstream_callback_handler_retains_forced_save_and_stop(tmp_path): pytest.importorskip("torch") transformers = pytest.importorskip("transformers") from transformers.trainer_callback import CallbackHandler, DefaultFlowCallback, TrainerControl, TrainerState reporter = Reporter() context = type("Context", (), {"managed_reporter": reporter, "report": {"training": {"max_steps": 8}}, "preemption_step": None})() args = transformers.TrainingArguments(output_dir=str(tmp_path), use_cpu=True, report_to=[], logging_steps=1, save_steps=4, max_steps=8) state = TrainerState(global_step=3, max_steps=8, logging_steps=1, save_steps=4) handler = CallbackHandler([DefaultFlowCallback(), long._managed_callback(context)], model=None, processing_class=None, optimizer=None, lr_scheduler=None) control = handler.on_step_end(args, state, TrainerControl()) control = handler.on_log(args, state, control, logs={"loss": 0.25}) assert control.should_save and control.should_training_stop assert context.preemption_step == 3 assert reporter.events == [("preempt_poll", 3), ("report", 3, {"loss": 0.25})] def test_periodic_64_128_saves_publish_while_final_256_waits_for_outer_cleanup(managed): setup, reporter = managed reporter.preempt_at = None result = long.run(setup.config(steps=256), URL, checkpoint_interval=64, managed_reporter=reporter) assert result["status"] == "completed" and result["optimizer_step"] == 256 assert [event[1] for event in reporter.events if event[0] == "publish"] == [64, 128, 192] assert all(environment._closed for environment in setup.state.environments) manifest = json.loads((setup.root / "run/training-manifest.json").read_text()) assert manifest["status"] == "completed" reporter.publish(setup.root / "run", Path(result["checkpoint"])) assert [event[1] for event in reporter.events if event[0] == "publish"] == [64, 128, 192, 256]