"""Versioned, resumable online NeMo Gym GRPO; historical trainers stay frozen. Upstream TRL owns generation, policy/reference log probabilities and optimization. Checkpoints preserve its adapter, optimizer, scheduler, RNG and trainer state. Only completed optimizer boundaries with a verified completion seal can resume. """ from __future__ import annotations import argparse import copy from datetime import datetime, timezone import fcntl import hashlib import importlib.metadata import inspect import json import math import os from pathlib import Path import random import shutil import statistics from types import SimpleNamespace from . import gym_training as legacy from .agent_rl_pilot import require_learning_update from .evaluation import checkpoint_sha256 from .training import read_json_snapshot, read_jsonl_snapshot, require_versions, sha256 FORMAT = "nemo-grpo-long-v1" SEAL = "checkpoint-complete.json" RECOMMENDED = {"max_steps": 256, "learning_rate": 5e-6, "beta": 0.1} def _atomic_json(path: Path, value: dict) -> None: temporary = path.with_name(path.name + ".tmp") with temporary.open("w", encoding="utf-8") as stream: json.dump(value, stream, indent=2, ensure_ascii=False, allow_nan=False) stream.write("\n") stream.flush() os.fsync(stream.fileno()) os.replace(temporary, path) descriptor = os.open(path.parent, os.O_DIRECTORY) try: os.fsync(descriptor) finally: os.close(descriptor) def _json_hash(value) -> str: return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False).encode()).hexdigest() def _sources() -> dict: directory = Path(__file__).parent return {name: sha256(directory / name) for name in ( "gym_training_long.py", "gym_training.py", "gym_client.py", "agent_rl_training.py", "agent_rl_tasks.py", "agent_rl_pilot.py", "training.py", "evaluation.py", "managed_trial.py", "managed_training.py", )} def task_schedule(rows: list[dict], steps: int, seed: int) -> list[dict]: """Freeze randomized task order once; native sampling then has one stable epoch. TRL's RepeatSampler has a private RNG without checkpoint serialization. A fixed one-epoch dataset avoids depending on that unsaved state across resumes. This RNG is local and never changes policy-generation RNG state. """ if not rows or type(steps) is not int or steps < 1: raise ValueError("A task schedule requires rows and positive integer steps") rng, schedule = random.Random(seed), [] while len(schedule) < steps: indices = list(range(len(rows))) rng.shuffle(indices) for index in indices: task = rows[index]["task"] schedule.append({"optimizer_step": len(schedule), "row_index": index, "task_id": task["task_id"], "episode_seed": task["episode_seed"]}) if len(schedule) == steps: break return schedule def _interval(value: int, steps: int) -> int: if type(value) is not int or not 1 <= value <= steps: raise ValueError("checkpoint_interval must be an integer in [1,max_steps]") return value def plan(config_path: str | Path, gym_url: str, *, checkpoint_interval: int = 64) -> dict: config, _ = read_json_snapshot(config_path) if not set(RECOMMENDED) <= set(config.get("training", {})): raise ValueError("Long recipe must explicitly freeze max_steps, learning_rate and beta") report = legacy.plan(config_path, gym_url) steps = report["training"]["max_steps"] interval = _interval(checkpoint_interval, steps) if report["training"]["beta"] <= 0: raise ValueError("Long recipe requires a positive KL coefficient against the input SFT policy") rows, digest = read_jsonl_snapshot(report["dataset"]["path"]) if digest != report["dataset"]["sha256"]: raise ValueError("Training tasks changed while freezing schedule") for row in rows: # The frozen environment validates all private task metadata before GPU work. from .agent_rl_tasks import GermanAgentEnv GermanAgentEnv().reset(task=row["task"]) environment = legacy.native.audit_task_group task = row["task"] environment([task] * report["training"]["num_generations"], [SimpleNamespace(_task=task)] * report["training"]["num_generations"], report["training"]["num_generations"]) if row["id"] != task["task_id"]: raise ValueError("Dataset row id must match authoritative task_id") schedule = task_schedule(rows, steps, report["seed"]) sources = _sources() report.update( method="native_trl_remote_nemo_gym_grpo_long", format_version=FORMAT, validation_only=False, release_eligible=False, qualification="Online gym training; release requires independent agent improvement and German-math retention gates", long_source_sha256=sources, source_commit=legacy._source_commit(sources), task_schedule=schedule, task_schedule_sha256=_json_hash(schedule), checkpointing={"interval": interval, "format_version": FORMAT, "optimizer_scheduler_rng": True, "total_horizon_steps": steps, "selection_steps": [step for step in (64, 128, 256) if step <= steps]}, retention_strategy={"type": "fixed_input_policy_kl", "beta": report["training"]["beta"], "learning_rate": report["training"]["learning_rate"], "reference": "input SFT checkpoint with RL adapter disabled", "limitation": "Sampled agent-state KL is not a guarantee of German-math retention; no MGSM/test feedback enters updates."}, ) report["fixed"].update(shuffle_dataset=False, save_strategy="steps", save_steps=interval, save_only_model=False, ignore_data_skip=False, lr_scheduler_type="linear", warmup_steps=0, dataloader_num_workers=0) return report def invalid_tool_diagnostics(completion: list[dict], rpc_calls: list[dict], score: dict) -> dict: """Observe native failures outside env.trace; deliberately do not change reward. An unexecuted call at the tool-iteration cap is not counted as an execution error. Only an actual tool-result message can be a native-only result. """ declared, unknown, invalid = 0, 0, 0 matched, native_only = 0, 0 for message in completion: if message.get("role") == "assistant": for call in message.get("tool_calls") or []: declared += 1 try: function = call["function"] name, arguments = function["name"], function["arguments"] if name not in {"lookup", "calculator"}: unknown += 1 continue arguments = json.loads(arguments) if isinstance(arguments, str) else arguments expected = {"record_id"} if name == "lookup" else {"expression"} if (not isinstance(arguments, dict) or set(arguments) != expected or any(not isinstance(value, str) for value in arguments.values())): invalid += 1 except (KeyError, TypeError, ValueError): invalid += 1 elif message.get("role") == "tool": expected = rpc_calls[matched] if matched < len(rpc_calls) else None if expected and (message.get("name"), message.get("content")) == ( expected["tool"], expected["observation"]): matched += 1 else: native_only += 1 return {"declared_calls": declared, "declared_unknown_names": unknown, "declared_invalid_arguments": invalid, "actual_rpc_calls": len(rpc_calls), "matched_rpc_results": matched, "native_only_tool_results": native_only, "rpc_results_missing_from_retained_messages": len(rpc_calls) - matched, "authoritative_tool_errors": sum(not row["ok"] for row in score["trace"]), "reward_unchanged": True} def _read_evidence(path: Path, size: int | None = None) -> list[dict]: data = path.read_bytes() if path.exists() else b"" data = data if size is None else data[:size] if data and not data.endswith(b"\n"): raise ValueError("Completed rollout evidence must end at a full JSONL record") def unique(pairs): result = {} for key, value in pairs: if key in result: raise ValueError("Duplicate rollout evidence key") result[key] = value return result return [json.loads(line, object_pairs_hook=unique, parse_constant=lambda value: (_ for _ in ()).throw(ValueError(value))) for line in data.splitlines()] def evidence_groups(records: list[dict], report: dict, steps: int) -> list[dict]: generations = report["training"]["num_generations"] if len(records) != steps * generations: raise ValueError("Checkpoint evidence must contain exactly one complete fresh group per optimizer step") groups = [] for step in range(steps): group = records[step * generations:(step + 1) * generations] expected = report["task_schedule"][step] for index, row in enumerate(group): if (row["optimizer_step"] != step or row["rollout_index"] != index or row["task_id"] != expected["task_id"] or row["episode_seed"] != expected["episode_seed"] or row["policy_version"] != {"global_step": step, "base_model_sha256": report["model"]["sha256"]}): raise ValueError("Evidence group differs from frozen task schedule or policy version") components = legacy.native.milestone_training_reward(row["score"]) if components != row["reward_components"] or components["training_reward"] != row["training_reward"]: raise ValueError("Recorded reward differs from frozen milestone protocol") audit = row["policy_audit"] spans = [{"prompt_ids": span["prompt_ids"], "completion_ids": span["sampled_completion_ids"], "call_index": span["call_index"]} for span in audit["sampled_spans"]] proven = legacy.native.audit_policy_tokens([audit["prompt_ids"]], [audit["completion_ids"]], [audit["tool_mask"]], spans)[0] if proven != audit: raise ValueError("Stored policy audit differs from exact sampled-token proof") rewards = [row["training_reward"] for row in group] groups.append({"optimizer_step": step, "rewards": rewards, "reward_std": statistics.pstdev(rewards), "strict_successes": sum(row["score"]["success"] for row in group), "strict_rewards": [row["score"]["reward"] for row in group], "task_ids": [row["task_id"] for row in group], "episode_seeds": [row["episode_seed"] for row in group], "distinct_task_count": 1, "policy_version": group[0]["policy_version"]}) return groups def _prefix(path: Path, size: int | None = None) -> dict: digest, consumed = hashlib.sha256(), 0 with path.open("rb") as stream: while size is None or consumed < size: block = stream.read(1024 * 1024 if size is None else min(1024 * 1024, size - consumed)) if not block: break digest.update(block) consumed += len(block) if size is not None and consumed != size: raise ValueError("Evidence prefix is shorter than the checkpoint seal") return {"bytes": consumed, "sha256": digest.hexdigest()} def _file_hashes(path: Path) -> dict: files = {} for entry in sorted(path.rglob("*")): if entry.is_symlink(): raise ValueError("Resumable checkpoints must not contain symlinks") if entry.is_file() and entry.name != SEAL: files[str(entry.relative_to(path))] = sha256(entry) return files def _managed_binding_hash(output: Path, frozen: dict) -> str | None: """Verify the immutable trial binding without requiring a live reporter.""" path = output / "managed-binding.json" expected = frozen.get("managed_training") if expected is None: if path.exists() or path.is_symlink(): raise ValueError("Unmanaged run contains an unexpected managed binding") return None if (not isinstance(expected, dict) or set(expected) != {"binding_sha256"} or path.is_symlink() or not path.is_file()): raise ValueError("Managed run requires its original trial binding") binding, _ = read_json_snapshot(path) if _json_hash(binding) != expected["binding_sha256"]: raise ValueError("Managed trial binding differs from the frozen run plan") return sha256(path) def _bind_managed_reporter(frozen: dict, reporter, *, resume: bool) -> None: """Authorize optimization only for the reporter that owns this managed run.""" output = Path(frozen["output_dir"]) if reporter is None: if (frozen.get("managed_training") is not None or (output / "managed-binding.json").exists() or (output / "managed-binding.json").is_symlink()): raise ValueError("Optimizing a managed run requires its original managed reporter") return binding = getattr(reporter, "binding", None) if (not isinstance(binding, dict) or not binding or any(not callable(getattr(reporter, name, None)) for name in ("report", "preempt_requested", "publish"))): raise ValueError("Managed reporter requires binding, report, preempt_requested and publish") expected = {"binding_sha256": _json_hash(binding)} if resume: if frozen.get("managed_training") != expected: raise ValueError("Managed reporter differs from the original trial binding") _managed_binding_hash(output, frozen) original, _ = read_json_snapshot(output / "managed-binding.json") if original != binding: raise ValueError("Managed reporter binding differs from stored trial binding") else: frozen["managed_training"] = expected def _verify_checkpoint(path: Path, output: Path, frozen: dict) -> dict: if path.parent != output or not path.name.startswith("checkpoint-") or path.is_symlink(): raise ValueError("Resume checkpoint must belong to this exact run output") if not (path / SEAL).is_file(): raise ValueError("Resume requires a completed checkpoint seal") seal, _ = read_json_snapshot(path / SEAL) if (seal.get("format_version") != FORMAT or seal.get("run_plan_sha256") != sha256(output / "run-plan.json") or seal.get("config_sha256") != frozen["config_sha256"] or seal.get("files_sha256") != _file_hashes(path) or seal.get("initial_adapter_sha256") != sha256(output / "initial-adapter.pt") or seal.get("runtime_binding_sha256") != sha256(output / "runtime-binding.json")): raise ValueError("Resume checkpoint provenance or file hashes differ from the sealed run") required = {"optimizer.pt", "scheduler.pt", "rng_state.pth", "trainer_state.json", "adapter_config.json", "adapter_model.safetensors", "trainable-state.pt", "learning-evidence.json"} if not required <= set(seal["files_sha256"]): raise ValueError("Checkpoint lacks optimizer, scheduler, RNG, trainer or adapter state") managed_hash = _managed_binding_hash(output, frozen) if seal.get("managed_binding_sha256") != managed_hash: raise ValueError("Checkpoint managed binding hash differs from its sealed run") reason = seal.get("checkpoint_reason", "scheduled") if reason not in {"scheduled", "managed_preemption"} or (reason == "managed_preemption" and managed_hash is None): raise ValueError("Checkpoint has an invalid optimizer boundary reason") state, _ = read_json_snapshot(path / "trainer_state.json") step = seal.get("optimizer_step") horizon = frozen["training"]["max_steps"] if (type(step) is not int or not 1 <= step <= horizon or state.get("global_step") != step or state.get("max_steps") != horizon or path.name != f"checkpoint-{step}" or (reason == "scheduled" and step != horizon and step % frozen["checkpointing"]["interval"])): raise ValueError("Checkpoint trainer state and sealed optimizer step disagree") evidence = output / "rollout-evidence.jsonl" if _prefix(evidence, seal["evidence_prefix"]["bytes"]) != seal["evidence_prefix"]: raise ValueError("Resume evidence prefix hash mismatch") diagnostics = output / "native-tool-diagnostics.jsonl" if _prefix(diagnostics, seal["diagnostics_prefix"]["bytes"]) != seal["diagnostics_prefix"]: raise ValueError("Resume native tool diagnostics prefix hash mismatch") evidence_groups(_read_evidence(evidence, seal["evidence_prefix"]["bytes"]), frozen, step) return seal def validate_checkpoint(output_dir: str | Path, checkpoint: str | Path, *, expected_step: int | None = None) -> dict: """Read-only artifact verification for driver adoption; no model/server calls. The actual resume additionally rechecks input weights, data and live native server identity. This helper never reads untrusted optimizer/RNG pickle data. """ output, path = Path(output_dir).resolve(), Path(checkpoint) path = (output / path if not path.is_absolute() else path).resolve() frozen, _ = read_json_snapshot(output / "run-plan.json") if (str(output) != frozen["output_dir"] or _sources() != frozen["long_source_sha256"] or sha256(frozen["config_path"]) != frozen["config_sha256"]): raise ValueError("Checkpoint consumed source or configuration changed") seal = _verify_checkpoint(path, output, frozen) if expected_step is not None and seal["optimizer_step"] != expected_step: raise ValueError("Sealed checkpoint is not the expected optimizer step") return seal | {"checkpoint": str(path), "checkpoint_sha256": checkpoint_sha256(path)} def discover_resume_checkpoint(output_dir: str | Path) -> Path | None: """Return the latest fully verified checkpoint, skipping interrupted saves.""" output = Path(output_dir).resolve() candidates = sorted((path for path in output.glob("checkpoint-*") if path.is_dir() and path.name.removeprefix("checkpoint-").isdigit()), key=lambda path: int(path.name.removeprefix("checkpoint-")), reverse=True) for path in candidates: try: validate_checkpoint(output, path) return path except (ValueError, OSError, KeyError, TypeError): continue return None def _resume_plan(config_path: Path, gym_url: str, checkpoint: Path, checkpoint_interval: int) -> tuple[dict, dict]: config, digest = read_json_snapshot(config_path) if int(os.environ.get("WORLD_SIZE", "1")) != 1: raise ValueError("Long Gym resume supports one CUDA process only") output = (config_path.parent / config["output_dir"]).resolve() frozen, _ = read_json_snapshot(output / "run-plan.json") if (frozen.get("format_version") != FORMAT or digest != frozen["config_sha256"] or str(config_path) != frozen["config_path"] or str(output) != frozen["output_dir"] or _sources() != frozen["long_source_sha256"] or checkpoint_interval != frozen["checkpointing"]["interval"] or _json_hash(frozen["task_schedule"]) != frozen["task_schedule_sha256"]): raise ValueError("Resume must preserve the exact configuration, schedule, source and checkpoint interval") if (sha256(frozen["dataset"]["path"]) != frozen["dataset"]["sha256"] or checkpoint_sha256(frozen["model"]["path"]) != frozen["model"]["sha256"]): raise ValueError("Resume input model or training data changed") if legacy._gym_api().fetch_gym_identity(gym_url) != frozen["gym"]["identity"]: raise ValueError("Resume gym server identity changed") return frozen, _verify_checkpoint(checkpoint, output, frozen) def _quarantine_suffix(output: Path, checkpoint: Path, seal: dict) -> dict | None: """Archive orphan work, retaining the exact evidence prefix used by optimizer state.""" evidence = output / "rollout-evidence.jsonl" raw = evidence.read_bytes() size = seal["evidence_prefix"]["bytes"] future = [path for path in output.glob("checkpoint-*") if path.is_dir() and path.name.removeprefix("checkpoint-").isdigit() and int(path.name.removeprefix("checkpoint-")) > seal["optimizer_step"]] history = output / "recovery" history.mkdir(exist_ok=True) index = len(list(history.iterdir())) + 1 destination = history / f"attempt-{index:04d}" destination.mkdir() if (output / "training-manifest.json").exists(): shutil.copy2(output / "training-manifest.json", destination / "previous-training-manifest.json") if len(raw) > size: (destination / "uncommitted-rollout-suffix.jsonl").write_bytes(raw[size:]) temporary = evidence.with_suffix(".recovering") temporary.write_bytes(raw[:size]) os.replace(temporary, evidence) diagnostics = output / "native-tool-diagnostics.jsonl" diagnostic_data = diagnostics.read_bytes() diagnostic_size = seal["diagnostics_prefix"]["bytes"] if len(diagnostic_data) > diagnostic_size: (destination / "uncommitted-native-tool-diagnostics.jsonl").write_bytes(diagnostic_data[diagnostic_size:]) temporary = diagnostics.with_suffix(".recovering") temporary.write_bytes(diagnostic_data[:diagnostic_size]) os.replace(temporary, diagnostics) for path in future: shutil.move(str(path), destination / path.name) return {"path": str(destination), "discarded_evidence_bytes": len(raw) - size, "quarantined_checkpoints": [path.name for path in future], "resume_from": str(checkpoint)} def _tensor_state(model) -> dict: return {name: tensor.detach().cpu().clone() for name, tensor in model.named_parameters() if tensor.requires_grad} def _response_parser_identity() -> dict: # Lazy imports keep source planning and CPU-only orchestration independent # of the optional training stack. Qualification binds the actual C engine. import importlib version = importlib.metadata.version("regex") if version != "2026.9.29": raise RuntimeError("Bounded Qwen parsing requires pinned regex==2026.9.29") files = {} for name in ("regex", "regex._main", "regex._regex_core", "regex._regex"): module = importlib.import_module(name) path = Path(module.__file__) if not path.is_file(): raise RuntimeError("Qwen parser engine has no inspectable source file") files[name] = sha256(path) return {"version": version, "source_sha256": files} def _runtime_identity(runtime, installed: dict) -> dict: parser = _response_parser_identity() files = {} for trainer_class in inspect.getmro(runtime.Trainer): if trainer_class is object: continue files[f"{trainer_class.__module__}.{trainer_class.__qualname__}"] = sha256(Path(inspect.getfile(trainer_class))) return {"dependencies": installed | {"jmespath": "1.0.1", "regex": parser["version"]}, "upstream_class_source_sha256": files, "response_parser_engine": parser} def _tensor_delta(current: dict, previous: dict, torch) -> dict: if not current or set(current) != set(previous): raise ValueError("Trainable adapter state changed shape or parameter identity") changed, maximum = 0, 0.0 for name, tensor in current.items(): if (tensor.shape != previous[name].shape or not bool(torch.isfinite(tensor).all()) or not bool(torch.isfinite(previous[name]).all())): raise ValueError("Nonfinite or incompatible trainable adapter parameter") delta = float((tensor.float() - previous[name].float()).abs().max()) changed += delta > 0 maximum = max(maximum, delta) return {"changed_trainable_tensors": changed, "max_abs_parameter_delta": maximum} def _learning_evidence(groups: list[dict], logs: list[dict], after: int = 0) -> dict: selected = [group for group in groups if group["optimizer_step"] >= after] rewards = [value for group in selected for value in group["rewards"]] norms = [float(log["grad_norm"]) for log in logs if log.get("step", 0) > after and "grad_norm" in log] if any(not math.isfinite(norm) or norm < 0 for norm in norms): raise ValueError("Nonfinite or negative gradient norm cannot qualify a resumable learning checkpoint") return {"optimizer_groups": len(selected), "rollouts": len(rewards), "mixed_reward_groups": sum(group["reward_std"] > 0 for group in selected), "nonzero_gradient_steps": sum(norm > 0 for norm in norms), "strict_successes": sum(group["strict_successes"] for group in selected), "training_reward_mean": statistics.mean(rewards) if rewards else None, "training_reward_note": "Shaped training reward is not strict task success or heldout performance", "logged_kl": [{"step": log["step"], "kl": log["kl"]} for log in logs if log.get("step", 0) > after and "kl" in log]} def _seal_checkpoint(trainer, context) -> dict: step, report = trainer.state.global_step, context.report checkpoint = context.output / f"checkpoint-{step}" if (checkpoint / SEAL).exists(): raise ValueError("Refusing to overwrite a sealed checkpoint") if _sources() != report["long_source_sha256"] or sha256(report["config_path"]) != report["config_sha256"]: raise RuntimeError("Consumed source or configuration changed during training") if legacy._gym_api().fetch_gym_identity(context.gym_url) != report["gym"]["identity"]: raise RuntimeError("Gym identity changed before checkpoint sealing") records = _read_evidence(context.evidence) groups = evidence_groups(records, report, step) current = _tensor_state(trainer.model) overall = _tensor_delta(current, context.initial, context.torch) interval = _tensor_delta(current, context.previous, context.torch) learning = {"optimizer_step": step, "from_initial_adapter": overall, "since_previous_checkpoint": interval, "previous_checkpoint_step": context.previous_step, "cumulative": _learning_evidence(groups, trainer.state.log_history), "interval": _learning_evidence(groups, trainer.state.log_history, context.previous_step), "limitations": "Adapter updates and training rewards do not establish merged BF16 changes or heldout improvement."} diagnostics = _read_evidence(context.diagnostics_path) if len(diagnostics) != step or [row["optimizer_step"] for row in diagnostics] != list(range(step)): raise ValueError("Native tool diagnostics must align with every committed optimizer group") learning["native_tool_diagnostics"] = { key: sum(rollout[key] for group in diagnostics for rollout in group["rollouts"]) for key in ("declared_calls", "declared_unknown_names", "declared_invalid_arguments", "actual_rpc_calls", "matched_rpc_results", "native_only_tool_results", "authoritative_tool_errors")} context.torch.save(current, checkpoint / "trainable-state.pt") _atomic_json(checkpoint / "learning-evidence.json", learning) reason = "managed_preemption" if getattr(context, "preemption_step", None) == step else "scheduled" managed_hash = _managed_binding_hash(context.output, report) reporter = getattr(context, "managed_reporter", None) if reporter is not None and _json_hash(reporter.binding) != report["managed_training"]["binding_sha256"]: raise ValueError("Managed reporter binding changed during optimization") seal = {"format_version": FORMAT, "optimizer_step": step, "checkpoint_reason": reason, "run_plan_sha256": sha256(context.output / "run-plan.json"), "config_sha256": report["config_sha256"], "initial_adapter_sha256": sha256(context.output / "initial-adapter.pt"), "runtime_binding_sha256": sha256(context.output / "runtime-binding.json"), "evidence_prefix": _prefix(context.evidence), "diagnostics_prefix": _prefix(context.diagnostics_path), "files_sha256": _file_hashes(checkpoint), "sealed_at": datetime.now(timezone.utc).isoformat()} if managed_hash is not None: seal["managed_binding_sha256"] = managed_hash for file in checkpoint.rglob("*"): if file.is_file(): with file.open("rb") as stream: os.fsync(stream.fileno()) # Written last: a partial upstream save or failed audit is never resumable. _atomic_json(checkpoint / SEAL, seal) _verify_checkpoint(checkpoint, context.output, report) context.previous, context.previous_step = current, step context.sealed.append({"path": str(checkpoint), "optimizer_step": step, "seal_sha256": sha256(checkpoint / SEAL), "learning": learning}) return seal def _long_trainer_class(base): class LongGymTrainer(legacy._gym_trainer_class(base)): def _save_checkpoint(self, model, trial): super()._save_checkpoint(model, trial) context = self._long_context _seal_checkpoint(self, context) step = self.state.global_step if (context.managed_reporter is not None and step < context.report["training"]["max_steps"] and context.preemption_step != step): # Stopping saves publish once, after outer cleanup and final manifest. context.managed_reporter.publish(context.output, context.output / f"checkpoint-{step}") if context.stop_after_step is not None and self.state.global_step >= context.stop_after_step: self.control.should_training_stop = True return LongGymTrainer def _managed_callback(context, *, callback_base=None): if callback_base is None: from transformers import TrainerCallback callback_base = TrainerCallback class ManagedCallback(callback_base): def on_step_end(self, args, state, control, **kwargs): step = state.global_step if step < context.report["training"]["max_steps"]: requested = context.managed_reporter.preempt_requested(step) if type(requested) is not bool: raise ValueError("Managed preemption status must be boolean") if requested: context.preemption_step = step control.should_save = True control.should_training_stop = True return control def on_log(self, args, state, control, logs=None, **kwargs): if logs and state.global_step > 0: context.managed_reporter.report(state.global_step, dict(logs)) return control return ManagedCallback() def _reward_callback(trainer_ref, context): base = legacy._reward_callback(trainer_ref, context.report["training"], context.report["model"]["sha256"], context.evidence, context.groups) def reward(completions, environments, **kwargs): trainer = trainer_ref[0] step = trainer.state.global_step expected = context.report["task_schedule"][step] metadata = legacy.native.audit_task_group(kwargs.get("task"), environments, context.report["training"]["num_generations"]) if any((row["task_id"], row["episode_seed"]) != (expected["task_id"], expected["episode_seed"]) for row in metadata): raise ValueError("Native GRPO group is not the frozen task for this optimizer step") if len(context.groups) != step: raise ValueError("Exactly one newly sampled reward group is required per optimizer step") # Legacy callback verifies actual server observations by independent replay, # checks sampled assistant IDs/masks, and retains the unchanged reward. values = base(completions, environments, **kwargs) diagnostics = [invalid_tool_diagnostics(completion, environment._calls, score) for completion, environment, score in zip( completions, environments, [environment._score(completion) for environment, completion in zip(environments, completions, strict=True)], strict=True)] context.diagnostics.append({"optimizer_step": step, "rollouts": diagnostics}) with context.diagnostics_path.open("a", encoding="utf-8") as stream: stream.write(json.dumps(context.diagnostics[-1], ensure_ascii=False, allow_nan=False) + "\n") stream.flush() os.fsync(stream.fileno()) # Durable before the corresponding optimizer/checkpoint is allowed to commit. with context.evidence.open("ab") as stream: stream.flush() os.fsync(stream.fileno()) return values return reward def run(config_path: str | Path, gym_url: str, *, resume_from_checkpoint: str | Path | None = None, stop_after_step: int | None = None, checkpoint_interval: int = 64, managed_reporter=None) -> dict: config_path = Path(config_path).resolve() config, _ = read_json_snapshot(config_path) managed_path = (config_path.parent / config["output_dir"]).resolve() / "managed-binding.json" if managed_reporter is None and (managed_path.exists() or managed_path.is_symlink()): raise ValueError("Optimizing a managed run requires its original managed reporter") resume = Path(resume_from_checkpoint).resolve() if resume_from_checkpoint else None frozen, resume_seal = (_resume_plan(config_path, gym_url, resume, checkpoint_interval) if resume else (plan(config_path, gym_url, checkpoint_interval=checkpoint_interval), None)) _bind_managed_reporter(frozen, managed_reporter, resume=bool(resume)) settings, start_step = frozen["training"], resume_seal["optimizer_step"] if resume else 0 if stop_after_step is not None and (type(stop_after_step) is not int or not start_step < stop_after_step <= settings["max_steps"] or (stop_after_step != settings["max_steps"] and stop_after_step % checkpoint_interval)): raise ValueError("stop_after_step must be a future saved checkpoint boundary within the fixed horizon") if start_step >= settings["max_steps"]: raise ValueError("The selected checkpoint has already completed this training horizon") rows, digest = read_jsonl_snapshot(frozen["dataset"]["path"]) if (digest != frozen["dataset"]["sha256"] or sha256(config_path) != frozen["config_sha256"] or _sources() != frozen["long_source_sha256"]): raise ValueError("Long recipe or training data changed after planning") installed = require_versions() if importlib.metadata.version("jmespath") != "1.0.1": raise RuntimeError("Native Qwen tools require pinned jmespath==1.0.1") runtime, gym = legacy._runtime(), legacy._gym_api() torch = runtime.torch if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported() or torch.cuda.device_count() != 1: raise RuntimeError("Use exactly one visible allocated CUDA GPU with BF16 support") binding = _runtime_identity(runtime, installed) if resume: original, _ = read_json_snapshot(Path(frozen["output_dir"]) / "runtime-binding.json") if binding != original: raise ValueError("Resume upstream trainer source or dependency versions changed") output = Path(frozen["output_dir"]) if not resume: output.mkdir(parents=True, exist_ok=False) lock = (output / ".training.lock").open("a") try: fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) except BlockingIOError: lock.close() raise RuntimeError("Another process holds this training output") from None report, environments, trainer_ref = copy.deepcopy(frozen), [], [None] manifest = output / "training-manifest.json" context = SimpleNamespace(output=output, evidence=output / "rollout-evidence.jsonl", report=frozen, gym_url=gym_url, torch=torch, groups=[], sealed=[], diagnostics=[], diagnostics_path=output / "native-tool-diagnostics.jsonl", previous_step=start_step, stop_after_step=stop_after_step, managed_reporter=managed_reporter, preemption_step=None) report.update(status="running", dependencies=installed | {"jmespath": "1.0.1"}, started_at=datetime.now(timezone.utc).isoformat(), gpu=torch.cuda.get_device_name(), evidence=str(context.evidence), resumed_from=str(resume) if resume else None, started_from_optimizer_step=start_step, requested_stop_after_step=stop_after_step, upstream_trainer_sha256=sha256(Path(inspect.getfile(runtime.Trainer))), resume_transport_url=gym_url) try: if resume: report["recovery"] = _quarantine_suffix(output, resume, resume_seal) context.groups = evidence_groups(_read_evidence(context.evidence), frozen, start_step) context.diagnostics = _read_evidence(context.diagnostics_path) else: if managed_reporter is not None: _atomic_json(output / "managed-binding.json", copy.deepcopy(managed_reporter.binding)) _atomic_json(output / "run-plan.json", frozen) if not resume: _atomic_json(output / "runtime-binding.json", binding) report["runtime_binding"] = binding _atomic_json(manifest, report) runtime.set_seed(report["seed"]) tokenizer = runtime.Tokenizer.from_pretrained(report["model"]["path"], local_files_only=True, trust_remote_code=False) report["policy_format"] = legacy.native.configure_policy_tokenizer(tokenizer) tokenizer.save_pretrained(str(output / "policy-tokenizer")) model = runtime.Model.from_pretrained(report["model"]["path"], local_files_only=True, trust_remote_code=False, use_safetensors=True, dtype=torch.bfloat16) if checkpoint_sha256(report["model"]["path"]) != report["model"]["sha256"]: raise ValueError("Base checkpoint changed while loading") model.config.use_cache = False model.generation_config.eos_token_id = tokenizer.eos_token_id def factory(): environment = gym.RemoteGymEnv(gym_url, identity=report["gym"]["identity"]) environments.append(environment) public = {name for name, _ in inspect.getmembers(environment, inspect.ismethod) if not name.startswith("_")} if public != {"reset", "lookup", "calculator"}: raise ValueError("Gym must expose only reset, lookup and calculator") return environment arguments = runtime.Config( output_dir=str(output), **settings, num_iterations=1, steps_per_generation=1, loss_type="grpo", full_determinism=True, bf16=True, fp16=False, gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False}, chat_template_kwargs=report["chat_template_kwargs"], use_vllm=False, disable_dropout=True, report_to="none", push_to_hub=False, seed=report["seed"], data_seed=report["seed"], optim="adamw_torch", shuffle_dataset=False, save_strategy="steps", save_steps=checkpoint_interval, save_total_limit=None, save_only_model=False, ignore_data_skip=False, lr_scheduler_type="linear", warmup_steps=0, logging_steps=1, dataloader_num_workers=0, dataloader_pin_memory=False, mask_truncated_completions=False, log_completions=False, ) scheduled = [rows[entry["row_index"]] for entry in frozen["task_schedule"]] trainer = _long_trainer_class(runtime.Trainer)( model=model, processing_class=tokenizer, args=arguments, train_dataset=runtime.Dataset.from_list([ {"id": row["id"], "prompt": [{"role": "user", "content": row["prompt"]}], "task": row["task"]} for row in scheduled]), reward_funcs=_reward_callback(trainer_ref, context), environment_factory=factory, callbacks=[_managed_callback(context)] if managed_reporter is not None else None, peft_config=runtime.LoraConfig(r=report["lora"]["r"], lora_alpha=report["lora"]["alpha"], lora_dropout=0.0, target_modules="all-linear", task_type="CAUSAL_LM"), ) trainer_ref[0], trainer._long_context = trainer, context if trainer.ref_model is not None: raise ValueError("Expected tracked input SFT reference with RL adapter disabled") if resume: context.initial = torch.load(output / "initial-adapter.pt", map_location="cpu", weights_only=True) context.previous = torch.load(resume / "trainable-state.pt", map_location="cpu", weights_only=True) # These counters do not change the one-use rollout policy, but retain # native step accounting. No stale generation buffer crosses resume. trainer._step, trainer._buffered_inputs = start_step, None else: context.initial = _tensor_state(trainer.model) if not context.initial: raise ValueError("LoRA created no trainable adapter parameters") context.previous = context.initial torch.save(context.initial, output / "initial-adapter.pt") trained = trainer.train(**({"resume_from_checkpoint": str(resume)} if resume else {})) step = trainer.state.global_step expected_end = context.preemption_step or stop_after_step or settings["max_steps"] if step != expected_end or not (output / f"checkpoint-{step}" / SEAL).is_file(): raise RuntimeError("Trainer did not reach and seal the requested optimizer boundary") for environment in environments: environment._raise_if_failed() groups = evidence_groups(_read_evidence(context.evidence), frozen, step) checkpoint = output / f"checkpoint-{step}" learning, _ = read_json_snapshot(checkpoint / "learning-evidence.json") status = ("completed" if step == settings["max_steps"] else "preempted_at_checkpoint" if context.preemption_step == step else "stopped_at_checkpoint") report.update(status=status, optimizer_step=step, checkpoint=str(checkpoint), checkpoint_sha256=checkpoint_sha256(checkpoint), metrics=trained.metrics, train_log_history=trainer.state.log_history, reward_groups=groups, sealed_checkpoints=context.sealed, learning_evidence=learning, update_evidence=learning["from_initial_adapter"] | { key: learning["cumulative"][key] for key in ("mixed_reward_groups", "nonzero_gradient_steps")}, task_coverage={"training_pool_rows": len(rows), "optimizer_groups": len(groups), "total_rollouts": len(groups) * settings["num_generations"], "consumed_unique_task_ids": sorted({identity for group in groups for identity in group["task_ids"]})}, completed_at=datetime.now(timezone.utc).isoformat()) if report["status"] == "completed": require_learning_update(report) return report except BaseException as error: report.update(status="failed", error={"type": type(error).__name__, "message": str(error)}) raise finally: cleanup = legacy._close_sessions(environments) cleanup_failure = bool(cleanup) and report["status"] in {"completed", "stopped_at_checkpoint", "preempted_at_checkpoint"} if cleanup: report["cleanup_errors"] = cleanup if cleanup_failure: report.update(status="failed", error={"type": "GymInfrastructureError", "message": "gym session cleanup failed"}) if context.evidence.exists(): report["evidence_sha256"] = sha256(context.evidence) if context.diagnostics_path.exists(): report["native_tool_diagnostics_sha256"] = sha256(context.diagnostics_path) try: _atomic_json(manifest, report) finally: fcntl.flock(lock, fcntl.LOCK_UN) lock.close() if cleanup_failure: raise gym.GymInfrastructureError("gym session cleanup failed") def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("config", type=Path) parser.add_argument("--gym-url", required=True) parser.add_argument("--checkpoint-interval", type=int, default=64) parser.add_argument("--resume-from-checkpoint", type=Path) parser.add_argument("--stop-after-step", type=int) parser.add_argument("--dry-run", action="store_true") args = parser.parse_args() result = (plan(args.config, args.gym_url, checkpoint_interval=args.checkpoint_interval) if args.dry_run else run(args.config, args.gym_url, checkpoint_interval=args.checkpoint_interval, resume_from_checkpoint=args.resume_from_checkpoint, stop_after_step=args.stop_after_step)) print(json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False)) if __name__ == "__main__": main()