"""CPU orchestration/checkpoint tests with an explicit scripted trainer fixture. These do not qualify real TRL/GPU resume equivalence. The predeclared GPU qualification must compare actual sampled evidence and adapter tensors separately. """ from __future__ import annotations import copy import json import math from pathlib import Path import pickle import random from types import SimpleNamespace import pytest from hessian_recipes import gym_training_long as long from hessian_recipes.agent_rl_tasks import GermanAgentEnv, generate_tasks, score_episode from hessian_recipes.evaluation import checkpoint_sha256 from hessian_recipes.training import sha256 class Tensor: requires_grad = True shape = () def __init__(self, value=0.0): self.value = value def detach(self): return self def cpu(self): return self def clone(self): return Tensor(self.value) def float(self): return self def __sub__(self, other): return Tensor(self.value - other.value) def abs(self): return Tensor(abs(self.value)) def max(self): return self.value class Model: def __init__(self): self.weight = Tensor() self.config = SimpleNamespace(use_cache=True) self.generation_config = SimpleNamespace(eos_token_id=7) def named_parameters(self): return [("adapter.weight", self.weight)] class Tokenizer: eos_token_id = 7 def save_pretrained(self, path): Path(path).mkdir(parents=True, exist_ok=True) (Path(path) / "tokenizer.json").write_text("fixture tokenizer") def save(value, path): Path(path).write_bytes(pickle.dumps(value)) def load(path, **kwargs): assert kwargs.get("weights_only") is True return pickle.loads(Path(path).read_bytes()) @pytest.fixture def setup_run(tmp_path, monkeypatch): base = tmp_path / "base" base.mkdir() (base / "model.safetensors").write_bytes(b"fixture base identity") data = tmp_path / "train.jsonl" rows = generate_tasks("train", 3, seed=877) data.write_text("".join(json.dumps(row) + "\n" for row in rows)) state = SimpleNamespace(heavy_calls=0, environments=[], mode="normal", cleanup_failure=False, trainer_args=[], identity={"backend": "nemo_gym", "fixture": True}) def config(name="run", steps=4): path = tmp_path / f"{name}.json" path.write_text(json.dumps({"schema_version": 1, "model": {"path": str(base), "sha256": checkpoint_sha256(base)}, "dataset": {"path": str(data), "sha256": sha256(data)}, "output_dir": str(tmp_path / name), "seed": 877, "training": {"max_steps": steps, "learning_rate": 5e-6, "beta": 0.1}})) return path class Remote: def __init__(self, base_url, identity): self._task = None self._calls = [] self._env = GermanAgentEnv() self._provenance = copy.deepcopy(identity) self._failed = None self._closed = False state.environments.append(self) def reset(self, task=None, **row): self._task, self._calls = copy.deepcopy(task), [] self._env.reset(task=task) def lookup(self, record_id: str) -> str: if state.mode == "transport": self._failed = RuntimeError("fixture Gym infrastructure failure") raise self._failed result = self._env.lookup(record_id) self._calls.append({"tool": "lookup", "arguments": {"record_id": record_id}, "observation": result}) return result def calculator(self, expression: str) -> str: result = self._env.calculator(expression) self._calls.append({"tool": "calculator", "arguments": {"expression": expression}, "observation": result}) return result def _score(self, completion): self._raise_if_failed() return score_episode(self._env, completion) def _raise_if_failed(self): if self._failed: raise self._failed def _close(self): self._closed = True if state.cleanup_failure: raise RuntimeError("fixture cleanup failure") class Trainer: def __init__(self, *, model, args, environment_factory, reward_funcs, train_dataset, **kwargs): self.model, self.args, self.reward, self.dataset = model, args, reward_funcs, train_dataset self.environments = [environment_factory() for _ in range(4)] self.state = SimpleNamespace(global_step=0, max_steps=args.max_steps, log_history=[]) self.control = SimpleNamespace(should_training_stop=False, should_save=False) self.callbacks = kwargs.get("callbacks") or [] self.ref_model = None self.momentum = 0.0 state.trainer_args.append(args) def _generate_single_turn(self, prompts): # Scripted sampled-ID fixture; not real tokenization/model inference. return [[1]] * 4, [[random.randrange(2, 500)] for _ in range(4)], None, {} def _generate(self, prompts): sampled = self._generate_single_turn(prompts) completions = [] for index, environment in enumerate(self.environments): task = environment._task messages = [] def call(name, arguments): messages.append({"role": "assistant", "content": "", "tool_calls": [ {"function": {"name": name, "arguments": json.dumps(arguments)}}]}) try: observation = getattr(environment, name)(**arguments) except RuntimeError: observation = "native fixture swallowed infrastructure error" messages.append({"role": "tool", "name": name, "content": observation}) return observation call("lookup", {"record_id": task["record_id"]}) answer = 0 if index % 2 == 0 and state.mode != "transport": facts = task["private_facts"] answer = json.loads(call("calculator", {"expression": f"{facts['quantity']}*{facts['unit_price']}"}))["result"] messages.append({"role": "assistant", "content": f"{answer}"}) completions.append(messages) return sampled[0], sampled[1], [[1]] * 4, completions, 4, None, {} def train(self, resume_from_checkpoint=None): if resume_from_checkpoint: checkpoint = Path(resume_from_checkpoint) self.model.weight.value = json.loads((checkpoint / "adapter_model.safetensors").read_text())["weight"] self.momentum = pickle.loads((checkpoint / "optimizer.pt").read_bytes())["momentum"] scheduler = pickle.loads((checkpoint / "scheduler.pt").read_bytes()) assert scheduler["horizon"] == self.args.max_steps previous = json.loads((checkpoint / "trainer_state.json").read_text()) self.state = SimpleNamespace(**previous) random.setstate(pickle.loads((checkpoint / "rng_state.pth").read_bytes())) for step in range(self.state.global_step, self.args.max_steps): row = self.dataset[step] for environment in self.environments: environment.reset(task=row["task"]) if state.mode == "wrong_task": for environment in self.environments: environment._task["task_id"] = "wrong-train-task" generated = self._generate([]) self.reward(generated[3], self.environments, task=[row["task"]] * 4) # A deterministic scalar momentum fixture exercises whether the # checkpoint restores optimizer AND global RNG, not just weights. gradient = random.uniform(0.1, 1.0) self.momentum = 0.9 * self.momentum + gradient self.model.weight.value += self.args.learning_rate * self.momentum self.state.global_step = step + 1 self.state.log_history.append({"step": step + 1, "grad_norm": gradient, "kl": 0.001}) for callback in self.callbacks: callback.on_step_end(self.args, self.state, self.control) for callback in self.callbacks: callback.on_log(self.args, self.state, self.control, logs=self.state.log_history[-1]) if (self.control.should_save or (step + 1) % self.args.save_steps == 0 or step + 1 == self.args.max_steps): self._save_checkpoint(self.model, None) self.control.should_save = False if self.control.should_training_stop: break return SimpleNamespace(metrics={"fixture": True}) def _save_checkpoint(self, model, trial): path = Path(self.args.output_dir) / f"checkpoint-{self.state.global_step}" path.mkdir() (path / "adapter_model.safetensors").write_text(json.dumps({"weight": model.weight.value})) if state.mode == "partial_save": raise OSError("fixture interrupted optimizer save") (path / "adapter_config.json").write_text("{}") save({"momentum": self.momentum}, path / "optimizer.pt") save({"step": self.state.global_step, "horizon": self.args.max_steps}, path / "scheduler.pt") save(random.getstate(), path / "rng_state.pth") (path / "trainer_state.json").write_text(json.dumps(vars(self.state))) def load_model(*args, **kwargs): state.heavy_calls += 1 random.random() # Construction must not perturb resumed generation RNG. assert kwargs["local_files_only"] and not kwargs["trust_remote_code"] return Model() runtime = SimpleNamespace( torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: True, is_bf16_supported=lambda: True, device_count=lambda: 1, get_device_name=lambda: "CPU mocked GPU"), bfloat16="bf16", isfinite=lambda tensor: SimpleNamespace(all=lambda: math.isfinite(tensor.value)), save=save, load=load), Dataset=SimpleNamespace(from_list=lambda rows: rows), LoraConfig=lambda **kwargs: kwargs, Model=SimpleNamespace(from_pretrained=load_model), Tokenizer=SimpleNamespace(from_pretrained=lambda *a, **k: Tokenizer()), set_seed=random.seed, Config=lambda **kwargs: SimpleNamespace(**kwargs), Trainer=Trainer, ) monkeypatch.setattr(long.legacy, "_runtime", lambda: runtime) monkeypatch.setattr(long.legacy, "_gym_api", lambda: SimpleNamespace( fetch_gym_identity=lambda url: copy.deepcopy(state.identity), RemoteGymEnv=Remote, GymInfrastructureError=RuntimeError)) monkeypatch.setattr(long, "require_versions", lambda: {"fixture": True}) monkeypatch.setattr(long, "_response_parser_identity", lambda: { "version": "2026.9.29", "source_sha256": {"fixture-engine": "0" * 64}}) monkeypatch.setattr(long.importlib.metadata, "version", lambda name: "1.0.1") monkeypatch.setattr(long.legacy.native, "configure_policy_tokenizer", lambda tokenizer: {"fixture": True}) return SimpleNamespace(state=state, config=config, root=tmp_path, data=data, model=base) URL = "http://127.0.0.1:9999" def test_plan_freezes_schedule_without_loading_ml(setup_run): setup = setup_run report = long.plan(setup.config(steps=256), URL) assert report["training"]["beta"] == 0.1 and report["training"]["learning_rate"] == 5e-6 assert len(report["task_schedule"]) == 256 assert report["fixed"]["shuffle_dataset"] is False assert report["checkpointing"]["selection_steps"] == [64, 128, 256] assert setup.state.heavy_calls == 0 and not (setup.root / "run").exists() def test_stop_resume_preserves_optimizer_rng_schedule_and_rollout_evidence(setup_run): setup = setup_run full = long.run(setup.config("full"), URL, checkpoint_interval=2) config = setup.config("segmented") first = long.run(config, URL, checkpoint_interval=2, stop_after_step=2) assert first["status"] == "stopped_at_checkpoint" assert long.validate_checkpoint(setup.root / "segmented", first["checkpoint"], expected_step=2) resumed = long.run(config, "http://127.0.0.1:9998", checkpoint_interval=2, resume_from_checkpoint=first["checkpoint"]) assert resumed["status"] == "completed" and resumed["optimizer_step"] == 4 full_weight = json.loads((Path(full["checkpoint"]) / "adapter_model.safetensors").read_text()) resumed_weight = json.loads((Path(resumed["checkpoint"]) / "adapter_model.safetensors").read_text()) assert full_weight == resumed_weight assert (setup.root / "full/rollout-evidence.jsonl").read_bytes() == (setup.root / "segmented/rollout-evidence.jsonl").read_bytes() assert resumed["learning_evidence"]["since_previous_checkpoint"]["changed_trainable_tensors"] == 1 assert resumed["learning_evidence"]["cumulative"]["nonzero_gradient_steps"] == 4 assert resumed["task_coverage"]["total_rollouts"] == 16 assert all(not args.shuffle_dataset and not args.save_only_model and not args.ignore_data_skip and args.max_steps == 4 and args.lr_scheduler_type == "linear" for args in setup.state.trainer_args) def test_orphan_evidence_and_partial_future_checkpoint_are_archived(setup_run): setup = setup_run config = setup.config() first = long.run(config, URL, checkpoint_interval=2, stop_after_step=2) output = setup.root / "run" with (output / "rollout-evidence.jsonl").open("ab") as stream: stream.write(b'{"unfinished":') (output / "checkpoint-4").mkdir() (output / "checkpoint-4/optimizer.pt").write_bytes(b"partial optimizer") assert long.discover_resume_checkpoint(output) == output / "checkpoint-2" result = long.run(config, URL, checkpoint_interval=2, resume_from_checkpoint=first["checkpoint"]) recovery = Path(result["recovery"]["path"]) assert (recovery / "uncommitted-rollout-suffix.jsonl").read_bytes() == b'{"unfinished":' assert (recovery / "checkpoint-4/optimizer.pt").read_bytes() == b"partial optimizer" assert long.discover_resume_checkpoint(output) == output / "checkpoint-4" @pytest.mark.parametrize("filename", ["optimizer.pt", "scheduler.pt", "rng_state.pth", "adapter_model.safetensors"]) def test_tampered_resume_payload_fails_before_model_load(setup_run, filename): setup = setup_run config = setup.config() first = long.run(config, URL, checkpoint_interval=2, stop_after_step=2) (Path(first["checkpoint"]) / filename).write_bytes(b"changed") before = setup.state.heavy_calls with pytest.raises(ValueError, match="hashes"): long.run(config, URL, checkpoint_interval=2, resume_from_checkpoint=first["checkpoint"]) assert setup.state.heavy_calls == before def test_modified_evidence_prefix_is_not_silently_recovered(setup_run): setup = setup_run config = setup.config() first = long.run(config, URL, checkpoint_interval=2, stop_after_step=2) evidence = setup.root / "run/rollout-evidence.jsonl" evidence.write_bytes(evidence.read_bytes().replace(b'"rollout_index": 0', b'"rollout_index": 9', 1)) with pytest.raises(ValueError, match="prefix hash mismatch"): long.run(config, URL, checkpoint_interval=2, resume_from_checkpoint=first["checkpoint"]) def test_upstream_source_change_rejected_before_recovery_mutation(setup_run, monkeypatch): setup = setup_run config = setup.config() first = long.run(config, URL, checkpoint_interval=2, stop_after_step=2) original = long._runtime_identity monkeypatch.setattr(long, "_runtime_identity", lambda runtime, installed: original(runtime, installed) | {"changed": True}) with pytest.raises(ValueError, match="upstream trainer source"): long.run(config, URL, checkpoint_interval=2, resume_from_checkpoint=first["checkpoint"]) assert not (setup.root / "run/recovery").exists() @pytest.mark.parametrize("change", ["config", "data", "gym_identity", "checkpoint_interval"]) def test_resume_protocol_changes_fail_before_model_load(setup_run, change): setup = setup_run config = setup.config() first = long.run(config, URL, checkpoint_interval=2, stop_after_step=2) if change == "config": value = json.loads(config.read_text()) value["training"]["beta"] = 0.2 config.write_text(json.dumps(value)) elif change == "data": setup.data.write_text(setup.data.read_text() + "\n") elif change == "gym_identity": setup.state.identity["fixture"] = False before = setup.state.heavy_calls with pytest.raises(ValueError): long.run(config, URL, checkpoint_interval=1 if change == "checkpoint_interval" else 2, resume_from_checkpoint=first["checkpoint"]) assert setup.state.heavy_calls == before def test_incomplete_upstream_save_never_gets_completion_seal(setup_run): setup = setup_run setup.state.mode = "partial_save" with pytest.raises(OSError, match="interrupted optimizer"): long.run(setup.config(), URL, checkpoint_interval=2) assert not (setup.root / "run/checkpoint-2" / long.SEAL).exists() assert long.discover_resume_checkpoint(setup.root / "run") is None assert json.loads((setup.root / "run/training-manifest.json").read_text())["status"] == "failed" @pytest.mark.parametrize("mode", ["transport", "wrong_task"]) def test_native_swallowed_rpc_or_wrong_group_never_becomes_training_zero(setup_run, mode): setup = setup_run setup.state.mode = mode with pytest.raises((RuntimeError, ValueError)): long.run(setup.config(), URL, checkpoint_interval=2) assert not (setup.root / "run/rollout-evidence.jsonl").exists() assert all(environment._closed for environment in setup.state.environments) def test_cleanup_cannot_mask_primary_failure(setup_run): setup = setup_run setup.state.mode, setup.state.cleanup_failure = "transport", True with pytest.raises(RuntimeError, match="infrastructure"): long.run(setup.config(), URL, checkpoint_interval=2) result = json.loads((setup.root / "run/training-manifest.json").read_text()) assert "infrastructure" in result["error"]["message"] and result["cleanup_errors"] def test_native_error_diagnostics_do_not_penalize_unexecuted_cap_call(): completion = [ {"role": "assistant", "tool_calls": [{"function": {"name": "bogus", "arguments": {}}}]}, {"role": "tool", "name": "bogus", "content": "{'error': 'Tool bogus not found.'}"}, {"role": "assistant", "tool_calls": [{"function": {"name": "lookup", "arguments": {"oops": "bad"}}}]}, ] result = long.invalid_tool_diagnostics(completion, [], {"trace": []}) assert result["declared_calls"] == 2 and result["declared_unknown_names"] == 1 assert result["declared_invalid_arguments"] == 1 and result["native_only_tool_results"] == 1 assert result["authoritative_tool_errors"] == 0 and result["reward_unchanged"] @pytest.mark.parametrize("norm", [float("inf"), float("nan"), -1.0]) def test_invalid_gradient_norms_cannot_qualify_learning_evidence(norm): with pytest.raises(ValueError, match="gradient norm"): long._learning_evidence([], [{"step": 1, "grad_norm": norm}]) def test_native_horizon_mismatch_is_not_a_resumable_checkpoint(setup_run): setup = setup_run config = setup.config() first = long.run(config, URL, checkpoint_interval=2, stop_after_step=2) checkpoint = Path(first["checkpoint"]) state_file = checkpoint / "trainer_state.json" value = json.loads(state_file.read_text()) value["max_steps"] = 2 state_file.write_text(json.dumps(value)) # Recompute the ordinary payload digest to isolate the horizon semantic # guard rather than only exercising the already-tested hash mismatch. seal = json.loads((checkpoint / long.SEAL).read_text()) seal["files_sha256"] = long._file_hashes(checkpoint) (checkpoint / long.SEAL).write_text(json.dumps(seal)) with pytest.raises(ValueError, match="trainer state"): long.validate_checkpoint(setup.root / "run", checkpoint) def test_schedule_shuffle_is_local_rng_and_repeats_only_after_exhaustion(): rows = generate_tasks("train", 3, seed=877) random.seed(7) before = random.getstate() schedule = long.task_schedule(rows, 8, 877) assert random.getstate() == before and schedule == long.task_schedule(rows, 8, 877) assert len({item["row_index"] for item in schedule[:3]}) == 3 assert len({item["row_index"] for item in schedule[3:6]}) == 3 def test_real_upstream_optimizer_scheduler_rng_roundtrip_without_model_inference(tmp_path): """Exercise pinned HF safe torch.load on CPU scalars, without an LM/forward.""" torch = pytest.importorskip("torch") transformers = pytest.importorskip("transformers") if torch.__version__.split("+")[0] != "2.10.0" or transformers.__version__ != "5.3.0": pytest.skip("this compatibility check targets the pinned training environment") import numpy as np from transformers import Trainer from transformers.training_args import ParallelMode def shell(parameter): trainer = object.__new__(Trainer) trainer.args = SimpleNamespace(world_size=1, process_index=0, parallel_mode=ParallelMode.NOT_PARALLEL, device=torch.device("cpu"), should_save=True) trainer.is_deepspeed_enabled = trainer.is_fsdp_enabled = trainer.is_fsdp_xla_v1_enabled = False trainer.optimizer = torch.optim.AdamW([parameter], lr=5e-6) trainer.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(trainer.optimizer, lambda step: max(0, 1 - step / 4)) return trainer def update(trainer, parameter): draw = (random.random(), float(np.random.random()), float(torch.rand(()))) parameter.grad = torch.tensor([sum(draw)], dtype=parameter.dtype) trainer.optimizer.step() trainer.lr_scheduler.step() trainer.optimizer.zero_grad() return draw random.seed(877) np.random.seed(877) torch.manual_seed(877) parameter = torch.nn.Parameter(torch.tensor([0.25])) original = shell(parameter) for _ in range(2): update(original, parameter) original._save_optimizer_and_scheduler(str(tmp_path)) original._save_rng_state(str(tmp_path)) saved_parameter = parameter.detach().clone() expected_draws = [update(original, parameter) for _ in range(2)] expected_parameter = parameter.detach().clone() restored_parameter = torch.nn.Parameter(saved_parameter) restored = shell(restored_parameter) random.seed(999) np.random.seed(999) torch.manual_seed(999) # These actual upstream methods use weights_only=True; the RNG method also # uses Transformers' narrowly scoped NumPy safe_globals compatibility list. restored._load_optimizer_and_scheduler(str(tmp_path)) restored._load_rng_state(str(tmp_path)) assert [update(restored, restored_parameter) for _ in range(2)] == expected_draws assert torch.equal(restored_parameter.detach(), expected_parameter) for key, value in original.optimizer.state[parameter].items(): assert torch.equal(value, restored.optimizer.state[restored_parameter][key])