"""Matched OpenEnv training arm with the frozen NeMo evaluation protocol. The original driver is compiled in a private namespace. Only its trainer imports, trainer CLI target, plan verification and service ownership are adapted. Evaluation, DEV selection, stop/resume qualification, merging and budget accounting are reused. No NeMo module, source file, prepared plan or journal is modified. """ from __future__ import annotations import argparse import ast from contextlib import contextmanager from datetime import datetime, timezone import hashlib import importlib.util import importlib.metadata import json import os from pathlib import Path import platform import subprocess from types import ModuleType from . import gym_experiment as nemo from . import gym_experiment_evaluation as evaluation from . import openenv_training as training from .evaluation import sha256_file FORMAT = "hessian-openenv-matched-experiment-v1" ENGINE_SHA256 = "1b8453d0fb6b3cce5b3009de9156ded6983db2a2fb900124f6f77f5331e6eaf4" REWRITES = {"trainer_imports": 6, "trainer_cli": 2} NATIVE_VERSION = "0.9.0" NATIVE_REVISION = "17c1edd22168fcb76b086def1ed90a7168fb8749" REUSE_POLICY = "reuse_complete_matching_suite_or_rerun_whole_suite" CPU_CHECKS = { "inherited_python_overlay_removed", "typed_native_reset_state_no_hidden_facts", "two_native_sessions_isolated_and_strict_score_parity", "unknown_names_kwargs_reserved_keys_badtypes_scored_zero", "verification_idempotent_terminal_immutable_reset_reusable", "32_concurrent_native_sessions_unique_and_correct", "native_capacity_rejection_and_recovery", "native_wire_duplicate_keys_and_nonfinite_json_rejected", "failed_controller_reset_discards_prior_task_and_recovers", "260_sequential_resets_no_session_limit_leak", "native_close_releases_all_environment_instances", "owned_service_process_reaped", } ROOT = Path(__file__).resolve().parents[2] def _snapshot(path): raw = Path(path).read_bytes() from .training import _unique_pairs value = json.loads(raw, object_pairs_hook=_unique_pairs, parse_constant=lambda value: (_ for _ in ()).throw(ValueError(value))) return value, raw, hashlib.sha256(raw).hexdigest() def _sources(): sources = nemo._sources() | training.source_hashes() for name in ("src/hessian_recipes/openenv_experiment.py", "examples/agent_gyms/openenv_smoke.py", "scripts/run-openenv-after-nemo.py", "cluster/determined/german-qwen-openenv-command.yaml"): sources[name] = sha256_file(ROOT / name) return sources def _stable_identity(identity): return {key: value for key, value in identity.items() if key != "active_environment_instances"} def _cpu_proof(path): proof, _, digest = _snapshot(path) checks = proof.get("checks") identity = proof.get("native_identity", {}) upstream = proof.get("upstream", {}) proof_sources = proof.get("source_sha256", {}) if (proof.get("schema_version") != 1 or proof.get("status") != "passed" or proof.get("backend") != "openenv" or proof.get("mode") != "native-sdk-cpu-scripted-fixtures" or proof.get("native_sdk_executed") is not True or proof.get("model_inference") is not False or proof.get("gpu") is not False or proof.get("strict_reward_unchanged") is not True or proof.get("server_released") is not True or not isinstance(checks, dict) or not CPU_CHECKS <= checks.keys() or any(value is not True for value in checks.values()) or upstream.get("revision") != NATIVE_REVISION or upstream.get("version") != NATIVE_VERSION or identity.get("backend") != "openenv" or identity.get("protocol") != "hessian-openenv-v1" or identity.get("transport") != "native-sdk-websocket" or identity.get("upstream_revision") != NATIVE_REVISION or identity.get("runtime", {}).get("version") != NATIVE_VERSION or identity.get("runtime", {}).get("commit") != NATIVE_REVISION or identity.get("server_sha256") != sha256_file(ROOT / "examples/agent_gyms/openenv_server.py") or identity.get("environment_sha256") != sha256_file(ROOT / "src/hessian_recipes/agent_rl_tasks.py") or identity.get("requirements_sha256") != sha256_file(ROOT / "requirements/openenv.txt") or any(proof_sources.get(key) != sha256_file(ROOT / relative) for key, relative in { "server": "examples/agent_gyms/openenv_server.py", "smoke": "examples/agent_gyms/openenv_smoke.py", "requirements": "requirements/openenv.txt"}.items())): raise ValueError("A complete matching native OpenEnv CPU qualification is required") return {"path": str(Path(path).resolve()), "sha256": digest, "native_identity": _stable_identity(identity)} def _verify_source_commit(plan): if plan["source_commit"] is None: return # Source hashes remain authoritative for an unpacked source archive. try: for name, digest in plan["source_sha256"].items(): blob = subprocess.check_output( ["git", "-C", str(ROOT), "show", f"{plan['source_commit']}:{name}"], stderr=subprocess.DEVNULL) if hashlib.sha256(blob).hexdigest() != digest: raise ValueError("OpenEnv source commit does not contain the consumed bytes") except (OSError, subprocess.CalledProcessError) as error: raise ValueError("OpenEnv source commit cannot be verified") from error def _nemo_plan(path): _, raw, digest = _snapshot(path) plan = nemo.verify_plan(path) if Path(path).read_bytes() != raw: raise ValueError("Original NeMo plan changed during verification") return plan, digest def prepare(nemo_plan, output, sdk_python, nativepin, *, reuse_baseline=True, budget_carry=None): """Freeze an arm; nativepin is the completed native CPU smoke report path. Copies the evaluation JSON and TRAIN JSONL bytes exactly. Their source data paths remain absolute paths to the original frozen NeMo inputs. No inference, GPU allocation, model/data upload or NeMo experiment mutation occurs. """ original, original_sha = _nemo_plan(nemo_plan) if original.get("budget_carry") is not None and (budget_carry is None or reuse_baseline): raise ValueError("Recovery OpenEnv requires its own carried budget and fresh baseline suites") root = Path(output).resolve() if root.is_relative_to(Path(original["root"])): raise ValueError("OpenEnv requires a separate arm root outside the NeMo experiment") for name in ("experiment-plan.json", "evaluation-plan", "train.jsonl", "training.json"): if (root / name).exists(): raise ValueError("Choose a fresh OpenEnv experiment directory") sdk_python = str(Path(sdk_python).absolute()) # Preserve the isolated venv symlink. if not Path(sdk_python).is_file() or not os.access(sdk_python, os.X_OK): raise ValueError("sdk_python must be an existing executable") pin = _cpu_proof(nativepin) from .openenv_client import sdk_identity sdk_runtime = sdk_identity(sdk_python) config, _, config_sha = _snapshot(original["training_config"]) train = Path(original["root"]) / "train.jsonl" train_bytes = train.read_bytes() protocol, protocol_bytes, protocol_sha = _snapshot(original["evaluation_plan"]) if (config_sha != original["training_config_sha256"] or hashlib.sha256(train_bytes).hexdigest() != original["train_sha256"] or protocol_sha != original["evaluation_plan_sha256"]): raise ValueError("Original NeMo input snapshot changed") root.mkdir(parents=True, exist_ok=True) (root / "evaluation-plan").mkdir() (root / "evaluation-plan/plan.json").write_bytes(protocol_bytes) (root / "train.jsonl").write_bytes(train_bytes) config.update(output_dir=str(root / "training"), dataset={"path": str(root / "train.jsonl"), "sha256": original["train_sha256"]}) nemo._write(root / "training.json", config) sources = _sources() plan = {"format": FORMAT, "backend": "openenv", "status": "prepared", "release_eligible": False, "created_at": datetime.now(timezone.utc).isoformat(), "root": str(root), "nemo_origin": {"path": str(Path(nemo_plan).resolve()), "sha256": original_sha}, "evaluation_plan": str(root / "evaluation-plan/plan.json"), "evaluation_plan_sha256": protocol_sha, "training_config": str(root / "training.json"), "training_config_sha256": sha256_file(root / "training.json"), "train_sha256": original["train_sha256"], "source_sha256": sources, "source_commit": training._source_commit(sources), "native_pin": pin, "sdk": {"python": sdk_python, "executable_sha256": sha256_file(sdk_python), "runtime": sdk_runtime}, "checkpoint_steps": list(nemo.CHECKPOINTS), "checkpoint_interval": 64, "max_gpu_wall_seconds": 8 * 3600, "resume_qualification": {"steps": 4, "stop_step": 2, "seed": 43}, "selection": protocol["selection_rule"], "baseline_reuse": {"enabled": reuse_baseline, "policy": REUSE_POLICY}, "notice": "Matched framework arm; one training seed, narrow synthetic task family; no broad quality or release claim."} if budget_carry is not None: plan["budget_carry"] = nemo.budget_lineage.reference(budget_carry) nemo.budget_lineage.load(plan) nemo._write(root / "experiment-plan.json", plan) verify_plan(root / "experiment-plan.json") return plan def verify_plan(path): path = Path(path).resolve() plan, _, _ = _snapshot(path) root = Path(plan["root"]) original, digest = _nemo_plan(plan["nemo_origin"]["path"]) if (plan.get("format") != FORMAT or plan.get("backend") != "openenv" or plan.get("release_eligible") is not False or path != root / "experiment-plan.json" or plan["nemo_origin"]["sha256"] != digest or root.is_relative_to(Path(original["root"])) or plan["source_sha256"] != _sources() or plan["native_pin"] != _cpu_proof(plan["native_pin"]["path"]) or plan["evaluation_plan"] != str(root / "evaluation-plan/plan.json") or Path(plan["evaluation_plan"]).read_bytes() != Path(original["evaluation_plan"]).read_bytes() or plan["evaluation_plan_sha256"] != original["evaluation_plan_sha256"] or plan["evaluation_plan_sha256"] != sha256_file(plan["evaluation_plan"]) or (root / "train.jsonl").read_bytes() != (Path(original["root"]) / "train.jsonl").read_bytes() or plan["train_sha256"] != original["train_sha256"] or plan["train_sha256"] != sha256_file(root / "train.jsonl") or plan["training_config"] != str(root / "training.json") or plan["training_config_sha256"] != sha256_file(plan["training_config"]) or plan["checkpoint_steps"] != list(nemo.CHECKPOINTS) or plan["checkpoint_interval"] != 64 or plan["max_gpu_wall_seconds"] != 8 * 3600 or plan["resume_qualification"] != {"steps": 4, "stop_step": 2, "seed": 43} or type(plan["baseline_reuse"].get("enabled")) is not bool or plan["baseline_reuse"].get("policy") != REUSE_POLICY or plan["sdk"]["python"] != str(Path(plan["sdk"]["python"]).absolute()) or sha256_file(plan["sdk"]["python"]) != plan["sdk"]["executable_sha256"]): raise ValueError("Frozen OpenEnv arm source or inputs changed") if original.get("budget_carry") is not None and (plan.get("budget_carry") is None or plan["baseline_reuse"]["enabled"]): raise ValueError("Recovery OpenEnv lost its carried budget or fresh baseline requirement") protocol, _ = evaluation.verify_plan(plan["evaluation_plan"]) config = nemo._read(original["training_config"]) config.update(output_dir=str(root / "training"), dataset={"path": str(root / "train.jsonl"), "sha256": plan["train_sha256"]}) if nemo._read(plan["training_config"]) != config or plan["selection"] != protocol["selection_rule"]: raise ValueError("OpenEnv training configuration differs from the matched NeMo arm") _verify_source_commit(plan) nemo.budget_lineage.load(plan) return plan class _Adapt(ast.NodeTransformer): """Only six explicit trainer imports and two CLI constants may differ.""" def __init__(self): self.counts = dict.fromkeys(REWRITES, 0) def visit_ImportFrom(self, node): if node.module in {"gym_training_long", "hessian_recipes.gym_training_long"}: replacements = [] for name in node.names: if name.name != "validate_checkpoint": raise RuntimeError("Unexpected frozen trainer import") replacements.append(ast.Assign( targets=[ast.Name(name.asname or name.name, ast.Store())], value=ast.Attribute(ast.Name("_openenv_training", ast.Load()), name.name, ast.Load()))) self.counts["trainer_imports"] += 1 return [ast.copy_location(row, node) for row in replacements] if node.module in {None, "hessian_recipes"}: names = [name for name in node.names if name.name == "gym_training_long"] if names: if len(node.names) != 1: raise RuntimeError("Unexpected combined trainer import") self.counts["trainer_imports"] += 1 return ast.copy_location(ast.Assign( targets=[ast.Name(names[0].asname or names[0].name, ast.Store())], value=ast.Name("_openenv_training", ast.Load())), node) return node def visit_Constant(self, node): if node.value == "hessian_recipes.gym_training_long": self.counts["trainer_cli"] += 1 return ast.copy_location(ast.Constant("hessian_recipes.openenv_training"), node) return node def _engine(sdk_python): source_path = Path(nemo.__file__) source = source_path.read_bytes() if hashlib.sha256(source).hexdigest() != ENGINE_SHA256: raise RuntimeError("Frozen experiment engine source changed") tree = ast.parse(source, filename=str(source_path)) adapter = _Adapt() tree = adapter.visit(tree) if adapter.counts != REWRITES: raise RuntimeError("Frozen experiment engine adaptation boundary changed") engine = ModuleType("hessian_recipes._openenv_experiment_engine") engine.__file__, engine.__package__ = str(source_path), "hessian_recipes" engine._openenv_training = training.engine(sdk_python) exec(compile(ast.fix_missing_locations(tree), str(source_path), "exec"), engine.__dict__) engine.FORMAT, engine.verify_plan = FORMAT, verify_plan engine.adaptation_counts = adapter.counts original_command = engine._Journal.command def command(journal, name, arguments): if arguments[:2] == ["-m", "hessian_recipes.openenv_training"]: arguments = [*arguments, "--sdk-python", str(sdk_python)] return original_command(journal, name, arguments) engine._Journal.command = command original_chain = engine._final_chain original_suite = engine._suite engine._frozen_shared_suite = original_suite def suite(journal, model, label, split, selection=None): if label == "baseline": stage = f"baseline-reuse-{split}" receipt = journal.root / "baseline-reuse" / f"{split}.json" if not journal.complete(stage): result = _reuse_baseline(engine, journal, model, split, selection) nemo._write(receipt, result) journal.done(stage, [receipt]) if result["status"] == "reused": manifest = Path(result["destination_manifest"]) for entry in nemo._read(manifest): name = f"baseline-{split}-{entry['kind']}-{entry['seed']}" journal.done(name, [entry["path"], entry["spec_path"], manifest, receipt]) return original_suite(journal, model, label, split, selection) engine._suite = suite def chain(journal): return original_chain(journal) | { "backend": "openenv", "nemo_origin": journal.plan["nemo_origin"], "native_pin": journal.plan["native_pin"], "sdk": journal.plan["sdk"], "source_commit": journal.plan["source_commit"], "frozen_engine_sha256": ENGINE_SHA256, "adaptation_counts": dict(REWRITES), "baseline_reuse": {split: { "path": str(journal.root / "baseline-reuse" / f"{split}.json"), "sha256": sha256_file(journal.root / "baseline-reuse" / f"{split}.json")} for split in ("dev", "test")}} engine._final_chain = chain engine._service = _service return engine def _reuse_baseline(engine, journal, model, split, selection): """Reuse a whole verified captured suite; scores never determine eligibility.""" result = {"status": "rerun", "policy": REUSE_POLICY, "split": split, "reason": "reuse disabled"} if not journal.plan["baseline_reuse"]["enabled"]: return result original, origin_sha = _nemo_plan(journal.plan["nemo_origin"]["path"]) source_root = Path(original["root"]) source_manifest = Path(original["evaluation_plan"]).parent / "specifications/baseline" / split / "manifest.json" try: source_selection = source_root / "selection/selection.json" if split == "test" else None if source_selection: evaluation._verified_selection(original["evaluation_plan"], source_selection) selection_sha = sha256_file(source_selection) if source_selection else None protocol, digest = evaluation.verify_plan(original["evaluation_plan"]) captured = evaluation._load_suite(protocol, digest, source_manifest, split, selection_digest=selection_sha) evaluation._compatible([captured]) if captured["model"] != protocol["baseline"]: raise ValueError("Only original SFT baseline evidence may be reused") for report in captured["reports"].values(): provenance = report["provenance"] if (any(importlib.metadata.version(name) != version for name, version in provenance["packages"].items()) or ("python" in provenance and provenance["python"] != platform.python_version())): raise ValueError("Captured baseline package/Python identity differs") entries = nemo._read(source_manifest) state = nemo._read(source_root / "experiment-status.json") if state["plan_sha256"] != origin_sha: raise ValueError("Source baseline journal belongs to another NeMo plan") snapshots = [] for row in entries: completed = state["completed"][f"baseline-{split}-{row['kind']}-{row['seed']}"] for file in (row["path"], row["spec_path"], str(source_manifest)): if completed.get(str(Path(file).resolve())) != sha256_file(file): raise ValueError("Source baseline stage is not durably completed") report_bytes, spec_bytes = Path(row["path"]).read_bytes(), Path(row["spec_path"]).read_bytes() report_sha = hashlib.sha256(report_bytes).hexdigest() expected_report = next(item["sha256"] for item in captured["evidence"] if (item["kind"], item["seed"]) == (row["kind"], row["seed"])) if report_sha != expected_report or hashlib.sha256(spec_bytes).hexdigest() != row["spec_sha256"]: raise ValueError("Source baseline snapshot changed during verification") snapshots.append((row, report_bytes, spec_bytes, report_sha)) except (OSError, ValueError, KeyError, TypeError, importlib.metadata.PackageNotFoundError) as error: return result | {"reason": f"No complete compatible source suite: {type(error).__name__}: {error}"} destination = Path(journal.plan["evaluation_plan"]).parent / "specifications/baseline" / split / "manifest.json" if not destination.exists(): engine._archive(journal.root, destination.parent) evaluation.write_specs(journal.plan["evaluation_plan"], model, "baseline", split, selection_path=selection) own_entries = nemo._read(destination) engine._validate_suite_manifest(journal, model, "baseline", split, selection, own_entries) own = {(row["kind"], row["seed"]): row for row in own_entries} # Validate the entire generated specification set before copying any report. if any(Path(own[row["kind"], row["seed"]]["spec_path"]).read_bytes() != spec for row, _, spec, _ in snapshots): return result | {"reason": "Generated baseline specifications are not byte-identical"} records = [] for row, report_bytes, _, report_sha in snapshots: target = Path(own[row["kind"], row["seed"]]["path"]) if target.exists() and target.read_bytes() != report_bytes: raise ValueError("Conflicting interrupted baseline evidence; refusing overwrite") target.parent.mkdir(parents=True, exist_ok=True) target.write_bytes(report_bytes) records.append({"kind": row["kind"], "seed": row["seed"], "source_report": row["path"], "report_sha256": report_sha, "source_spec": row["spec_path"], "spec_sha256": row["spec_sha256"]}) own_plan, own_digest = evaluation.verify_plan(journal.plan["evaluation_plan"]) copied = evaluation._load_suite(own_plan, own_digest, destination, split, selection_digest=sha256_file(selection) if selection else None) evaluation._compatible([captured, copied]) return result | {"status": "reused", "reason": "Complete source suite and identical effective specifications verified", "nemo_origin_sha256": origin_sha, "source_manifest": str(source_manifest), "source_manifest_sha256": sha256_file(source_manifest), "source_selection_sha256": selection_sha, "destination_selection_sha256": sha256_file(selection) if selection else None, "destination_manifest": str(destination), "reports": records, "timing_note": "Baseline generation timing/usage describe the original NeMo baseline capture, not new OpenEnv-arm compute."} def _launcher(): path = ROOT / "examples/agent_gyms/openenv_smoke.py" spec = importlib.util.spec_from_file_location("hessian_recipes._openenv_launcher", path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module.launch_openenv_server @contextmanager def _service(root, sdk_python): plan = verify_plan(Path(root) / "experiment-plan.json") with _launcher()(sdk_python, Path(root) / "services" / datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%f")) as owned: from .openenv_client import fetch_gym_identity identity = fetch_gym_identity(owned["url"]) if _stable_identity(identity) != plan["native_pin"]["native_identity"]: raise RuntimeError("Native OpenEnv runtime differs from the CPU-qualified identity") yield owned["url"], identity def run(plan_path, *, resume=False, qualification_only=False): plan = verify_plan(plan_path) from .openenv_client import sdk_identity if sdk_identity(plan["sdk"]["python"]) != plan["sdk"]["runtime"]: raise ValueError("Installed OpenEnv SDK changed since preparation") return _engine(plan["sdk"]["python"]).run( plan_path, plan["sdk"]["python"], resume=resume, qualification_only=qualification_only) def main(): parser = argparse.ArgumentParser(description=__doc__) commands = parser.add_subparsers(dest="command", required=True) prep = commands.add_parser("prepare") for key in ("nemo-plan", "output", "sdk-python", "native-pin"): prep.add_argument("--" + key, required=True, type=Path) prep.add_argument("--rerun-baseline", action="store_true") prep.add_argument("--budget-carry", type=Path) execute = commands.add_parser("run") execute.add_argument("--plan", required=True, type=Path) execute.add_argument("--resume", action="store_true") execute.add_argument("--qualification-only", action="store_true") args = parser.parse_args() result = (prepare(args.nemo_plan, args.output, args.sdk_python, args.native_pin, reuse_baseline=not args.rerun_baseline, budget_carry=args.budget_carry) if args.command == "prepare" else run(args.plan, resume=args.resume, qualification_only=args.qualification_only)) print(json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False)) if __name__ == "__main__": main()