From 99cc68148a3706a41879817c526d1dda115a12cc Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:29:00 +0000 Subject: [PATCH 001/233] Harbor layer for Miles/verl images: hold base packages, allow listed moves, assert at start Harbor is imported in-process by both backends, so it shares the engine environment. docker/harbor-layer/versions.json lists the only package moves Harbor's bounds force (Miles: openai, pyyaml, websockets, importlib-metadata; verl: websockets, importlib-metadata) and pins litellm 1.95.0 so filelock (Miles) and aiohttp/yarl (verl) stay at the base version. harbor_layer.py builds the uv constraints, refuses any other change at build time and writes the start lock; yeto.launch.image_check asserts it from verl agentic_main and on the first real Harbor Trial load (Modal skips the image entrypoint). Co-Authored-By: Claude Opus 5.5 --- docker/harbor-layer/harbor_layer.py | 183 ++++++++++++++++++++ docker/harbor-layer/versions.json | 86 ++++++++++ tests/test_harbor_image_layer.py | 238 ++++++++++++++++++++++++++ yeto/launch/image_check.py | 57 ++++++ yeto/rl/adapters/verl/agentic_main.py | 3 + yeto/rl/harness/harbor/image_layer.py | 63 +++++++ yeto/rl/harness/harbor/trial.py | 3 + 7 files changed, 633 insertions(+) create mode 100644 docker/harbor-layer/harbor_layer.py create mode 100644 docker/harbor-layer/versions.json create mode 100644 tests/test_harbor_image_layer.py create mode 100644 yeto/launch/image_check.py create mode 100644 yeto/rl/harness/harbor/image_layer.py diff --git a/docker/harbor-layer/harbor_layer.py b/docker/harbor-layer/harbor_layer.py new file mode 100644 index 000000000..255b562f9 --- /dev/null +++ b/docker/harbor-layer/harbor_layer.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +"""Install Harbor into an engine image without moving the engine's packages (S26 E0.4). + +Harbor runs in the engine's own Python (yeto's HarborAgentLoop and the Miles +Harbor example import it in-process), so it shares one environment with +SGLang / Miles / vLLM / verl / Ray. The expected result of that install is +kept in ``versions.json`` next to this file: + +* every package of the base image is held at its base version (uv + constraint), except the ones listed under ``backends..changes``; Harbor's + dependency bounds force those to move, each to one exact version; +* ``pins`` fixes Harbor-only packages whose choice decides the moves above. + +Commands (run inside the image; stdlib only): + +* ``constraints --backend B BEFORE``: the uv constraint file, to stdout; +* ``check --backend B BEFORE AFTER``: exit 1 when the install moved, removed + or downgraded anything ``versions.json`` does not allow; +* ``lock --backend B AFTER --out LOCK``: the start lock that + ``docker/image-assert/assert_image_versions.py`` checks at container start. + +BEFORE / AFTER are ``uv pip freeze`` outputs. +""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from pathlib import Path + +SPEC_FILE = Path(__file__).with_name("versions.json") +_PINNED = re.compile(r"^([A-Za-z0-9][A-Za-z0-9._-]*)==(\S+)$") +_DIRECT = re.compile(r"^([A-Za-z0-9][A-Za-z0-9._-]*)\s*@\s*(\S+)$") +# Module each asserted package is imported as, where it is light enough to +# import at container start (no CUDA context, no engine side effects). +START_IMPORTS = { + "harbor": "harbor", "boat-sdk": "boat_sdk", "openai": "openai", "pyyaml": "yaml", + "websockets": "websockets", "importlib-metadata": "importlib_metadata", "filelock": "filelock", + "litellm": None, # imports slowly and fetches a cost table; version only +} + + +def canonical(name: str) -> str: + """PEP 503 name: ``PyYAML`` / ``importlib_metadata`` -> ``pyyaml`` / ``importlib-metadata``.""" + return re.sub(r"[-_.]+", "-", name).lower() + + +def parse_freeze(text: str) -> dict[str, str]: + """``uv pip freeze`` -> ``{name: version}``; a direct (VCS, path) install maps to ``@ ``. + + Editable lines (``-e ...``) carry no name and are skipped: the engines + install their own sources editable and Harbor's install leaves them alone + (``check`` would not see such a move; the start lock's commits do). + """ + found: dict[str, str] = {} + for raw in text.splitlines(): + line = raw.split(" #", 1)[0].strip() + if not line or line.startswith(("#", "-")): + continue + if match := _PINNED.match(line): + found[canonical(match.group(1))] = match.group(2) + elif match := _DIRECT.match(line): + found[canonical(match.group(1))] = "@ " + match.group(2) + return found + + +def load_spec(backend: str, spec_file: Path = SPEC_FILE) -> tuple[dict, dict]: + spec = json.loads(spec_file.read_text()) + if backend not in spec["backends"]: + raise SystemExit(f"[harbor-layer] unknown backend {backend!r}; known: {sorted(spec['backends'])}") + return spec, spec["backends"][backend] + + +def constraints(before: dict[str, str], spec: dict, backend_spec: dict) -> list[str]: + """Hold every pinned base package, except the allowed changes (at their layer version) and the pins.""" + changes = {canonical(n): c for n, c in backend_spec["changes"].items()} + pins = {canonical(n): p for n, p in spec["pins"].items()} + lines = [f"{name}=={version}" for name, version in sorted(before.items()) + if not version.startswith("@") and name not in changes and name not in pins] + lines += [f"{name}=={change['layer']}" for name, change in sorted(changes.items())] + lines += [f"{name}=={pin['version']}" for name, pin in sorted(pins.items())] + return lines + + +def check(before: dict[str, str], after: dict[str, str], spec: dict, backend_spec: dict) -> list[str]: + """Every way ``after`` differs from what ``versions.json`` allows, one line each (empty: as expected).""" + changes = {canonical(n): c for n, c in backend_spec["changes"].items()} + problems: list[str] = [] + for name, change in sorted(changes.items()): + if before.get(name) != change["base"]: + problems.append(f"{name}: base image has {before.get(name)}, versions.json expects base " + f"{change['base']} (the base image moved; review the change and update versions.json)") + if after.get(name) != change["layer"]: + problems.append(f"{name}: installed {after.get(name)}, versions.json allows only {change['layer']}") + for name, version in sorted(before.items()): + if name in changes: + continue + if name not in after: + problems.append(f"{name}: removed by the Harbor install (base {version})") + elif after[name] != version: + problems.append(f"{name}: moved {version} -> {after[name]} by the Harbor install, " + "not allowed by versions.json") + for name, pin in sorted(spec["pins"].items()): + if after.get(canonical(name)) != pin["version"]: + problems.append(f"{canonical(name)}: installed {after.get(canonical(name))}, pinned {pin['version']}") + harbor = after.get("harbor", "") + if not harbor.endswith("@" + spec["harbor"]["commit"]): + problems.append(f"harbor: installed {harbor or 'nothing'}, pinned commit {spec['harbor']['commit']}") + if after.get("boat-sdk") != spec["boat_sdk"]: + problems.append(f"boat-sdk: installed {after.get('boat-sdk')}, pinned {spec['boat_sdk']}") + return problems + + +def start_lock(after: dict[str, str], spec: dict, backend: str, *, cpu_arch: str = "x86_64") -> dict: + """The lock ``assert_image_versions.py`` checks at start: engine packages at their built version, + the allowed changes and pins at their layer version, Harbor and boat-sdk importable.""" + backend_spec = spec["backends"][backend] + packages: dict[str, dict] = {} + for name in spec["assert_packages"]: + version = after.get(canonical(name)) + if version is not None and not version.startswith("@"): + packages[canonical(name)] = {"version": version} + expected = {canonical(n): c["layer"] for n, c in backend_spec["changes"].items()} + expected |= {canonical(n): p["version"] for n, p in spec["pins"].items()} + expected["filelock"] = after.get("filelock") + expected["harbor"] = spec["harbor"]["version"] + expected["boat-sdk"] = spec["boat_sdk"] + for name, version in expected.items(): + if version is None: + continue + row = {"version": version} + if START_IMPORTS.get(name): + row["import"] = START_IMPORTS[name] + packages[name] = row + return { + "image": backend_spec["image"], + "cpu_arch": cpu_arch, + "device_family": "cuda", + "python": None, + "packages": packages, + "commits": {}, + "required_files": [], + "decisions": [f"{n}: {c['base']} -> {c['layer']} ({c['why']})" for n, c in backend_spec["changes"].items()] + + [f"{n}: pinned {p['version']} ({p['why']})" for n, p in spec["pins"].items()], + } + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("command", choices=("constraints", "check", "lock")) + ap.add_argument("--backend", required=True) + ap.add_argument("freeze", type=Path, nargs="+", help="BEFORE (constraints, check) and/or AFTER (check, lock)") + ap.add_argument("--out", type=Path, help="lock: where to write the start lock") + ap.add_argument("--spec", type=Path, default=SPEC_FILE) + args = ap.parse_args(argv) + spec, backend_spec = load_spec(args.backend, args.spec) + freezes = [parse_freeze(path.read_text()) for path in args.freeze] + if args.command == "constraints": + print("\n".join(constraints(freezes[0], spec, backend_spec))) + return 0 + if args.command == "check": + if len(freezes) != 2: + ap.error("check takes BEFORE and AFTER") + problems = check(freezes[0], freezes[1], spec, backend_spec) + for line in problems: + print(f"[harbor-layer] FAIL {line}", file=sys.stderr) + if problems: + print(f"[harbor-layer] {len(problems)} problem(s): the Harbor install changed the {args.backend} " + f"image beyond {args.spec.name}; refusing the image", file=sys.stderr) + return 1 + print(f"[harbor-layer] ok: {args.backend} image changed only as {args.spec.name} allows", file=sys.stderr) + return 0 + if args.out is None: + ap.error("lock needs --out") + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(start_lock(freezes[-1], spec, args.backend), indent=2) + "\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/docker/harbor-layer/versions.json b/docker/harbor-layer/versions.json new file mode 100644 index 000000000..831884235 --- /dev/null +++ b/docker/harbor-layer/versions.json @@ -0,0 +1,86 @@ +{ + "harbor": { + "requirement": "harbor @ git+https://github.com/harbor-framework/harbor@8f88560ce9daae9b4cc0f6ff7bac8779ae9cfa15", + "commit": "8f88560ce9daae9b4cc0f6ff7bac8779ae9cfa15", + "version": "0.20.0" + }, + "boat_sdk": "1.10.0", + "uv": "0.9.5", + "assert_packages": [ + "torch", + "triton", + "transformers", + "tokenizers", + "safetensors", + "accelerate", + "peft", + "ray", + "numpy", + "pydantic", + "pydantic-core", + "huggingface-hub", + "aiohttp", + "yarl", + "fastapi", + "uvicorn", + "sglang", + "sglang-router", + "sgl-kernel", + "flashinfer-python", + "megatron-core", + "transformer-engine", + "miles", + "vllm", + "verl", + "tensordict", + "flash-attn" + ], + "backends": { + "miles": { + "image": "yeto-miles-ports+harbor", + "changes": { + "openai": { + "base": "2.6.1", + "layer": "2.54.0", + "why": "litellm >= 1.83.14 (Harbor floor) needs openai >= 2.20; SGLang 2fa88018 pins openai==2.6.1" + }, + "pyyaml": { + "base": "6.0.1", + "layer": "6.0.3", + "why": "Harbor 0.20.0 needs pyyaml >= 6.0.2; SGLang 2fa88018 pins PyYAML==6.0.1" + }, + "websockets": { + "base": "17.1", + "layer": "15.0.1", + "why": "supabase -> realtime >= 2.28.2 needs websockets < 16" + }, + "importlib-metadata": { + "base": "9.0.1", + "layer": "8.9.0", + "why": "litellm >= 1.84.0 needs importlib-metadata < 9.0" + } + } + }, + "verl": { + "image": "verl-build+harbor", + "changes": { + "websockets": { + "base": "16.0", + "layer": "15.0.1", + "why": "supabase -> realtime >= 2.28.2 needs websockets < 16" + }, + "importlib-metadata": { + "base": "9.0.0", + "layer": "8.9.0", + "why": "litellm >= 1.84.0 needs importlib-metadata < 9.0" + } + } + } + }, + "pins": { + "litellm": { + "version": "1.95.0", + "why": "newest litellm that resolves with filelock 4.0.0 (Miles base), aiohttp 3.14.1 (verl base) and openai 2.54.0 (uv pip compile, 2026-10-20); 1.106.0 needs filelock < 4 and aiohttp >= 3.14.2" + } + } +} diff --git a/tests/test_harbor_image_layer.py b/tests/test_harbor_image_layer.py new file mode 100644 index 000000000..9e4f27400 --- /dev/null +++ b/tests/test_harbor_image_layer.py @@ -0,0 +1,238 @@ +"""S26 E0.4: the Harbor layer of the Miles / verl images and the start check.""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path + +import pytest + +from yeto.launch import image_check +from yeto.rl.harness.harbor import image_layer +from yeto.rl.harness.harbor.pins import HARBOR_COMMIT, HARBOR_VERSION + +ROOT = Path(__file__).resolve().parents[1] +LAYER = ROOT / "docker" / "harbor-layer" +ASSERT_SCRIPT = ROOT / "docker" / "image-assert" / "assert_image_versions.py" + + +def _load(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +layer = _load(LAYER / "harbor_layer.py", "harbor_layer") +SPEC = json.loads((LAYER / "versions.json").read_text()) +HARBOR_LINE = f"harbor @ git+https://github.com/harbor-framework/harbor@{HARBOR_COMMIT}" + +# Base versions the S25 N10.9 check read from the two images (install_diff in +# SESSION25/logs/T-N109/check-miles2.log and check-verl.log). +MILES_BEFORE = """\ +filelock==4.0.0 +importlib-metadata==9.0.1 +openai==2.6.1 +PyYAML==6.0.1 +websockets==17.1 +aiohttp==3.14.5 +torch==2.11.0+cu130 +-e file:///sgl-workspace/sglang/python +""" +VERL_BEFORE = """\ +aiohttp==3.14.1 +yarl==1.24.2 +importlib_metadata==9.0.0 +openai==2.41.1 +filelock==3.29.4 +pyyaml==6.0.3 +websockets==16.0 +torch==2.13.0+cu130 +vllm==0.29.0 +""" + + +def _after(before: str, backend: str, **moves: str) -> dict[str, str]: + """The freeze an install exactly as versions.json allows would leave (plus ``moves``).""" + after = layer.parse_freeze(before) + for name, change in SPEC["backends"][backend]["changes"].items(): + after[name] = change["layer"] + after |= {"litellm": SPEC["pins"]["litellm"]["version"], "boat-sdk": SPEC["boat_sdk"], + "harbor": "@ git+https://github.com/harbor-framework/harbor@" + HARBOR_COMMIT, + "supabase": "2.33.0"} + after |= {layer.canonical(k): v for k, v in moves.items()} + return after + + +# --- versions.json ----------------------------------------------------------- + +def test_versions_json_pins_the_same_harbor_as_the_yeto_pins(): + assert SPEC["harbor"]["commit"] == HARBOR_COMMIT and SPEC["harbor"]["version"] == HARBOR_VERSION + assert SPEC["harbor"]["requirement"] == HARBOR_LINE + + +def test_versions_json_records_the_moves_n109_saw_and_holds_filelock(): + miles, verl = (SPEC["backends"][b]["changes"] for b in ("miles", "verl")) + assert {n: (c["base"], c["layer"]) for n, c in miles.items()} == { + "openai": ("2.6.1", "2.54.0"), "pyyaml": ("6.0.1", "6.0.3"), + "websockets": ("17.1", "15.0.1"), "importlib-metadata": ("9.0.1", "8.9.0")} + assert {n: (c["base"], c["layer"]) for n, c in verl.items()} == { + "websockets": ("16.0", "15.0.1"), "importlib-metadata": ("9.0.0", "8.9.0")} + assert "filelock" not in miles and "aiohttp" not in verl # held at the base version + assert all(c["why"] for c in [*miles.values(), *verl.values(), *SPEC["pins"].values()]) + + +# --- freeze parsing and constraints ----------------------------------------- + +def test_parse_freeze_canonicalizes_names_and_skips_editables(): + parsed = layer.parse_freeze(MILES_BEFORE + HARBOR_LINE + "\n# comment\n") + assert parsed["pyyaml"] == "6.0.1" and parsed["importlib-metadata"] == "9.0.1" + assert parsed["harbor"] == "@ git+https://github.com/harbor-framework/harbor@" + HARBOR_COMMIT + assert not any("sglang" in name for name in parsed) + + +def test_constraints_hold_every_base_package_except_the_allowed_moves(): + spec, backend = layer.load_spec("miles") + lines = layer.constraints(layer.parse_freeze(MILES_BEFORE), spec, backend) + assert "filelock==4.0.0" in lines and "torch==2.11.0+cu130" in lines and "aiohttp==3.14.5" in lines + assert "openai==2.54.0" in lines and "openai==2.6.1" not in lines + assert "pyyaml==6.0.3" in lines and "litellm==1.95.0" in lines + assert len([line for line in lines if line.startswith("openai==")]) == 1 + + +def test_unknown_backend_is_refused(): + with pytest.raises(SystemExit, match="unknown backend"): + layer.load_spec("trl") + with pytest.raises(ValueError, match="unknown backend"): + image_layer.layer_commands("trl", "python") + + +# --- build check ------------------------------------------------------------- + +@pytest.mark.parametrize("backend, before", [("miles", MILES_BEFORE), ("verl", VERL_BEFORE)]) +def test_check_accepts_an_install_that_moved_only_what_versions_json_allows(backend, before): + spec, backend_spec = layer.load_spec(backend) + assert layer.check(layer.parse_freeze(before), _after(before, backend), spec, backend_spec) == [] + + +def test_check_refuses_the_unlisted_filelock_move_n109_made_on_miles(): + spec, backend = layer.load_spec("miles") + problems = layer.check(layer.parse_freeze(MILES_BEFORE), _after(MILES_BEFORE, "miles", filelock="3.32.7"), + spec, backend) + assert problems == ["filelock: moved 4.0.0 -> 3.32.7 by the Harbor install, not allowed by versions.json"] + + +def test_check_refuses_a_removed_package_a_wrong_layer_version_and_a_moved_base(): + spec, backend = layer.load_spec("verl") + before = layer.parse_freeze(VERL_BEFORE.replace("websockets==16.0", "websockets==16.1")) + after = _after(VERL_BEFORE, "verl", **{"importlib-metadata": "8.5.0"}) + del after["vllm"] + problems = layer.check(before, after, spec, backend) + assert any(p.startswith("vllm: removed") for p in problems) + assert any(p.startswith("importlib-metadata: installed 8.5.0, versions.json allows only 8.9.0") for p in problems) + assert any(p.startswith("websockets: base image has 16.1") for p in problems) + + +def test_check_refuses_a_wrong_harbor_commit_or_litellm_pin(): + spec, backend = layer.load_spec("verl") + after = _after(VERL_BEFORE, "verl", litellm="1.106.0", harbor="@ git+https://x@deadbeef") + problems = layer.check(layer.parse_freeze(VERL_BEFORE), after, spec, backend) + assert {p.split(":")[0] for p in problems} == {"litellm", "harbor"} + + +def test_check_cli_exits_1_with_one_line_per_problem(tmp_path, capsys): + before, after = tmp_path / "before.txt", tmp_path / "after.txt" + before.write_text(MILES_BEFORE) + moved = _after(MILES_BEFORE, "miles", filelock="3.32.7") + after.write_text("\n".join(f"{n}=={v}" if not v.startswith("@") else f"{n} {v}" for n, v in moved.items())) + assert layer.main(["check", "--backend", "miles", str(before), str(after)]) == 1 + err = capsys.readouterr().err + assert "[harbor-layer] FAIL filelock: moved 4.0.0 -> 3.32.7" in err and "refusing the image" in err + + +# --- start lock and start check ---------------------------------------------- + +def test_start_lock_asserts_engine_packages_and_the_layer_versions(): + lock = layer.start_lock(_after(VERL_BEFORE, "verl"), SPEC, "verl") + packages = lock["packages"] + assert packages["torch"] == {"version": "2.13.0+cu130"} and packages["vllm"] == {"version": "0.29.0"} + assert packages["websockets"] == {"version": "15.0.1", "import": "websockets"} + assert packages["filelock"] == {"version": "3.29.4", "import": "filelock"} + assert packages["harbor"] == {"version": HARBOR_VERSION, "import": "harbor"} + assert packages["litellm"] == {"version": "1.95.0"} + assert lock["image"] == "verl-build+harbor" and lock["decisions"] + + +def _matching_lock(tmp_path: Path, **override: dict) -> Path: + checker = _load(ASSERT_SCRIPT, "assert_image_versions") + packages = {"pytest": {"version": pytest.__version__, "import": "pytest"}} | override + lock = {"image": "test", "cpu_arch": checker.platform.machine(), "python": None, "packages": packages} + path = tmp_path / "versions.lock.json" + path.write_text(json.dumps(lock)) + return path + + +def test_check_image_passes_on_a_matching_environment(tmp_path): + rows = image_check.check_image(_matching_lock(tmp_path), ASSERT_SCRIPT) + assert {row["status"] for row in rows} == {"ok"} + + +def test_check_image_names_every_failed_check(tmp_path): + lock = _matching_lock(tmp_path, openai={"version": "0.0.1"}, **{"no-such-dist": {"version": "1"}}) + with pytest.raises(image_check.ImageVersionMismatch) as err: + image_check.check_image(lock, ASSERT_SCRIPT) + text = str(err.value) + assert "2 failed check(s)" in text and "package no-such-dist: not installed" in text + assert "package openai" in text and "docker/harbor-layer/versions.json" in text + + +def test_check_image_at_start_does_nothing_outside_a_yeto_image(tmp_path, monkeypatch): + monkeypatch.setenv(image_check.LOCK_ENV, str(tmp_path / "absent.json")) + image_check.check_image_at_start.cache_clear() + try: + assert image_check.check_image_at_start() == () + finally: + image_check.check_image_at_start.cache_clear() + + +def test_check_image_at_start_raises_when_the_image_has_a_lock(tmp_path, monkeypatch): + lock = _matching_lock(tmp_path, **{"no-such-dist": {"version": "1"}}) + monkeypatch.setenv(image_check.LOCK_ENV, str(lock)) + monkeypatch.setattr(image_check, "IMAGE_DIR", ASSERT_SCRIPT.parent) + image_check.check_image_at_start.cache_clear() + try: + with pytest.raises(image_check.ImageVersionMismatch, match="no-such-dist"): + image_check.check_image_at_start() + finally: + image_check.check_image_at_start.cache_clear() + + +# --- image build commands ---------------------------------------------------- + +def test_layer_commands_constrain_check_lock_and_assert_in_that_order(): + commands = image_layer.layer_commands("miles", "/usr/bin/python") + joined = "\n".join(commands) + assert f"-c {image_layer.IMAGE_LAYER_DIR}/constraints.txt '{HARBOR_LINE}' 'boat-sdk==1.10.0'" in joined + steps = ["constraints --backend miles", "uv pip install", "check --backend miles", + f"--out {image_layer.START_LOCK}", "assert_image_versions.py"] + positions = [next(i for i, c in enumerate(commands) if step in c) for step in steps] + assert positions == sorted(positions) + + +def test_modal_harbor_layer_ships_the_layer_files_and_the_checker(): + calls: list[tuple] = [] + + class FakeImage: + def __getattr__(self, name): + def record(*args, **kwargs): + calls.append((name, args, kwargs)) + return self + return record + + image_layer.modal_harbor_layer(FakeImage(), "verl", "/workspace/verl/.venv/bin/python") + names = [c[0] for c in calls] + assert names == ["add_local_dir", "add_local_file", "run_commands"] + assert calls[0][1] == (str(LAYER), image_layer.IMAGE_LAYER_DIR) and calls[0][2] == {"copy": True} + assert calls[1][1][1] == f"{image_layer.IMAGE_DIR}/assert_image_versions.py" + assert any("check --backend verl" in c for c in calls[2][1]) diff --git a/yeto/launch/image_check.py b/yeto/launch/image_check.py new file mode 100644 index 000000000..0e9d628d2 --- /dev/null +++ b/yeto/launch/image_check.py @@ -0,0 +1,57 @@ +"""Start check of a yeto engine image from Python (S26 E0.4). + +A Docker image runs ``docker/image-assert/entrypoint.sh`` before its +command. Modal does not run the image entrypoint, so the Harbor-layer images +(:mod:`yeto.rl.harness.harbor.image_layer`) are checked from yeto's own entry +points through :func:`check_image_at_start`: the same stdlib checker, against +the lock the image build wrote. Outside a yeto image (no lock file: CPU +tests, a developer machine) the check does nothing. +""" + +from __future__ import annotations + +import importlib.util +import json +import os +from functools import cache +from pathlib import Path + +IMAGE_DIR = Path("/opt/yeto/image") +LOCK_ENV = "YETO_IMAGE_LOCK" + + +class ImageVersionMismatch(RuntimeError): + """The running image differs from its version lock; the run must not start.""" + + +def _checker(path: Path): + spec = importlib.util.spec_from_file_location("yeto_image_assert", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def check_image(lock_path: Path, checker_path: Path) -> list[dict]: + """Rows of the check; raise :class:`ImageVersionMismatch` naming every failed row.""" + checker = _checker(checker_path) + lock = json.loads(lock_path.read_text()) + report = checker.run(lock, commits_file=checker_path.parent / "commits.json") + if report.failed: + lines = "\n".join(f" {row['check']}: {row['detail']}" for row in report.failed) + raise ImageVersionMismatch( + f"image {lock.get('image')} does not match its version lock {lock_path} " + f"({len(report.failed)} failed check(s)):\n{lines}\n" + "Rebuild the image, or review and update docker/harbor-layer/versions.json.") + return report.rows + + +@cache +def check_image_at_start() -> tuple[dict, ...]: + """Check the image once per process; empty when the process does not run in a yeto image.""" + lock_path = Path(os.environ.get(LOCK_ENV) or IMAGE_DIR / "versions.lock.json") + checker_path = IMAGE_DIR / "assert_image_versions.py" + if not lock_path.is_file(): + return () + if not checker_path.is_file(): + raise ImageVersionMismatch(f"{lock_path} exists but the checker {checker_path} is missing") + return tuple(check_image(lock_path, checker_path)) diff --git a/yeto/rl/adapters/verl/agentic_main.py b/yeto/rl/adapters/verl/agentic_main.py index ea7d69e31..dd026ba26 100644 --- a/yeto/rl/adapters/verl/agentic_main.py +++ b/yeto/rl/adapters/verl/agentic_main.py @@ -93,6 +93,9 @@ def runtime_env(environ: dict[str, str] | None = None) -> dict: def main(argv: list[str] | None = None) -> None: + from yeto.launch.image_check import check_image_at_start + + check_image_at_start() # before Ray: a mismatched image stops here with the failed checks import ray import verl.trainer.main_ppo as main_ppo diff --git a/yeto/rl/harness/harbor/image_layer.py b/yeto/rl/harness/harbor/image_layer.py new file mode 100644 index 000000000..30ec4d1b6 --- /dev/null +++ b/yeto/rl/harness/harbor/image_layer.py @@ -0,0 +1,63 @@ +"""The Harbor layer of the Miles and verl engine images (S26 E0.4). + +Harbor is imported in-process by both backends (verl ``HarborAgentLoop``, +Miles' Harbor example and ``records_generate``), so it cannot live in its own +venv without a process bridge. It is installed into the engine's Python under +a uv constraint file that holds every base package; only the moves listed in +``docker/harbor-layer/versions.json`` are allowed. The build then checks the +result and writes the start lock ``/opt/yeto/image/versions.lock.json`` that +:func:`yeto.launch.image_check.check_image_at_start` asserts before training. +""" + +from __future__ import annotations + +from pathlib import Path + +from .pins import HARBOR_COMMIT + +LAYER_DIR = Path(__file__).resolve().parents[4] / "docker" / "harbor-layer" +ASSERT_SCRIPT = LAYER_DIR.parent / "image-assert" / "assert_image_versions.py" +IMAGE_LAYER_DIR = "/opt/yeto/harbor-layer" +IMAGE_DIR = "/opt/yeto/image" +START_LOCK = f"{IMAGE_DIR}/versions.lock.json" +BACKENDS = ("miles", "verl") + + +def layer_commands(backend: str, python: str, *, uv_version: str = "0.9.5", + harbor_requirement: str | None = None, boat_sdk: str = "1.10.0") -> list[str]: + """Shell commands (one image layer) that install Harbor + boat-sdk into ``python``'s environment. + + ``python`` is the engine interpreter (verl: the venv's python; Miles: the + image's ``python``). Any failed step fails the image build. + """ + if backend not in BACKENDS: + raise ValueError(f"unknown backend {backend!r}; known: {list(BACKENDS)}") + requirement = harbor_requirement or f"harbor @ git+https://github.com/harbor-framework/harbor@{HARBOR_COMMIT}" + tool = f"{python} {IMAGE_LAYER_DIR}/harbor_layer.py" + before, after = f"{IMAGE_LAYER_DIR}/freeze-before.txt", f"{IMAGE_LAYER_DIR}/freeze-after.txt" + return [ + f"command -v uv || python3 -m pip install uv=={uv_version}", + f"uv pip freeze --python {python} > {before}", + f"{tool} constraints --backend {backend} {before} > {IMAGE_LAYER_DIR}/constraints.txt", + f"uv pip install --python {python} -c {IMAGE_LAYER_DIR}/constraints.txt '{requirement}' 'boat-sdk=={boat_sdk}'", + f"uv pip freeze --python {python} > {after}", + f"{tool} check --backend {backend} {before} {after}", + f"{tool} lock --backend {backend} {after} --out {START_LOCK}", + f"{python} {IMAGE_DIR}/assert_image_versions.py {START_LOCK} --report {IMAGE_LAYER_DIR}/build-assert.json", + ] + + +def modal_harbor_layer(image, backend: str, python: str): + """``image`` (a ``modal.Image``) plus the Harbor layer; Modal ignores the Docker entrypoint, + so the start check runs from yeto's entry points instead (``check_image_at_start``).""" + import json + + spec = json.loads((LAYER_DIR / "versions.json").read_text()) + if spec["harbor"]["commit"] != HARBOR_COMMIT: + raise ValueError(f"versions.json pins Harbor {spec['harbor']['commit']}, pins.py {HARBOR_COMMIT}") + return (image + .add_local_dir(str(LAYER_DIR), IMAGE_LAYER_DIR, copy=True) + .add_local_file(str(ASSERT_SCRIPT), f"{IMAGE_DIR}/assert_image_versions.py", copy=True) + .run_commands(*layer_commands(backend, python, uv_version=spec["uv"], + harbor_requirement=spec["harbor"]["requirement"], + boat_sdk=spec["boat_sdk"]))) diff --git a/yeto/rl/harness/harbor/trial.py b/yeto/rl/harness/harbor/trial.py index f506d57e3..fa1653c0d 100644 --- a/yeto/rl/harness/harbor/trial.py +++ b/yeto/rl/harness/harbor/trial.py @@ -174,6 +174,9 @@ def __init__(self, build_config: TrialConfigBuilder, *, trial_cls: Any = None) - def _trial(self) -> Any: if self._trial_cls is None: + from yeto.launch.image_check import check_image_at_start + + check_image_at_start() # the image's Harbor layer must match docker/harbor-layer (S26 E0.4) from harbor.trial.trial import Trial self._trial_cls = Trial From 7faeaae13ad03ab74c2a836937b0b02b0f7dd05e Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:44:59 +0000 Subject: [PATCH 002/233] Remove the self-written Codex harness Harbor (task loading, sandbox via the boat provider, agent, verifier) plus the records proxy replace it. Removed: yeto/rl/harness/codex (agent, sandbox broker, TB2 provider, tool exec, grading, Codex bridge), the Miles codex glue, SecRLEnv retry, the RLF family, the TB2 reward-env adapter and Modal reward sandbox, the old eval runners, the launcher Codex bundle/contract path, the SSH-harness Codex/SecRLEnv daemon path, and their tests. Kept or moved: the Responses SSE writer into the records proxy, refresh_segment_counts into yeto.rl.compaction, the SecRLEnv outcome check into yeto.rl.secrlenv_outcome (trajectory evidence of that kind still verifies), the CompactionRL launch check into yeto.launch.harness_env. Removed Codex/SecRLEnv launcher flags are refused. The last commit with the old harness is tagged s26-pre-harness-removal. Co-Authored-By: Claude Opus 5.5 --- docs/RLF_CODEX_RUN.md | 177 -- scripts/fetch_codex_bundle.py | 184 -- tests/decoupling_golden.py | 43 - .../generate_signed_era.py | 2 +- tests/golden/decoupling/codex_harness.json | 273 -- tests/import_boundary_allowlist.txt | 5 - tests/test_codex_bridge_compaction.py | 505 ---- tests/test_codex_compaction_wiring.py | 356 --- tests/test_codex_end_kind.py | 83 - tests/test_codex_evaluate_time.py | 52 - tests/test_codex_head_staging.py | 40 - tests/test_harness_codex_exec_replay.py | 89 - tests/test_harness_codex_openenv.py | 935 ------- tests/test_harness_codex_preflight_member.py | 106 - tests/test_harness_session_protocol.py | 88 - tests/test_harness_tb2_provider.py | 707 ----- tests/test_import_boundaries.py | 2 +- tests/test_qwen3_cyber_profile.py | 49 +- tests/test_reasoning_loss_mask.py | 22 +- tests/test_reward_env.py | 380 +-- tests/test_rl_agentic_suspend.py | 491 ---- tests/test_rl_agentic_tool_call_args.py | 4 +- tests/test_rl_algorithm_capabilities.py | 8 +- tests/test_rl_binding_check.py | 4 +- tests/test_rl_codex_schema.py | 116 - tests/test_rl_compactionrl_counts_receipt.py | 2 +- tests/test_rl_engine_ports.py | 32 - tests/test_rl_eval_infer.py | 450 ---- tests/test_rl_eval_island.py | 59 +- tests/test_rl_ir_harness.py | 20 +- tests/test_rl_launcher.py | 97 +- tests/test_rl_launcher_codex_bundle.py | 401 --- tests/test_rl_launcher_dense_full.py | 21 - tests/test_rl_miles_adapter_config.py | 23 +- tests/test_rl_miles_carry_over.py | 4 +- tests/test_rl_neutral_rewards.py | 67 - tests/test_rl_policy_versions.py | 16 - tests/test_rl_reward_success.py | 22 - tests/test_rl_sao_streaming_runtime.py | 2 +- tests/test_rl_secrlenv_ports_retry.py | 509 ---- tests/test_rl_ssh_harness.py | 926 +------ tests/test_rl_synthetic_segments.py | 6 +- tests/test_rl_tb2_baseline.py | 127 - tests/test_rl_tbench_direct_preflight.py | 137 - tests/test_rl_trajectory_evidence.py | 2 +- tests/test_rl_trajectory_timing.py | 13 - tests/test_rlf_baseline.py | 484 ---- tests/test_rlf_bridge.py | 33 - tests/test_rlf_compaction_profile.py | 7 - tests/test_rlf_env_lease.py | 446 --- tests/test_rlf_gate_evidence.py | 151 -- tests/test_rlf_rl_gate.py | 350 --- tests/test_rlf_rollout_entry.py | 226 -- tests/test_run_flags.py | 4 +- tests/test_sandbox_isolation.py | 142 - tests/test_secret_handling_hardening.py | 31 - tests/test_secrlenv_adapter.py | 1937 ------------- tests/test_secrlenv_codex_harness.py | 1790 ------------ tests/test_tbench21_sao_diloco_plan.py | 234 -- tools/build_eval_data.py | 124 - .../probes/build_tbench21_sao_diloco_plan.py | 520 ---- tools/probes/s17_eval_island_gpu_smoke.py | 85 - tools/reward_env/build_images.py | 85 - tools/reward_env/holdout.py | 52 +- yeto/cli.py | 15 - yeto/cloud/modal_ckpt_store.py | 2 +- yeto/cloud/modal_eval_island.py | 9 +- yeto/cloud/modal_reward_env.py | 142 - yeto/cloud/modal_rlf_baseline.py | 193 -- yeto/cloud/modal_tb2_baseline.py | 129 - yeto/launch/args.py | 2 +- yeto/launch/codex.py | 314 --- yeto/launch/harness_env.py | 90 + yeto/launch/modal_islands.py | 13 +- yeto/launch/rl_args.py | 140 +- yeto/launch/rl_island_task.py | 42 +- yeto/launcher.py | 9 +- yeto/modal_runner.py | 13 - yeto/rl/__init__.py | 93 - yeto/rl/adapters/miles/binding.py | 16 +- yeto/rl/adapters/miles/carry_over.py | 4 +- yeto/rl/adapters/miles/config.py | 72 +- yeto/rl/adapters/miles/entry.py | 33 +- .../miles/harness_glue/codex_generate.py | 129 - .../miles/harness_glue/codex_openenv.py | 40 - yeto/rl/adapters/miles/in_flight.py | 29 +- yeto/rl/adapters/miles/island_entry.py | 394 +-- yeto/rl/adapters/miles/policy_age.py | 6 +- yeto/rl/adapters/miles/rewards.py | 25 +- yeto/rl/adapters/miles/rollout_meta_hook.py | 2 +- yeto/rl/adapters/miles/sampler_island.py | 2 +- yeto/rl/adapters/miles/secrlenv_miles.py | 74 - yeto/rl/adapters/miles/secrlenv_retry.py | 318 --- yeto/rl/algos/compactionrl.py | 12 +- yeto/rl/compaction.py | 37 + yeto/rl/engine/algorithm.py | 18 +- yeto/rl/engine/backends.py | 2 - yeto/rl/eval/codex_infer.py | 420 --- yeto/rl/eval/rlf_baseline.py | 714 ----- yeto/rl/eval/rlf_rl_gate.py | 500 ---- yeto/rl/eval/tb2_attempt.py | 63 - yeto/rl/eval/tb2_baseline.py | 169 -- yeto/rl/harness/README.md | 55 +- yeto/rl/harness/codex/__init__.py | 1 - yeto/rl/harness/codex/agent.py | 1070 -------- yeto/rl/harness/codex/alignment.py | 75 - yeto/rl/harness/codex/chain_judge.py | 78 - yeto/rl/harness/codex/client.py | 215 -- yeto/rl/harness/codex/codex_harness_agent.py | 2389 ----------------- .../codex/codex_openenv_agent_function.py | 681 ----- .../codex/codex_openenv_agent_worker.py | 81 - .../harness/codex/codex_openenv_generate.py | 420 --- ...codex_openenv_subprocess_agent_function.py | 616 ----- yeto/rl/harness/codex/compaction_bridge.py | 379 --- yeto/rl/harness/codex/environment.py | 181 -- yeto/rl/harness/codex/generate.py | 23 - yeto/rl/harness/codex/pins.py | 33 - yeto/rl/harness/codex/preflight.py | 360 --- yeto/rl/harness/codex/reward.py | 377 --- yeto/rl/harness/codex/tb2_network_grants.py | 104 - yeto/rl/harness/codex/tb2_provider.py | 980 ------- yeto/rl/harness/codex/tbench_reward.py | 130 - yeto/rl/harness/gateway/proxy.py | 69 +- yeto/rl/harness/harbor/README.md | 4 +- yeto/rl/harness/harbor/__init__.py | 6 +- yeto/rl/harness/reward_env/__init__.py | 5 +- yeto/rl/harness/reward_env/benchmark.py | 4 +- yeto/rl/harness/reward_env/tb2.py | 252 -- yeto/rl/harness/session_protocol.md | 2 + yeto/rl/island_tasks.py | 2 +- yeto/rl/rlf/__init__.py | 10 - yeto/rl/rlf/catalog.py | 116 - yeto/rl/rlf/contracts.py | 98 - yeto/rl/rlf/modal.py | 579 ---- yeto/rl/rlf/preflight.py | 41 - yeto/rl/rlf/trajectory.py | 14 - yeto/rl/secrlenv_outcome.py | 111 + yeto/rl/ssh_harness.py | 1035 +------ yeto/rl/tbench_direct_preflight.py | 129 - yeto/rl/trajectory_evidence.py | 8 +- 140 files changed, 459 insertions(+), 28917 deletions(-) delete mode 100644 docs/RLF_CODEX_RUN.md delete mode 100644 scripts/fetch_codex_bundle.py delete mode 100644 tests/golden/decoupling/codex_harness.json delete mode 100644 tests/test_codex_bridge_compaction.py delete mode 100644 tests/test_codex_compaction_wiring.py delete mode 100644 tests/test_codex_end_kind.py delete mode 100644 tests/test_codex_evaluate_time.py delete mode 100644 tests/test_codex_head_staging.py delete mode 100644 tests/test_harness_codex_exec_replay.py delete mode 100644 tests/test_harness_codex_openenv.py delete mode 100644 tests/test_harness_codex_preflight_member.py delete mode 100644 tests/test_harness_session_protocol.py delete mode 100644 tests/test_harness_tb2_provider.py delete mode 100644 tests/test_rl_agentic_suspend.py delete mode 100644 tests/test_rl_codex_schema.py delete mode 100644 tests/test_rl_eval_infer.py delete mode 100644 tests/test_rl_launcher_codex_bundle.py delete mode 100644 tests/test_rl_secrlenv_ports_retry.py delete mode 100644 tests/test_rl_tb2_baseline.py delete mode 100644 tests/test_rl_tbench_direct_preflight.py delete mode 100644 tests/test_rlf_baseline.py delete mode 100644 tests/test_rlf_bridge.py delete mode 100644 tests/test_rlf_compaction_profile.py delete mode 100644 tests/test_rlf_env_lease.py delete mode 100644 tests/test_rlf_gate_evidence.py delete mode 100644 tests/test_rlf_rl_gate.py delete mode 100644 tests/test_rlf_rollout_entry.py delete mode 100644 tests/test_sandbox_isolation.py delete mode 100644 tests/test_secrlenv_adapter.py delete mode 100644 tests/test_secrlenv_codex_harness.py delete mode 100644 tests/test_tbench21_sao_diloco_plan.py delete mode 100644 tools/build_eval_data.py delete mode 100755 tools/probes/build_tbench21_sao_diloco_plan.py delete mode 100644 tools/probes/s17_eval_island_gpu_smoke.py delete mode 100644 tools/reward_env/build_images.py delete mode 100644 yeto/cloud/modal_reward_env.py delete mode 100644 yeto/cloud/modal_rlf_baseline.py delete mode 100644 yeto/cloud/modal_tb2_baseline.py delete mode 100644 yeto/launch/codex.py create mode 100644 yeto/launch/harness_env.py delete mode 100644 yeto/rl/adapters/miles/harness_glue/codex_generate.py delete mode 100644 yeto/rl/adapters/miles/harness_glue/codex_openenv.py delete mode 100644 yeto/rl/adapters/miles/secrlenv_miles.py delete mode 100644 yeto/rl/adapters/miles/secrlenv_retry.py delete mode 100644 yeto/rl/eval/codex_infer.py delete mode 100644 yeto/rl/eval/rlf_baseline.py delete mode 100644 yeto/rl/eval/rlf_rl_gate.py delete mode 100644 yeto/rl/eval/tb2_attempt.py delete mode 100644 yeto/rl/eval/tb2_baseline.py delete mode 100644 yeto/rl/harness/codex/__init__.py delete mode 100644 yeto/rl/harness/codex/agent.py delete mode 100644 yeto/rl/harness/codex/alignment.py delete mode 100644 yeto/rl/harness/codex/chain_judge.py delete mode 100644 yeto/rl/harness/codex/client.py delete mode 100644 yeto/rl/harness/codex/codex_harness_agent.py delete mode 100644 yeto/rl/harness/codex/codex_openenv_agent_function.py delete mode 100644 yeto/rl/harness/codex/codex_openenv_agent_worker.py delete mode 100644 yeto/rl/harness/codex/codex_openenv_generate.py delete mode 100644 yeto/rl/harness/codex/codex_openenv_subprocess_agent_function.py delete mode 100644 yeto/rl/harness/codex/compaction_bridge.py delete mode 100644 yeto/rl/harness/codex/environment.py delete mode 100644 yeto/rl/harness/codex/generate.py delete mode 100644 yeto/rl/harness/codex/pins.py delete mode 100644 yeto/rl/harness/codex/preflight.py delete mode 100644 yeto/rl/harness/codex/reward.py delete mode 100644 yeto/rl/harness/codex/tb2_network_grants.py delete mode 100644 yeto/rl/harness/codex/tb2_provider.py delete mode 100644 yeto/rl/harness/codex/tbench_reward.py delete mode 100644 yeto/rl/harness/reward_env/tb2.py delete mode 100644 yeto/rl/rlf/__init__.py delete mode 100644 yeto/rl/rlf/catalog.py delete mode 100644 yeto/rl/rlf/contracts.py delete mode 100644 yeto/rl/rlf/modal.py delete mode 100644 yeto/rl/rlf/preflight.py delete mode 100644 yeto/rl/rlf/trajectory.py create mode 100644 yeto/rl/secrlenv_outcome.py delete mode 100644 yeto/rl/tbench_direct_preflight.py diff --git a/docs/RLF_CODEX_RUN.md b/docs/RLF_CODEX_RUN.md deleted file mode 100644 index 1c82a2d50..000000000 --- a/docs/RLF_CODEX_RUN.md +++ /dev/null @@ -1,177 +0,0 @@ -# RLF → Yeto Codex RL runbook - -The online RL harness is Codex only. rlenvfactory/AgentFlow may produce or review offline task packages, but its node output is never used as task identity, grader evidence, reward, or rollout input. SWE and the legacy `rl-agentic-reward-env` change remain deferred. - -## Apply order - -1. Generate or install an rlenvfactory CTF package and retain `public/`, `task.yaml`, `MANIFEST.json`, source commit, package/image digests, and runtime quality-gate evidence. -2. Admit the package with `yeto.rl.rlf.admit_task_package`; reject packages without a trusted runtime gate. -3. Acquire one Modal VM lease per `(task_id, episode_id, trajectory_id)`. Keep secrets on the trusted side and tear down the lease after grading. -4. Run the existing Codex subprocess harness. Copy RLF identities and policy identity into every chain/compaction segment; preserve tokens, logprobs, loss masks, and segment metadata. -5. Finish and grade on the trusted side, then pass the resulting metadata through the Miles or generic verl preflight. Bare float rewards are rejected. -6. Run the CPU gate before any paid GPU run. A healthy RL gate must show nonzero reward variance, success count, absolute advantage sum, and policy update. - -The baseline profile is Qwen3-30B-A3B thinking mode with a 32K window, 1024 generated tokens per turn, at most 48 turns, and 8192-byte tool output. - -## Qwen3 cyber baseline profile (5.6) - -The registered backend profile is `qwen3_30b_a3b` (`yeto.rl.codex_backend`, -allowlisted in `yeto/rl/harness/codex/pins.py`). Pass it per launch with -`--codex-backend-profile qwen3_30b_a3b`. Its budgets are declared on the -profile; three of them are launch-time knobs that must be set explicitly -because their global defaults deliberately stay unchanged: - -| Budget | Profile value | How it takes effect | -| --- | --- | --- | -| context window | 32768 tokens | `--agent-max-seq-len 32768` (run config `agent_max_seq_len`) | -| generated tokens per turn | 1024 | `YETO_CODEX_BACKEND_MAX_TOKENS=1024` | -| turns | 48 | `SECRLENV_MAX_TURNS=48` (worker-passthrough env; global default stays 40) | -| tool output | 8192 bytes | read from the profile by `codex_harness_agent` (`MAX_TOOL_OUTPUT_BYTES`); profiles without the field keep 32768 | - -Keep the 20-30 minute episode deadline (`OPENENV_MAX_ROLLOUT_TIME_SECONDS`, -default 1800s) and record `max_seq_len_hit`, `end_reason`, turns, model -tokens, tool calls, and tool-output bytes per episode for the baseline -analysis. If a flag is reached only after truncation, enable the registered compaction plan: 24K active window, 4K–8K reserve, and at most 2–3 compactions while preserving segment identity. - -## Baseline runner (5.6): RLF 形态的 Qwen3-30B-A3B - -The runner is `yeto.rl.eval.rlf_baseline` (RLF counterpart of -`yeto/rl/eval/codex_infer.py`, which covers the TB2 attempt shape). Per run: -admit the task packages into the RLF catalog, render prompt rows, bring up -serving, then per row acquire one RLF CTF lease (W4 `RlfModalProvider`), -drive the training Codex subprocess worker, finish on the trusted side, and -write the trace. Codex-only; AgentFlow artifacts are refused at admission. - -### Serving shape (inside the GPU container) - -- Modal `A100-80GB` x1, image = the pinned Miles-ports image - (`MILES_NEXT_IMAGE`), model read from the `yeto-fn-models` Volume: - `/mnt/yeto-models/hf/Qwen3-30B-A3B/ad44e777` (bf16). -- `python -m sglang.launch_server --model-path /mnt/yeto-models/hf/Qwen3-30B-A3B/ad44e777 \ - --trust-remote-code --context-length 32768 --mem-fraction-static 0.9` - plus the Miles session server (`session_server_config`, LoRA-free). The - profile's parsers/chat template come from Miles' `resolve_serving("qwen3")`. -- Serving that is not healthy within 15 min aborts the whole run and the - summary is written with `aborted.classification = "infrastructure"` - (never scored). HARD per-task deadline: 30 min - (`--episode-seconds`, default 1800; a tighter lease deadline wins). - -### Run - -Driver side (needs a local modal client): - - export YETO_CODEX_BUNDLE_DIR=/path/to/codex-bundle # scripts/fetch_codex_bundle.py - export YETO_RLF_CTF_TASKS_DIR=/path/to/tasks # admitted packages (task_id subdirs) - export YETO_RLF_BASELINE_ENVS_JSON=/path/to/envs.json # see below - modal run yeto.cloud.modal_rlf_baseline --config-json run.json - -`envs.json` holds the container env, same inputs as the eval island smoke: -the Codex contract env (`yeto.launcher.codex_container_env( -codex_bundle_contract(...))`), `TBENCH_REWARD_HMAC_KEY`, the sandbox-only -`MODAL_TOKEN_ID`/`MODAL_TOKEN_SECRET` (the RLF provider opens one VM per -episode from inside the container), and `YETO_RLF_CTF_FLAGS_JSON` -(`{"task_id": "flag{...}"}` — trusted; lives only in the Modal Secret, -never in the image, the traces, or the logs). `run.json`: - - {"run_name": "qwen3-30b-baseline-01", "repeat": 1, - "evidence": {"runtime_quality_gate": {"healthy": true}}, - "baseline": {"episode_seconds": 1800}} - -Evidence lands in the `yeto-rlf-baseline-runs` Volume under the run name. -In-container / CPU-smoke entry without Modal: - - python -m yeto.rl.eval.rlf_baseline --tasks-dir /root/rlf-tasks \ - --flags-json @flags.json --evidence '{"runtime_quality_gate": {"healthy": true}}' \ - --out /tmp/rlf-baseline-out [--repeat N] [--task ID] [--serve external --session-url URL] - -### Evidence checklist (per run) - -- `traces.jsonl`: one record per episode — RLF identity quadruple, - `policy_id` (`Qwen/Qwen3-30B-A3B@ad44e777...`), the task prompt, per-turn - completion/context tokens and tool-output bytes, per-tool events - (scrubbed command, bounded scrubbed output preview, `flag_seen` marker), - submission markers (`after_tool_calls`, `has_flag`, `flag_correct`; the - submitted flag value is never recorded), `submit_calls`, - `max_model_total_tokens`, `max_seq_len_hit`, compaction count, - `end_reason`/`end_kind`, `grader_evidence`, reward/success and the - `classification` (`success` / `failed` / `infrastructure` / `rejected`). -- `summary.json`: success rate over evaluated episodes (infrastructure and - rejected episodes are excluded, never scored 0), reward distribution, - end-reason counts, `max_seq_len_hit` rate, total compactions, - turns/tool-calls/token means and maxima, per-task breakdown. - -Stop conditions: serving start timeout (infrastructure abort), the 30-min -per-task deadline (`timeout` episodes keep reward 0 via the signed outcome -envelope and are never graded), or every prompt row processed. A -`max_seq_len_hit` rate above zero is the evidence for enabling the -registered 24K/4K compaction plan (5.7). - -Known unverified (first GPU run must confirm): Miles `resolve_serving` -coverage of the `qwen3` tito family, LoRA-free `SessionServerConfig` -(`lora_rank=None`), and A100-80GB scheduling latency on Modal. - - -## Real-GPU RL gate (5.8): Miles GRPO on `qwen35_08b` - -The gate proves the RLF -> Miles GRPO -> GAE -> policy-update chain on real -GPU with the small `qwen35_08b` profile (0.8B, the S17-verified harness -stack; LoRA, 1xGPU). The 30B capability surface is covered separately by the -5.6 baseline; 30B RL training is a follow-up item (design decision -2026-10-11). Tasks: `fixture-ctf-ezshell` (solvable in ~2 turns, provides -success) and `fixture-ctf-ret2win` (provides reward variance). Codex-only; -AgentFlow is offline-only and refused at admission. - -### Runner and commands - -The gate runner is `yeto.rl.eval.rlf_rl_gate`. Island side, in order: - - # 1. admit the packages, write Miles --prompt-data (flags never logged) - python -m yeto.rl.eval.rlf_rl_gate prompt-data \ - --tasks-dir "$YETO_RLF_CTF_TASKS_DIR" --flags-json /trusted/flags.json \ - --out /run/rlf_gate_prompt.jsonl --repeat 2 - - # 2. launch the Miles GRPO island (qwen35_08b, LoRA, 1xGPU): - # --data /run/rlf_gate_prompt.jsonl --codex-backend-profile qwen35_08b - # --reward-function rlf_success (-> --custom-rm-path yeto.rl.adapters.miles.rewards.rlf_success) - # --advantage-estimator grpo (fixed by island_entry) - # env: YETO_HARNESS_ENVIRONMENT_PROVIDER=yeto.cloud.modal_reward_env:rlf_modal_provider - # YETO_RLF_CTF_TASKS_DIR / YETO_RLF_CTF_FLAGS_JSON (Modal Secret, never in logs) - - # 3. after training, collect real artifacts, judge, write evidence - python -m yeto.rl.eval.rlf_rl_gate gate \ - --completed-groups /run/completed_groups.pt \ - --adapter-before /run/adapter_init --adapter-after /run/adapter_final \ - --step 2 --flags-json /trusted/flags.json --out /run/rlf_gate_evidence.json - -### Hard criteria (`yeto.rl.rlf.contracts.rl_gate`) - -Non-empty reward, `reward_std > 0`, `success_count > 0`, -`advantage_abs_sum > 0`, `policy_update == true`. All-zero or constant -reward MUST fail. The verdict is announced only for evidence collected -from the real island artifacts (`collector="miles_island"`); CPU/injected -collections exercise the orchestration but can never announce a pass. - -### Evidence chain (`yeto.rlf.rl_gate.evidence/v1`) - -- reward/success: the island completed-groups checkpoint - (`--completed-groups-path`) — the samples the learner actually trained - on, each with the `rlf_success` reward and `metadata.success` (W4 grade); -- advantage: recomputed per GRPO group with `contracts.group_advantages` - (the designated computation of record; never hand-filled); -- policy update: `contracts.policy_update_evidence` over a per-tensor - summary of the initial vs final LoRA adapter (代拍板: the Miles fork - exports no per-step snapshot hook; this proves parameters changed without - a fork patch, at the cost of not attributing the change to a specific - optimizer step, which the gate does not require); -- the JSON carries reward detail, success, `advantage_abs_sum`, and the - before/after policy digests; trusted flags are scrubbed - (`rlf.modal.scrub`) and only their sha256 appears in `config_digest`. - -### Stop conditions - -Stop and treat as infrastructure (never scored): serving not healthy within -15 min, or per-episode HARD deadline (30 min) exceeded repeatedly. Stop -with gate failure (the chain is broken, investigate before re-running): -`rl_gate` raising on any criterion, a sample missing `metadata.success`, -or unchanged adapter digests after training. A passing run archives the -evidence JSON next to the island artifacts. diff --git a/scripts/fetch_codex_bundle.py b/scripts/fetch_codex_bundle.py deleted file mode 100644 index c063a8aac..000000000 --- a/scripts/fetch_codex_bundle.py +++ /dev/null @@ -1,184 +0,0 @@ -#!/usr/bin/env python3 -"""Fetch and attest the pinned stock Codex Linux artifact (design R-SCOPE). - -Reproducible, pin-driven: every expected digest comes from ``yeto.rl`` -(``CODEX_NPM_PACKAGE``, ``CODEX_NPM_TARBALL_SHA256``, ``CODEX_LINUX_BINARY_SHA256``, -``CODEX_LINUX_BINARY_SIZE_BYTES``, ``CODEX_PACKAGE_MANIFEST_SHA256``, -``CODEX_APP_SERVER_SCHEMA_SHA256``). Steps: - -1. ``npm pack `` (or ``--tarball`` to reuse a download) and - check the tarball sha256. -2. Extract ``package/vendor//bin/codex`` and - ``package/vendor//codex-package.json``; check size + sha256. -3. Run ``codex --version`` and ``codex app-server generate-json-schema - --experimental`` (the same command ``yeto.rl.adapters.miles.island_entry._preflight_codex_harness`` - runs in the container) and write the v2 schema into the bundle. Codex - serialises the schema with unstable key order, so the file is written in - canonical form (``sort_keys=True, indent=2``) and compared against the pin - by sha256; a mismatch is reported (exit 3) unless ``--allow-schema-drift``. - -Output layout (``--out``, default ``~/work/codex-bundle``), the run asset dir -the launcher mounts at ``/opt/yeto/codex`` (``YETO_CODEX_BUNDLE_DIR``): - - /npm/.tgz - /codex/codex-x86_64-unknown-linux-musl - /codex/codex-package.json - /codex/codex_app_server_protocol.v2.schemas.json - /codex/BUNDLE-SHA256.json (what was verified, with the pins) -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import os -import shutil -import subprocess -import sys -import tarfile -import tempfile -from pathlib import Path - -REPO_ROOT = Path(__file__).resolve().parents[1] -sys.path.insert(0, str(REPO_ROOT)) - -from yeto.rl import ( # noqa: E402 - CODEX_APP_SERVER_SCHEMA_SHA256, - CODEX_CLI_VERSION, - CODEX_LINUX_BINARY_SHA256, - CODEX_LINUX_BINARY_SIZE_BYTES, - CODEX_LINUX_TARGET, - CODEX_NPM_PACKAGE, - CODEX_NPM_TARBALL_SHA256, - CODEX_PACKAGE_MANIFEST_SHA256, -) - -BINARY_NAME = f"codex-{CODEX_LINUX_TARGET}" -MANIFEST_NAME = "codex-package.json" -SCHEMA_NAME = "codex_app_server_protocol.v2.schemas.json" -RECORD_NAME = "BUNDLE-SHA256.json" - - -def sha256_of(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - for chunk in iter(lambda: handle.read(1 << 20), b""): - digest.update(chunk) - return digest.hexdigest() - - -def npm_pack(package: str, into: Path, npm: str) -> Path: - into.mkdir(parents=True, exist_ok=True) - out = subprocess.run( - [npm, "pack", package, "--pack-destination", str(into)], - check=True, - capture_output=True, - text=True, - ).stdout.strip().splitlines() - name = out[-1].strip() if out else "" - tarball = into / name - if not name.endswith(".tgz") or not tarball.is_file(): - raise SystemExit(f"npm pack did not report a tarball (stdout: {out!r})") - return tarball - - -def extract_member(tarball: Path, member: str, destination: Path) -> None: - with tarfile.open(tarball, "r:gz") as archive: - try: - info = archive.getmember(member) - except KeyError as exc: - raise SystemExit(f"{member} missing from {tarball}") from exc - if not info.isfile(): - raise SystemExit(f"{member} is not a regular file in {tarball}") - source = archive.extractfile(info) - assert source is not None - destination.parent.mkdir(parents=True, exist_ok=True) - with destination.open("wb") as handle: - shutil.copyfileobj(source, handle) - - -def generate_schema(binary: Path) -> dict: - with tempfile.TemporaryDirectory(prefix="yeto-codex-schema-") as temporary: - out = Path(temporary) / "schema" - subprocess.run( - [str(binary), "app-server", "generate-json-schema", "--experimental", "--out", str(out)], - check=True, - capture_output=True, - text=True, - timeout=120, - env={**os.environ, "HOME": temporary}, - ) - with (out / SCHEMA_NAME).open(encoding="utf-8") as handle: - return json.load(handle) - - -def main(argv: list[str] | None = None) -> int: - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--out", default=os.path.expanduser("~/work/codex-bundle")) - parser.add_argument("--tarball", default=None, help="reuse an already downloaded npm tarball") - parser.add_argument("--npm", default=shutil.which("npm") or "npm") - parser.add_argument("--allow-schema-drift", action="store_true", - help="write the bundle even if the canonical schema sha256 differs from the pin") - args = parser.parse_args(argv) - - out = Path(args.out).expanduser().resolve() - codex_dir = out / "codex" - tarball = Path(args.tarball).expanduser().resolve() if args.tarball else npm_pack(CODEX_NPM_PACKAGE, out / "npm", args.npm) - tarball_sha = sha256_of(tarball) - if tarball_sha != CODEX_NPM_TARBALL_SHA256: - raise SystemExit(f"tarball sha256 {tarball_sha} != pin {CODEX_NPM_TARBALL_SHA256} ({tarball})") - - binary = codex_dir / BINARY_NAME - manifest = codex_dir / MANIFEST_NAME - extract_member(tarball, f"package/vendor/{CODEX_LINUX_TARGET}/bin/codex", binary) - extract_member(tarball, f"package/vendor/{CODEX_LINUX_TARGET}/{MANIFEST_NAME}", manifest) - binary.chmod(0o755) - size = binary.stat().st_size - if size != CODEX_LINUX_BINARY_SIZE_BYTES: - raise SystemExit(f"binary size {size} != pin {CODEX_LINUX_BINARY_SIZE_BYTES}") - binary_sha = sha256_of(binary) - if binary_sha != CODEX_LINUX_BINARY_SHA256: - raise SystemExit(f"binary sha256 {binary_sha} != pin {CODEX_LINUX_BINARY_SHA256}") - manifest_sha = sha256_of(manifest) - if manifest_sha != CODEX_PACKAGE_MANIFEST_SHA256: - raise SystemExit(f"manifest sha256 {manifest_sha} != pin {CODEX_PACKAGE_MANIFEST_SHA256}") - - version = subprocess.run([str(binary), "--version"], check=True, capture_output=True, text=True, timeout=30).stdout.strip() - if version != CODEX_CLI_VERSION: - raise SystemExit(f"codex --version {version!r} != pin {CODEX_CLI_VERSION!r}") - - schema_path = codex_dir / SCHEMA_NAME - schema_path.write_text(json.dumps(generate_schema(binary), sort_keys=True, indent=2) + "\n", encoding="utf-8") - schema_sha = sha256_of(schema_path) - schema_ok = schema_sha == CODEX_APP_SERVER_SCHEMA_SHA256 - - record = { - "npm_package": CODEX_NPM_PACKAGE, - "tarball": str(tarball), - "tarball_sha256": tarball_sha, - "binary": str(binary), - "binary_sha256": binary_sha, - "binary_size_bytes": size, - "cli_version": version, - "package_manifest_sha256": manifest_sha, - "app_server_schema_sha256": schema_sha, - "app_server_schema_sha256_pin": CODEX_APP_SERVER_SCHEMA_SHA256, - "app_server_schema_matches_pin": schema_ok, - } - (codex_dir / RECORD_NAME).write_text(json.dumps(record, indent=2, sort_keys=True) + "\n", encoding="utf-8") - print(json.dumps(record, indent=2, sort_keys=True)) - if not schema_ok: - print( - f"WARNING: canonical app-server schema sha256 {schema_sha} != pin " - f"{CODEX_APP_SERVER_SCHEMA_SHA256}; the container preflight " - "(yeto.rl.adapters.miles.island_entry._preflight_codex_harness) will reject this bundle until the pin is re-derived", - file=sys.stderr, - ) - if not args.allow_schema_drift: - return 3 - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/decoupling_golden.py b/tests/decoupling_golden.py index 83563992c..021b34d9c 100644 --- a/tests/decoupling_golden.py +++ b/tests/decoupling_golden.py @@ -234,47 +234,6 @@ def _fn_cli(_tmp: Path) -> tuple[str, ...]: return tuple(fta._argv("2x8")) -def _codex_cli(_tmp: Path) -> tuple[str, ...]: - from yeto.rl import CODEX_OPENENV_AGENT - from yeto.rl.codex_backend import QWEN35_08B_MODEL, QWEN35_08B_REVISION - - return BASE + ( - "--model", QWEN35_08B_MODEL, "--model-revision", QWEN35_08B_REVISION, - "--lora-targets", "attention", - "--custom-generate-function-path", "miles.rollout.generate_hub.agentic_tool_call.generate", - "--custom-agent-function-path", CODEX_OPENENV_AGENT, - "--codex-reasoning-effort", "xhigh", "--codex-backend-profile", "qwen35_08b", - "--use-session-server", "--tito-model", "qwen35", - "--tito-allowed-append-roles", "tool", "user", "--agent-max-seq-len", "128", - ) - - -def _codex_bundle(mp, tmp: Path) -> None: - """The fake Codex bundle of ``test_rl_launcher_codex_bundle.bundle`` (contract - attestation stubbed to the pinned values; no 310 MB binary, no network).""" - import test_rl_launcher_codex_bundle as cb - import yeto.rl.ssh_harness as sh - from yeto import launcher as L - - root = tmp / "codex-bundle" / "codex" - root.mkdir(parents=True) - for name in ("codex-x86_64-unknown-linux-musl", "codex-package.json", - "codex_app_server_protocol.v2.schemas.json"): - (root / name).write_bytes(b"x") - mp.setattr(sh, "_codex_harness_contract", lambda _ns, args: cb._synthetic_contract(args)) - mp.setenv(L.CODEX_BUNDLE_DIR_ENV, str(root)) - mp.setenv(L.HARNESS_ENVIRONMENT_PROVIDER_ENV, cb.PROVIDER) - mp.delenv(L.CODEX_COMPACTION_ENV, raising=False) - mp.delenv(L.HARNESS_PREFLIGHT_ENV, raising=False) - # The reasoning / tool-call parser names are resolved by Miles - # (miles.utils.chat_template_utils), which this host must not import: the two - # values are recorded as explicit placeholders, everything else is yeto's own. - from yeto.rl.adapters.miles import config as mc - - mp.setattr(mc, "_resolve_tito_parsers", lambda model: ( - "", "")) - - CONFIGS: tuple[Config, ...] = ( Config("grpo_default", "GRPO 默认（ports，固定分区，strict-avg）", lambda t: BASE), Config("grpo_tis", "GRPO + TIS 失配修正（tis_clip 2.0）", @@ -285,8 +244,6 @@ def _codex_bundle(mp, tmp: Path) -> None: lambda t: BASE + ("--rl-algorithm-spec", str(REPO / DRGRPO_EVIDENCE))), Config("seq_adv_maxrl", "seq_adv：MaxRL 优势变换（含奖励流水线插件）", lambda t: BASE + ("--rl-algorithm-spec", _spec_file(t, "maxrl", _maxrl_spec))), - Config("codex_harness", "codex harness（ports，Qwen3.5-0.8B，假 bundle 合同）", _codex_cli, - patch=_codex_bundle), Config("elastic", "岛内弹性模式（--rl-elastic，单配置资源表）", lambda t: BASE + _elastic(t)), Config("fn_2x8", "Flash-Next 2×8 正式训练形状（fntrain.sh print 2x8）", _fn_cli, fn=True), diff --git a/tests/fixtures/trajectory_evidence/generate_signed_era.py b/tests/fixtures/trajectory_evidence/generate_signed_era.py index 0868e8c29..9a3d20278 100644 --- a/tests/fixtures/trajectory_evidence/generate_signed_era.py +++ b/tests/fixtures/trajectory_evidence/generate_signed_era.py @@ -22,7 +22,7 @@ os.environ.pop("TBENCH_REWARD_HMAC_KEY_FILE", None) from yeto.rl import tbench_outcome # noqa: E402 -from yeto.rl.harness.codex import reward # noqa: E402 +from yeto.rl import secrlenv_outcome as reward # noqa: E402 from yeto.rl.trajectory_evidence import ( # noqa: E402 build_trajectory_batch_evidence, write_trajectory_batch_evidence, diff --git a/tests/golden/decoupling/codex_harness.json b/tests/golden/decoupling/codex_harness.json deleted file mode 100644 index b957d6a5d..000000000 --- a/tests/golden/decoupling/codex_harness.json +++ /dev/null @@ -1,273 +0,0 @@ -{ - "algorithm_semantic_sha256": "1e843db7fd7e17df813ae42634f2b1f724082cbc29b2a475104e793a820d7eca", - "algorithm_sha256": "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea", - "algorithm_sha256_launcher_expected": "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea", - "algorithm_sha256_learner_ports_spec": "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea", - "algorithm_spec": { - "advantage_estimator": "grpo", - "dynamic_sampling_filter": null, - "dynamic_sampling_max_replacements": null, - "kl_coef": null, - "loss": "policy_loss", - "schema": "yeto-rl-algorithm-spec-v1" - }, - "backend_identity": { - "compat_group": "nvidia-h100", - "device_family": "nvidia", - "engine": "miles", - "param_map_sha256": "36c37d691b6372ee86485cf0f7d4f8d5e185a6346b875aac33380518ee8f259a", - "schema": "yeto-backend-identity-v2" - }, - "execution_profile": { - "contract_hash": "sha256:67aee42a1a530e9b32695039266e1cef5393816354f3dd9363a6467a31823d43", - "profile": { - "algorithm_contract": "on-policy", - "algorithm_spec_sha256": "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea", - "allowed_overlap": [], - "execution_mode": "partitioned-serial", - "groups_per_batch": 4, - "lr_schedule_sha256": "1a5872e05b262db8c1b8fd48b5b7249b41e479181097e724ae6200ed122d0e99", - "max_inflight_batches": 1, - "max_policy_age": 0, - "name": "miles-lora-partitioned-serial", - "optimizer_steps_per_round": 1, - "outer_protocol": "strict-avg", - "publish_rule": "after-every-update", - "ready_buffer_groups": 0, - "samples_per_group": 2, - "schema": "yeto-rl-execution-profile-v1" - } - }, - "launch_algorithm_sha256": "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea", - "launch_cli": [ - "--rl-placement", - "fixed-partition", - "--rl-rollout-gpus", - "1", - "--gpu", - "aws:2xa100@us-east-1", - "--model", - "Qwen/Qwen3.5-0.8B", - "--model-revision", - "2fc06364715b967f1860aea9cf38778875588b17", - "--lora-targets", - "attention", - "--custom-generate-function-path", - "miles.rollout.generate_hub.agentic_tool_call.generate", - "--custom-agent-function-path", - "yeto.rl.harness.codex.codex_openenv_subprocess_agent_function.run", - "--codex-reasoning-effort", - "xhigh", - "--codex-backend-profile", - "qwen35_08b", - "--use-session-server", - "--tito-model", - "qwen35", - "--tito-allowed-append-roles", - "tool", - "user", - "--agent-max-seq-len", - "128" - ], - "miles_argv": [ - "train.py", - "--train-backend", - "megatron", - "--hf-checkpoint", - "/model", - "--ref-load", - "/model", - "--megatron-to-hf-mode", - "bridge", - "--model-name", - "SimpleNamespace", - "--num-layers", - "2", - "--hidden-size", - "16", - "--num-attention-heads", - "4", - "--num-query-groups", - "2", - "--kv-channels", - "4", - "--ffn-hidden-size", - "32", - "--max-position-embeddings", - "32768", - "--seq-length", - "2048", - "--normalization", - "RMSNorm", - "--norm-epsilon", - "1e-06", - "--position-embedding-type", - "rope", - "--rotary-base", - "1000000", - "--rotary-percent", - "1.0", - "--vocab-size", - "64", - "--lora-rank", - "16", - "--lora-alpha", - "16", - "--lora-dropout", - "0", - "--lora-type", - "canonical_lora", - "--target-modules", - "q_proj,k_proj,v_proj,o_proj", - "--sglang-max-lora-rank", - "16", - "--actor-num-nodes", - "1", - "--actor-num-gpus-per-node", - "1", - "--num-gpus-per-node", - "2", - "--rollout-num-gpus-per-engine", - "1", - "--rollout-num-gpus", - "1", - "--update-weight-transfer-mode", - "broadcast", - "--no-offload-train", - "--sglang-mem-fraction-static", - "0.4", - "--tensor-model-parallel-size", - "1", - "--pipeline-model-parallel-size", - "1", - "--context-parallel-size", - "1", - "--expert-model-parallel-size", - "1", - "--expert-tensor-parallel-size", - "1", - "--prompt-data", - "/root/yeto-rl/prompts.jsonl", - "--input-key", - "messages", - "--label-key", - "label", - "--metadata-key", - "metadata", - "--rollout-seed", - "0", - "--num-rollout", - "3", - "--rollout-batch-size", - "4", - "--n-samples-per-prompt", - "2", - "--over-sampling-batch-size", - "4", - "--num-steps-per-rollout", - "1", - "--global-batch-size", - "8", - "--balance-data", - "--rollout-max-context-len", - "2048", - "--rollout-max-response-len", - "128", - "--rollout-sample-filter-path", - "yeto.rl.adapters.miles.rollout_meta_hook.record_trained_groups", - "--rollout-all-samples-process-path", - "yeto.rl.adapters.miles.rollout_meta_hook.extract_rollout_metadata", - "--buffer-filter-path", - "yeto.rl.adapters.miles.rollout_meta_hook.policy_buffer_filter", - "--custom-rm-path", - "pkg.reward.score", - "--advantage-estimator", - "grpo", - "--lr", - "0.0003", - "--accumulate-allreduce-grads-in-fp32", - "--attention-softmax-in-fp32", - "--attention-backend", - "unfused", - "--no-gradient-accumulation-fusion", - "--bf16", - "--no-load-optim", - "--no-load-rng", - "--no-save-optim", - "--no-save-rng", - "--finetune", - "--seed", - "0", - "--pin-rollout-manager-to-head", - "--lr-decay-style", - "linear", - "--lr-decay-iters", - "3", - "--lr-warmup-iters", - "0", - "--min-lr", - "0", - "--sglang-enable-deterministic-inference", - "--distributed-timeout-minutes", - "10", - "--apply-chat-template-kwargs", - "{\"clear_thinking\":false}", - "--custom-generate-function-path", - "miles.rollout.generate_hub.agentic_tool_call.generate", - "--custom-agent-function-path", - "yeto.rl.harness.codex.codex_openenv_subprocess_agent_function.run", - "--max-seq-len", - "128", - "--use-session-server", - "--session-server-port", - "9301", - "--tito-model", - "qwen35", - "--sglang-reasoning-parser", - "", - "--sglang-tool-call-parser", - "", - "--group-query-attention" - ], - "miles_argv_sha256": "279d3bae4c7c8381c9dfbea4e0e83ffdbfcd0afbeabb0f8931d7c3312984f771", - "name": "codex_harness", - "placement": { - "allow_cross_node_engine_tp": false, - "allow_cross_node_tp": false, - "bundle_map": null, - "expert_parallel": 1, - "gpus_per_engine": 1, - "gpus_per_node": null, - "kind": "fixed-partition", - "model_parallel": 1, - "node_parallel": 1, - "rollout_cell_names": [], - "rollout_gpus": 1, - "standby_gpus": 0, - "trainer_gpus": 1 - }, - "plugins": [ - { - "module_file": "yeto/rl/adapters/miles/rollout_meta_hook.py", - "path": "yeto.rl.adapters.miles.rollout_meta_hook.extract_rollout_metadata" - }, - { - "module_file": "yeto/rl/adapters/miles/rollout_meta_hook.py", - "path": "yeto.rl.adapters.miles.rollout_meta_hook.policy_buffer_filter" - }, - { - "module_file": "yeto/rl/adapters/miles/rollout_meta_hook.py", - "path": "yeto.rl.adapters.miles.rollout_meta_hook.record_trained_groups" - }, - { - "module_file": "yeto/rl/harness/codex/codex_openenv_subprocess_agent_function.py", - "path": "yeto.rl.harness.codex.codex_openenv_subprocess_agent_function.run" - } - ], - "runtime_attrs": { - "yeto_rl_dynamic_sampling_max_replacements": null - }, - "session_contract_hash": null, - "session_contract_hash_note": "运行时由 LoRA 张量布局算出（yeto.protocol.layout_fingerprint），再与岛契约 v3（后端身份、学习率计划、语义哈希、代码包摘要）绑定（yeto.rl.engine.backend_identity），离线不可得", - "title": "codex harness（ports，Qwen3.5-0.8B，假 bundle 合同）" -} diff --git a/tests/import_boundary_allowlist.txt b/tests/import_boundary_allowlist.txt index 9d257f9a8..6d8455c07 100644 --- a/tests/import_boundary_allowlist.txt +++ b/tests/import_boundary_allowlist.txt @@ -12,11 +12,6 @@ yeto/rl/algos/reducers.py megatron # A3 :35,41 —— 暂不动（任务 4.8， yeto/rl/algos/vendor/miles_mis.py miles # A3 :350 cp_utils —— 暂不动（任务 4.8） yeto/rl/grad_audit.py megatron # A3 :209,320 —— 暂不动（任务 4.8） # -# --- harness（A6、A7、R4-R7、R14） -yeto/rl/harness/codex/reward.py miles # R4 :159,222 Sample.Status / DynamicFilterOutput -yeto/rl/harness/codex/tbench_reward.py miles # R5 :32,101 -yeto/rl/harness/codex/tb2_provider.py modal # 审计未列（扫描新发现）：:362,368,679 tb2 环境用 Modal 沙箱；云层接口时再定归属 -# # --- 奖励与过滤（R1-R3） yeto/rl/math_reward.py miles # R1 :24 math_utils —— 保留（D9c，运行时 math 奖励需装 Miles） yeto/rl/filters.py miles # R3 :30 DynamicFilterOutput diff --git a/tests/test_codex_bridge_compaction.py b/tests/test_codex_bridge_compaction.py deleted file mode 100644 index 504c61fa6..000000000 --- a/tests/test_codex_bridge_compaction.py +++ /dev/null @@ -1,505 +0,0 @@ -"""CompactionRL compaction at the Codex Responses bridge (design D8; CPU only). - -Fake Codex client (direct POSTs to the bridge, as Codex app-server would send -them) + fake session server (aiohttp, per-session routes). No Ray, no Miles. -""" - -from __future__ import annotations - -import asyncio -import copy -import json -from dataclasses import asdict, dataclass -from types import SimpleNamespace -from typing import Any - -import aiohttp -import pytest -from aiohttp import web - -from test_secrlenv_codex_harness import _codex_body, _completion, _fake_miles_sse, _sse_events -from yeto.rl.compaction import RESUME_TEMPLATE, SUMMARY_PROMPT -from yeto.rl.harness.codex import codex_harness_agent as harness -from yeto.rl.harness.codex import codex_openenv_agent_function as adapter -from yeto.rl.harness.codex import codex_openenv_generate as generate_wrapper -from yeto.rl.harness.codex import compaction_bridge as cb - -C, T = 8000, 4000 - - -def _run(coro): - return asyncio.run(coro) - - -def _summary(text: str, total: int, finish: str = "stop") -> dict[str, Any]: - return { - "id": "chatcmpl-summary", - "object": "chat.completion", - "created": 1, - "model": harness.BACKEND_MODEL, - "choices": [{"index": 0, "finish_reason": finish, "message": { - "role": "assistant", "content": text, "reasoning_content": "compress"}}], - "usage": {"prompt_tokens": total - 10, "completion_tokens": 10, "total_tokens": total}, - } - - -class FakeSessionServer: - """Per-session Chat Completions; records (session, body, headers).""" - - def __init__(self, replies: list[dict[str, Any]]): - self.replies = list(replies) - self.requests: list[tuple[str, dict[str, Any], dict[str, str]]] = [] - - async def __aenter__(self): - async def chat(request: web.Request) -> web.Response: - self.requests.append((request.match_info["sid"], await request.json(), dict(request.headers))) - return web.Response(body=_fake_miles_sse(self.replies.pop(0)), content_type="text/event-stream") - - app = web.Application() - app.router.add_post("/sessions/{sid}/v1/chat/completions", chat) - self.runner = web.AppRunner(app, access_log=None) - await self.runner.setup() - site = web.TCPSite(self.runner, "127.0.0.1", 0) - await site.start() - self.router = f"http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}" - return self - - async def __aexit__(self, *_): - await self.runner.cleanup() - - def url(self, sid: str) -> str: - return f"{self.router}/sessions/{sid}" - - -class FakeCodex: - """Sends what Codex app-server sends: the full append-only history every turn.""" - - def __init__(self, bridge: harness._ResponsesBridge, http: aiohttp.ClientSession): - self.bridge, self.http, self.raw = bridge, http, [] - - async def turn(self, tool_output: str | None = None) -> tuple[int, dict[str, Any] | None]: - if tool_output is not None: - self.bridge.expect_tool_output(self.bridge._pending.call_id, tool_output) - history = copy.deepcopy(self.bridge._expected_input) - else: - history = [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "solve target"}]}] - async with self.http.post(f"{self.bridge.url}/v1/responses", json=_codex_body(history), - headers={"Authorization": f"Bearer {self.bridge.token}"}) as response: - text = await response.text() - self.raw.append((response.status, text)) - if response.status != 200: - return response.status, None - return 200, _sse_events(text)[-1]["response"] - - -def _env(monkeypatch, *, on: bool = True, t_comp: int = T, max_tokens: int = 512): - monkeypatch.setenv("YETO_CODEX_BACKEND_MAX_TOKENS", str(max_tokens)) - monkeypatch.delenv("YETO_CODEX_COMPACTION_ENABLED", raising=False) - if on: - monkeypatch.setenv(cb.COMPACTIONRL_ENV, "1") - monkeypatch.setenv(cb.T_COMP_ENV, str(t_comp)) - else: - monkeypatch.delenv(cb.COMPACTIONRL_ENV, raising=False) - monkeypatch.delenv(cb.T_COMP_ENV, raising=False) - - -def _exec(i: int, total: int) -> dict[str, Any]: - return _completion("terminal.exec", json.dumps({"command": f"step {i}"}), f"call-{i}", f"think {i}", total) - - -def _submit(total: int) -> dict[str, Any]: - return _completion("submit", '{"evidence":"done"}', "call-submit", "finish", total) - - -def _bridge(server, metrics, *, cls=cb.CompactionRLBridge, **kw): - if cls is cb.CompactionRLBridge: - kw.setdefault("segment_base_urls", [server.url(f"S{i}") for i in (1, 2, 3)]) - kw.setdefault("config", cb.compaction_config(C)) - return cls(server.url("S0"), "solve target", {"max_tokens": 512}, metrics, max_seq_len=C, **kw) - - -def _by_session(server): - out: dict[str, list[dict[str, Any]]] = {} - for sid, body, _headers in server.requests: - out.setdefault(sid, []).append(body) - return out - - -# --- trigger, trainable summary, rebuild, session switch -------------------- - -def test_trigger_summary_rebuild_and_new_session(monkeypatch): - _env(monkeypatch) - replies = [_exec(1, 500), _exec(2, 3900), _summary("scratchS-ONE", 4300), - _submit(900)] - - async def scenario(): - metrics = harness.legacy.AgentMetrics() - async with FakeSessionServer(replies) as server: - async with _bridge(server, metrics) as bridge, aiohttp.ClientSession() as http: - codex = FakeCodex(bridge, http) - assert (await codex.turn())[0] == 200 - assert (await codex.turn('{"output":"a"}'))[0] == 200 # 500+bound: no trigger - status, resp = await codex.turn('{"output":"b"}') # 3900+bound > C-T: compaction - assert status == 200 and resp["output"][1]["name"] == "submit" - assert bridge._expected_input is not None and len(bridge._history_after_model) == 1 + 3 * 2 + 2 # Codex keeps the full history - return server, metrics - - server, metrics = _run(scenario()) - sessions = _by_session(server) - assert [len(sessions[s]) for s in ("S0", "S1")] == [3, 1] and set(sessions) == {"S0", "S1"} - first, second, summary_req = sessions["S0"] - # Summary: same policy, same session (trainable turn with logprobs in the session record). - assert summary_req["messages"][:-1] == second["messages"] + summary_req["messages"][len(second["messages"]):-1] - assert summary_req["messages"][-1] == {"role": "user", "content": SUMMARY_PROMPT} - assert summary_req["tool_choice"] == "none" and summary_req["model"] == second["model"] - # Explicit D8 choice: the summary request carries no tool table (execution turns do). - assert summary_req["tools"] == list(cb.SUMMARY_REQUEST_TOOLS) == [] - assert summary_req["tool_choice"] == cb.SUMMARY_REQUEST_TOOL_CHOICE and second["tools"] - assert all(not k.lower().startswith("x-miles-compaction") for _s, _b, h in server.requests for k in h) - # Rebuilt context = system + u_resume(summary without ) + last k=2 atomic steps, in S1. - resumed = sessions["S1"][0]["messages"] - assert resumed[0] == {"role": "system", "content": harness.BASE_INSTRUCTIONS} - assert resumed[1] == {"role": "user", "content": RESUME_TEMPLATE.format(summary="S-ONE")} - assert resumed[2:] == summary_req["messages"][2:-1] - assert len(resumed[2:]) == 4 and "scratch" not in json.dumps(resumed) - record = cb.compaction_counters(metrics)[cb.METRICS_KEY] - assert record["compactions"] == 1 and record["t_comp"] == T and record["context_budget"] == C - assert record["segments"][0]["summary_ok"] is True and record["segments"][0]["kept_steps"] == 2 - assert metrics.turns == 3 # the summary is not a Codex turn - - -def test_trigger_boundary_is_remaining_context_below_t_comp(monkeypatch): - _env(monkeypatch) - bridge = cb.CompactionRLBridge("http://127.0.0.1:1/sessions/S0", "p", {"max_tokens": 512}, - harness.legacy.AgentMetrics(), max_seq_len=C, - segment_base_urls=["http://127.0.0.1:1/sessions/S1"] * 3) - bridge._atomic_steps = [({}, {})] - assert not bridge._should_compact(C - T) # C-|h| == T_comp: no trigger - assert bridge._should_compact(C - T + 1) - bridge._compaction_count = 3 - assert not bridge._should_compact(C) # cap reached - with pytest.raises(harness.CodexHarnessError, match="one pre-created session"): - cb.CompactionRLBridge("http://h/sessions/S0", "p", {"max_tokens": 512}, harness.legacy.AgentMetrics(), - max_seq_len=C, segment_base_urls=["http://h/sessions/S1"]) - with pytest.raises(harness.CodexHarnessError, match="config rejected"): - cb.compaction_config(T) # C must exceed T_comp - - -def test_kept_steps_shrink_when_rebuilt_context_would_trigger(monkeypatch): - _env(monkeypatch) - big = "x" * 1500 - replies = [_exec(1, 300), _exec(2, 2500), _summary("S", 5900), _submit(1200)] - - async def scenario(): - metrics = harness.legacy.AgentMetrics() - async with FakeSessionServer(replies) as server: - async with _bridge(server, metrics) as bridge, aiohttp.ClientSession() as http: - codex = FakeCodex(bridge, http) - await codex.turn() - await codex.turn(big) - assert (await codex.turn(big))[0] == 200 - return server, metrics - - server, metrics = _run(scenario()) - resumed = _by_session(server)["S1"][0]["messages"] - summary_req = _by_session(server)["S0"][-1]["messages"] - k = cb.compaction_counters(metrics)[cb.METRICS_KEY]["segments"][0]["kept_steps"] - assert k < 2 # two 1.5 kB steps would make the rebuilt context trigger again - assert resumed[2:] == (summary_req[2:-1][-2 * k:] if k else []) - - -def test_cap_of_three_then_run_to_full_truncation(monkeypatch): - _env(monkeypatch) - replies = [_exec(1, 3900)] - for i in range(3): - replies += [_summary(f"S{i}", 4200), _exec(i + 2, 3900)] - replies += [_exec(9, 7995)] # 4th trigger is ignored; the context then fills up - - async def scenario(): - metrics = harness.legacy.AgentMetrics() - async with FakeSessionServer(replies) as server: - async with _bridge(server, metrics) as bridge, aiohttp.ClientSession() as http: - codex = FakeCodex(bridge, http) - assert (await codex.turn())[0] == 200 - for _ in range(4): - assert (await codex.turn('{"output":"o"}'))[0] == 200 - status, _resp = await codex.turn('{"output":"o"}') - fatal = bridge.fatal.result() if bridge.fatal.done() else None - return server, metrics, status, fatal - - server, metrics, status, fatal = _run(scenario()) - assert status == 400 and isinstance(fatal, harness.CodexSequenceLimit) - assert metrics.max_seq_len_hit == 1 # -> policy status max_seq_len (truncated, still rewarded) - sessions = _by_session(server) - assert [len(sessions[f"S{i}"]) for i in range(4)] == [2, 2, 2, 2] - assert sessions["S3"][0]["messages"][1]["content"] == RESUME_TEMPLATE.format(summary="S2") - record = cb.compaction_counters(metrics)[cb.METRICS_KEY] - assert record["compactions"] == 3 and len(record["segments"]) == 3 - - -def test_summary_without_tag_is_kept_and_flagged(monkeypatch): - _env(monkeypatch) - replies = [_exec(1, 3900), _summary("aplain state", 4400, finish="length"), _submit(500)] - - async def scenario(): - metrics = harness.legacy.AgentMetrics() - async with FakeSessionServer(replies) as server: - async with _bridge(server, metrics) as bridge, aiohttp.ClientSession() as http: - codex = FakeCodex(bridge, http) - await codex.turn() - assert (await codex.turn('{"output":"o"}'))[0] == 200 - return server, metrics - - server, metrics = _run(scenario()) - seg = cb.compaction_counters(metrics)[cb.METRICS_KEY]["segments"][0] - assert seg["summary_ok"] is False and seg["summary_finish_reason"] == "length" - assert _by_session(server)["S1"][0]["messages"][1]["content"] == RESUME_TEMPLATE.format(summary="plain state") - - -def test_no_room_for_summary_truncates(monkeypatch): - _env(monkeypatch) - replies = [_exec(1, 7000)] - - async def scenario(): - async with FakeSessionServer(replies) as server: - async with _bridge(server, harness.legacy.AgentMetrics()) as bridge, aiohttp.ClientSession() as http: - codex = FakeCodex(bridge, http) - await codex.turn() - status, _ = await codex.turn("y" * 800) - return status, bridge.fatal.result(), len(server.requests) - - status, fatal, n = _run(scenario()) - assert status == 400 and isinstance(fatal, harness.CodexSequenceLimit) and n == 1 - - -def test_codex_argv_disables_codex_auto_compaction(monkeypatch): - _env(monkeypatch) - bridge = SimpleNamespace(url="http://127.0.0.1:9") - argv = cb.codex_argv(harness.Path("/codex"), bridge) - stock = harness._codex_argv(harness.Path("/codex"), bridge) - assert argv == [*stock, "-c", "model_auto_compact_token_limit=9223372036854775807"] - assert not any("auto_compact" in a for a in stock) - - -# --- compaction off: stock path byte-identical ------------------------------ - -def _script(): - return [_exec(1, 500), _exec(2, 2900), _exec(3, 3300), _submit(3600)] - - -async def _drive(cls, monkeypatch, **kw): - async with FakeSessionServer(_script()) as server: - metrics = harness.legacy.AgentMetrics() - async with _bridge(server, metrics, cls=cls, **kw) as bridge, aiohttp.ClientSession() as http: - codex = FakeCodex(bridge, http) - await codex.turn() - for out in ('{"o":1}', '{"o":2}', '{"o":3}'): - await codex.turn(out) - return [(s, b) for s, b, _h in server.requests], codex.raw, asdict(metrics) - - -def test_off_stock_bridge_unchanged_and_compactionrl_without_trigger_is_byte_identical(monkeypatch): - # created_at is wall-clock seconds; a run that crosses a second boundary - # would differ, so pin the clock this module reads (monotonic stays real). - monkeypatch.setattr(harness, "time", SimpleNamespace(time=lambda: 1_700_000_000.0, - monotonic=harness.time.monotonic)) - _env(monkeypatch, on=False) - stock_requests, stock_raw, stock_metrics = _run(_drive(harness._ResponsesBridge, monkeypatch)) - _env(monkeypatch, on=True, t_comp=10) # never triggers below C-10 - rl_requests, rl_raw, rl_metrics = _run(_drive(cb.CompactionRLBridge, monkeypatch, - config=cb.compaction_config(C, {cb.T_COMP_ENV: "10"}))) - stock_metrics.pop("total_generation_time"), rl_metrics.pop("total_generation_time") # wall clock - assert rl_requests == stock_requests and rl_raw == stock_raw and rl_metrics == stock_metrics - - -def test_off_drive_untrusted_trusted_side_and_metrics_are_stock(monkeypatch): - _env(monkeypatch, on=False) - calls = [] - - async def fake_drive(*args, **kwargs): - calls.append((args[1], args[4], kwargs)) - return "completed" - - monkeypatch.setattr(harness, "_drive_codex", fake_drive) - job = {"base_url": "http://r/sessions/S0", "prompt": "p", "request_kwargs": {"temperature": 1.0}, - "episode_id": "ep-1", "max_seq_len": C} - before = copy.deepcopy(job) - assert _run(adapter.prepare_segment_sessions(job, post=None)) == {} and job == before - result = _run(adapter.drive_untrusted(job, SimpleNamespace(), binary=harness.Path("/codex"))) - assert calls == [("http://r/sessions/S0", {"temperature": 1.0}, {"max_seq_len": C})] - metrics = harness.legacy.AgentMetrics() - assert set(result["metrics"]) == set(asdict(metrics)) | set(harness.tito_counters(metrics)) - with pytest.raises(harness.CodexHarnessError, match="without CompactionRL"): - _run(adapter.drive_untrusted({**job, cb.SEGMENT_URLS_KEY: ["u"] * 3}, SimpleNamespace(), - binary=harness.Path("/codex"))) - - -def test_on_trusted_side_precreates_sessions_and_worker_requires_them(monkeypatch): - _env(monkeypatch) - posts = [] - - async def post(url, body): - posts.append((url, body)) - return {"session_id": f"seg{len(posts)}"} - - job = {"base_url": "http://r:1/sessions/S0", "prompt": "p", - "request_kwargs": {"temperature": 0.7, "top_p": 1, "top_k": 20, "max_tokens": 9}, "episode_id": "e", - "max_seq_len": C} - meta = _run(adapter.prepare_segment_sessions(job, post=post)) - assert meta == {cb.SESSIONS_METADATA_KEY: ["seg1", "seg2", "seg3"], cb.SESSIONS_ROUTER_METADATA_KEY: "http://r:1"} - assert cb.SESSIONS_ROUTER_METADATA_KEY == generate_wrapper.ROUTER_KEY - assert job[cb.SEGMENT_URLS_KEY] == [f"http://r:1/sessions/seg{i}" for i in (1, 2, 3)] - assert posts == [("http://r:1/sessions", {"evaluation": False, "temperature": 0.7, "top_p": 1.0, "top_k": 20})] * 3 - with pytest.raises(harness.CodexHarnessError, match="max_seq_len"): - _run(adapter.prepare_segment_sessions({**job, "max_seq_len": None}, post=post)) - with pytest.raises(harness.CodexHarnessError, match="no pre-created segment sessions"): - _run(adapter.drive_untrusted({k: v for k, v in job.items() if k != cb.SEGMENT_URLS_KEY}, - SimpleNamespace(), binary=harness.Path("/codex"))) - - -# --- generate wrapper: one sample per segment ------------------------------- - -def _sample(n_prompt, n_gen, version="pv", **meta): - tokens = list(range(n_prompt + n_gen)) - span = SimpleNamespace(version=version, abs_start=n_prompt, abs_end=n_prompt + n_gen) - return SimpleNamespace(tokens=tokens, loss_mask=[1] * n_gen, rollout_log_probs=[-0.1] * n_gen, - response_length=n_gen, weight_versions=[SimpleNamespace(spans=[span])], - metadata=dict(meta), status=None) - - -def _seg0_meta(compactions, sessions=("a", "b", "c"), **extra): - return {"expected_policy_version": "pv", "session_server_id": "10.0.0.1:9", "exit_status": "completed", - "tbench_trusted_outcome": {"reward": 1.0}, - "chain_index": 0, "chains_total": 1, "segment_id": 0, "chain_break_reason": None, - generate_wrapper.SESSIONS_KEY: list(sessions), - "agent_metrics": {cb.METRICS_KEY: {"compactions": compactions}}, **extra} - - -@dataclass(frozen=True) -class _Out: - samples: Any - - -def _wrap(seg0, collected, *, raise_on=None): - inp = SimpleNamespace(sample=SimpleNamespace(group_index=7, index=11, metadata={"expected_policy_version": "pv"})) - seen = [] - - async def upstream(_input): - return _Out(samples=seg0) - - async def collect(_input, router, sid): - seen.append((router, sid)) - if sid == raise_on: - raise TimeoutError("boom") - return collected.get(sid, ([], {})) - - return _run(generate_wrapper.generate(inp, upstream=upstream, collect=collect)), seen - - -def test_wrapper_splits_segments_with_metadata_and_shared_reward(): - seg0 = _sample(10, 30, **_seg0_meta(2)) - s1, s2 = _sample(8, 20), _sample(8, 5) - out, seen = _wrap(seg0, {"a": ([s1], {"sm": 1}), "b": ([s2], {})}) - assert seen == [("http://10.0.0.1:9", "a"), ("http://10.0.0.1:9", "b"), ("http://10.0.0.1:9", "c")] - assert out.samples == [seg0, s1, s2] and all(s.status is None for s in out.samples) - md = [s.metadata for s in out.samples] - assert [m["segment_index"] for m in md] == [0, 1, 2] and {m["num_segments"] for m in md} == {3} - assert [m["segment_tokens"] for m in md] == [30, 20, 5] - assert [m["tokens_after"] for m in md] == [25, 5, 0] and {m["gae_length"] for m in md} == {55} - assert [m["chain_index"] for m in md] == [0, 1, 2] and {m["chains_total"] for m in md} == {3} - assert [m["chain_break_reason"] for m in md] == [None, "compaction_window", "compaction_window"] - assert {json.dumps(m["tbench_trusted_outcome"]) for m in md} == {'{"reward": 1.0}'} - assert {s.rollout_id for s in out.samples} == {11} and {s.group_index for s in out.samples} == {7} - assert md[1]["sm"] == 1 and all(m["truncated"] is False and m["compactions"] == 2 for m in md) - - -def test_wrapper_no_compaction_keeps_single_sample_and_drains_sessions(): - seg0 = _sample(10, 30, **_seg0_meta(0, exit_status="max_seq_len")) - out, seen = _wrap(seg0, {}) - assert out.samples is seg0 and len(seen) == 3 - assert seg0.metadata["num_segments"] == 1 and seg0.metadata["tokens_after"] == 0 - assert seg0.metadata["truncated"] is True - - -@pytest.mark.parametrize("case", ["mismatch_used_empty", "mismatch_unused_nonempty", "collect_error", "no_record"]) -def test_wrapper_aborts_on_inconsistent_segments(case): - meta = _seg0_meta(1) - collected = {"a": ([_sample(4, 4)], {})} - raise_on = None - if case == "mismatch_used_empty": - collected = {} - elif case == "mismatch_unused_nonempty": - collected["b"] = ([_sample(4, 4)], {}) - elif case == "collect_error": - raise_on = "a" - else: - meta["agent_metrics"] = {} - seg0 = _sample(10, 30, **meta) - out, seen = _wrap(seg0, collected, raise_on=raise_on) - assert len(seen) == 3 # every pre-created session is still collected (deleted) - assert out.samples is seg0 and seg0.status == "ABORTED" - assert "tbench_trusted_outcome" not in seg0.metadata - - -def test_wrapper_infrastructure_failure_only_drains(): - seg0 = _sample(10, 30, **_seg0_meta(1, **{generate_wrapper.INFRASTRUCTURE_KEY: "x"})) - out, seen = _wrap(seg0, {"a": ([_sample(4, 4)], {})}) - assert len(seen) == 3 and out.samples is seg0 and "segment_index" not in seg0.metadata - - -def test_wrapper_without_compaction_metadata_is_untouched(): - seg0 = _sample(10, 30, expected_policy_version="pv") - before = copy.deepcopy(seg0.metadata) - out, seen = _wrap(seg0, {}) - assert seen == [] and out.samples is seg0 - assert {k: v for k, v in seg0.metadata.items() if k not in before} == { - "chain_index": 0, "chains_total": 1, "policy_versions_actual": ["pv"]} - - -def test_wrapper_collects_on_the_recorded_router_and_aborts_without_one(): - """s19-compaction-g1-20261010d: the returned metadata had no session_server_id; - the collect URL was "http:///sessions/" and retried forever.""" - meta = _seg0_meta(1, **{generate_wrapper.ROUTER_KEY: "http://10.0.0.2:7"}) - del meta["session_server_id"] - seg0, s1 = _sample(10, 30, **meta), _sample(8, 20) - out, seen = _wrap(seg0, {"a": ([s1], {})}) - assert {r for r, _ in seen} == {"http://10.0.0.2:7"} and [m.metadata["tokens_after"] for m in out.samples] == [20, 0] - bare = _seg0_meta(1) - del bare["session_server_id"] - seg0 = _sample(10, 30, **bare) - out, seen = _wrap(seg0, {"a": ([s1], {})}) - assert seen == [] and generate_wrapper.INFRASTRUCTURE_KEY in seg0.metadata - - -def test_wrapper_bounds_each_segment_collect(monkeypatch): - """A hanging collect (Miles retries) aborts the rollout after the budget.""" - import asyncio - - monkeypatch.setenv(generate_wrapper.SEGMENT_COLLECT_TIMEOUT_ENV, "0.05") - seg0 = _sample(10, 30, **_seg0_meta(1)) - inp = SimpleNamespace(sample=SimpleNamespace(group_index=7, index=11, metadata={"expected_policy_version": "pv"})) - - async def upstream(_input): - return _Out(samples=seg0) - - async def hang(_input, _router, _sid): - await asyncio.sleep(30) - - out = _run(generate_wrapper.generate(inp, upstream=upstream, collect=hang)) - assert "timed out" in seg0.metadata[generate_wrapper.INFRASTRUCTURE_KEY] - assert generate_wrapper.segment_collect_timeout_s() == 0.05 - monkeypatch.setenv(generate_wrapper.SEGMENT_COLLECT_TIMEOUT_ENV, "bad") - assert generate_wrapper.segment_collect_timeout_s() == generate_wrapper.DEFAULT_SEGMENT_COLLECT_TIMEOUT_S - - -def test_wrapper_returns_a_list_for_every_rollout_when_compactionrl_is_on(monkeypatch): - """s19-compaction-g1-20261010e: mixed list / bare-Sample outputs broke Miles' - flattening ("'Sample' object is not iterable").""" - seg0 = _sample(10, 30, **_seg0_meta(0)) - monkeypatch.setenv(cb.COMPACTIONRL_ENV, "1") - out, _ = _wrap(seg0, {}) - assert out.samples == [seg0] - seg0 = _sample(10, 30, **_seg0_meta(0)) - monkeypatch.setenv(cb.COMPACTIONRL_ENV, "0") - out, _ = _wrap(seg0, {}) - assert out.samples is seg0 # CompactionRL off: unchanged diff --git a/tests/test_codex_compaction_wiring.py b/tests/test_codex_compaction_wiring.py deleted file mode 100644 index 2b37fe33c..000000000 --- a/tests/test_codex_compaction_wiring.py +++ /dev/null @@ -1,356 +0,0 @@ -"""CompactionRL switch wiring (progress.md "S13 Codex 压缩接线"; CPU only, no Ray). - -- spec <-> switch: launcher env injection, preflight consistency, control arm; -- pre-created segment sessions are deleted when they cannot reach the - generate wrapper (half-way creation, worker crash, finish_trusted error); -- R-D5a chain-count judge with/without compaction; -- module list / summary tool-table constants. -""" - -from __future__ import annotations - -import asyncio -from types import SimpleNamespace - -import pytest -from aiohttp import web - -from yeto import launcher as L -from yeto.rl import CODEX_HARNESS_AGENT, CODEX_OPENENV_AGENT_MODULES -from yeto.rl.algos import compactionrl as crl -from yeto.rl.harness.codex import chain_judge -from yeto.rl.harness.codex import codex_openenv_agent_function as adapter -from yeto.rl.harness.codex import codex_openenv_subprocess_agent_function as subprocess_agent -from yeto.rl.harness.codex import compaction_bridge as cb -from yeto.rl.harness.codex import preflight -from yeto.rl.harness.codex.environment import FakeTerminalEnvironment - -from test_harness_codex_openenv import _Provider, _metadata -from test_rl_launcher_codex_bundle import PROVIDER, _codex_args, bundle # noqa: F401 - fixture - -SWITCH, TCOMP = crl.COMPACTION_SWITCH_ENV, crl.COMPACTION_T_COMP_ENV - - -def _run(coro): - return asyncio.run(coro) - - -def _env(root, **extra): - return {L.CODEX_BUNDLE_DIR_ENV: str(root), L.HARNESS_ENVIRONMENT_PROVIDER_ENV: PROVIDER, **extra} - - -def _with_spec(args, spec): - args.rl_algorithm_spec_json = spec.canonical_json() - return args - - -# --- constants shared by bridge / spec ------------------------------------- - -def test_switch_names_match_the_bridge_and_module_list_includes_it(): - assert crl.COMPACTION_SWITCH_ENV == cb.COMPACTIONRL_ENV and crl.COMPACTION_T_COMP_ENV == cb.T_COMP_ENV - assert crl.compactionrl_spec().canonical_json() and crl.COMPACTED_GAE_VARIANT in crl.compactionrl_spec().canonical_json() - assert "compaction_bridge.py" in CODEX_OPENENV_AGENT_MODULES - assert preflight.required_pin_updates()["CODEX_OPENENV_AGENT_MODULES"] == CODEX_OPENENV_AGENT_MODULES - assert cb.SUMMARY_REQUEST_TOOLS == () and cb.SUMMARY_REQUEST_TOOL_CHOICE == "none" - - -# --- launcher --------------------------------------------------------------- - -def test_launcher_compactionrl_spec_switches_rollout_compaction_on(bundle): # noqa: F811 - root, _ = bundle - args = _with_spec(_codex_args(), crl.compactionrl_spec()) - # upstream agentic_tool_call.generate never collects the segment sessions - # (s19-compaction-g1-20261010b: no tokens_after, trainer failed at round 0) - with pytest.raises(ValueError, match="codex_openenv_generate"): - L.codex_harness_launch(args, environ=_env(root)) - from yeto.rl.adapters.miles import config as miles_config - - args.custom_generate_function_path = miles_config.CODEX_OPENENV_GENERATE - check = miles_config._requires_agentic_generate("--custom-agent-function-path") - assert check("x.y", SimpleNamespace(agent=args)) is None # the wrapper reads the agent flags - assert check("x.y", SimpleNamespace(agent=SimpleNamespace(custom_generate_function_path="a.b"))) - _flags, envs, _ = L.codex_harness_launch(args, environ=_env(root)) - assert envs[SWITCH] == "1" and TCOMP not in envs - # The island preflight forwards YETO_CODEX_* to the Ray rollout workers. - assert preflight.worker_runtime_env(SimpleNamespace(), envs)[SWITCH] == "1" - _flags, envs, _ = L.codex_harness_launch(args, environ=_env(root, **{TCOMP: "4096", SWITCH: "1"})) - assert envs[SWITCH] == "1" and envs[TCOMP] == "4096" - with pytest.raises(ValueError, match="contradicts"): - L.codex_harness_launch(args, environ=_env(root, **{SWITCH: "0"})) - with pytest.raises(ValueError, match="positive integer"): - L.codex_harness_launch(args, environ=_env(root, **{TCOMP: "-3"})) - - -def test_launcher_default_specs_unchanged_and_switch_rejected(bundle): # noqa: F811 - root, _ = bundle - args = _codex_args() # default spec - _f, base, _ = L.codex_harness_launch(args, environ=_env(root)) - assert SWITCH not in base and TCOMP not in base - _f, explicit_off, _ = L.codex_harness_launch(args, environ=_env(root, **{SWITCH: "0"})) - assert explicit_off == base - with pytest.raises(ValueError, match="only valid with cross_segment_per_sample"): - L.codex_harness_launch(args, environ=_env(root, **{SWITCH: "1"})) - with pytest.raises(ValueError, match="without"): - L.codex_harness_launch(args, environ=_env(root, **{TCOMP: "4096"})) - with pytest.raises(ValueError, match="boolean"): - L.codex_harness_launch(args, environ=_env(root, **{SWITCH: "maybe"})) - - -def test_launcher_control_arm_rejected_on_codex(bundle): # noqa: F811 - root, _ = bundle - args = _with_spec(_codex_args(), crl.compactionrl_whole_rollout_control_spec()) - for extra in ({}, {SWITCH: "1"}): - with pytest.raises(ValueError, match="cannot be concatenated"): - L.codex_harness_launch(args, environ=_env(root, **extra)) - - -def test_launcher_compactionrl_needs_the_compacting_harness(bundle): # noqa: F811 - root, _ = bundle - args = _with_spec(_codex_args(), crl.compactionrl_spec()) - args.custom_agent_function_path = CODEX_HARNESS_AGENT # signed, but cannot compact - with pytest.raises(ValueError, match="only the Codex OpenEnv harness"): - L.codex_harness_launch(args, environ=_env(root)) - args.custom_agent_function_path = None # not a Codex run at all - with pytest.raises(ValueError, match="only the Codex OpenEnv harness"): - L.codex_harness_launch(args, environ=_env(root)) - - -# --- island preflight ------------------------------------------------------- - -@pytest.mark.parametrize( - "variant,env,error", - [ - ("cross_segment_per_sample", {SWITCH: "1"}, None), - ("cross_segment_per_sample", {SWITCH: "1", TCOMP: "10240"}, None), - (None, {}, None), - ("vanilla", {SWITCH: "0"}, None), - ("cross_segment_per_sample", {}, "requires YETO_CODEX_COMPACTIONRL=1"), - (None, {SWITCH: "1"}, "only valid with"), - ("vanilla", {TCOMP: "100"}, "without"), - ("cross_segment_whole_rollout", {SWITCH: "1"}, "cannot be concatenated"), - ("cross_segment_whole_rollout", {}, "cannot be concatenated"), - ("cross_segment_per_sample", {SWITCH: "1", TCOMP: "x"}, "positive integer"), - ], -) -def test_preflight_switch_matches_gae_variant(variant, env, error): - args = SimpleNamespace(gae_variant=variant) - if error is None: - preflight.assert_compactionrl_consistent(args, env) - else: - with pytest.raises(preflight.PreflightError, match=error): - preflight.assert_compactionrl_consistent(args, env) - - -# --- segment session cleanup ------------------------------------------------ - -class _SessionServer: - """Fake router: POST /sessions + DELETE /sessions/{id} (pin: 204).""" - - def __init__(self, fail_after: int | None = None): - self.fail_after = fail_after - self.live: set[str] = set() - self.deleted: list[str] = [] - self.n = 0 - - async def __aenter__(self): - async def create(_request): - self.n += 1 - if self.fail_after is not None and self.n > self.fail_after: - return web.Response(status=500) - sid = f"seg{self.n}" - self.live.add(sid) - return web.json_response({"session_id": sid}) - - async def delete(request): - sid = request.match_info["sid"] - if sid not in self.live: - return web.Response(status=404) - self.live.discard(sid) - self.deleted.append(sid) - return web.Response(status=204) - - app = web.Application() - app.router.add_post("/sessions", create) - app.router.add_delete("/sessions/{sid}", delete) - self.runner = web.AppRunner(app, access_log=None) - await self.runner.setup() - site = web.TCPSite(self.runner, "127.0.0.1", 0) - await site.start() - self.url = f"http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}" - return self - - async def __aexit__(self, *exc): - await self.runner.cleanup() - - -def _job(url): - return {"base_url": f"{url}/sessions/S0", "prompt": "p", "request_kwargs": {}, "episode_id": "e", - "max_seq_len": 8000} - - -def test_half_way_precreation_deletes_what_was_created(monkeypatch): - monkeypatch.setenv(SWITCH, "1") - monkeypatch.setenv(TCOMP, "4000") - - async def scenario(): - async with _SessionServer(fail_after=2) as server: - with pytest.raises(RuntimeError, match="HTTP 500"): - await adapter.prepare_segment_sessions(_job(server.url)) - return server - - server = _run(scenario()) - assert server.deleted == ["seg1", "seg2"] and server.live == set() - - -def test_delete_segment_sessions_is_best_effort(): - calls = [] - - async def delete(url): - calls.append(url) - if url.endswith("/b"): - raise OSError("down") - - failed = _run(adapter.delete_segment_sessions("http://r:1/sessions/S0", ["a", "b", "c"], delete=delete)) - assert failed == ["b"] and calls == [f"http://r:1/sessions/{s}" for s in "abc"] - assert _run(adapter.delete_segment_sessions("http://r:1/sessions/S0", None, delete=delete)) == [] - - -def _in_process(monkeypatch, server, *, drive=None, finish=None): - monkeypatch.setenv(SWITCH, "1") - monkeypatch.setenv(TCOMP, "4000") - if drive is not None: - monkeypatch.setattr(adapter, "drive_untrusted", drive) - if finish is not None: - monkeypatch.setattr(adapter, "finish_trusted", finish) - env = FakeTerminalEnvironment(passed=True) - return adapter.run(f"{server.url}/sessions/S0", "p", {}, {"task_id": "fix-git", "trajectory_id": "t", - "max_seq_len": 8000}, - environment=env, verifier=env) - - -@pytest.mark.parametrize("where", ["drive", "finish"]) -def test_in_process_run_deletes_sessions_on_unexpected_errors(monkeypatch, where): - async def crash(*_a, **_k): - raise ValueError("worker crashed") - - async def ok_drive(*_a, **_k): - return {"status": "completed", "episode_id": "e", "metrics": {}} - - async def scenario(): - async with _SessionServer() as server: - with pytest.raises(ValueError, match="worker crashed"): - await _in_process(monkeypatch, server, drive=crash if where == "drive" else ok_drive, - finish=crash if where == "finish" else None) - return server - - server = _run(scenario()) - assert sorted(server.deleted) == ["seg1", "seg2", "seg3"] and server.live == set() - - -def test_in_process_run_hands_sessions_off_on_success(monkeypatch): - async def ok_drive(*_a, **_k): - return {"status": "completed", "episode_id": "e", "metrics": {}} - - async def finish(*_a, **_k): - return {"signed": 1} - - async def scenario(): - async with _SessionServer() as server: - result = await _in_process(monkeypatch, server, drive=ok_drive, finish=finish) - return server, result - - server, result = _run(scenario()) - assert result[cb.SESSIONS_METADATA_KEY] == ["seg1", "seg2", "seg3"] - assert server.deleted == [] and server.live == {"seg1", "seg2", "seg3"} # the wrapper collects them - - -def _subprocess(monkeypatch, *, drive, finish=None): - monkeypatch.setenv(SWITCH, "1") - monkeypatch.setenv(TCOMP, "4000") - provider = _Provider(FakeTerminalEnvironment(passed=True)) - subprocess_agent.configure(provider=provider, tool_wait_board=None, harness_board=None, member="m0") - monkeypatch.setattr(subprocess_agent, "_drive_worker", drive) - if finish is not None: - monkeypatch.setattr(adapter, "finish_trusted", finish) - return provider - - -@pytest.mark.parametrize("where", ["worker", "finish", "harness_error"]) -def test_subprocess_run_cleans_up_or_hands_off_segment_sessions(monkeypatch, where): - async def crash(*_a, **_k): - raise RuntimeError("worker crashed") - - async def harness_error(*_a, **_k): - raise adapter.harness.CodexHarnessError("worker exited without a result") - - async def ok_drive(*_a, **_k): - return {"status": "completed", "episode_id": "e", "metrics": {}} - - drive = {"worker": crash, "finish": ok_drive, "harness_error": harness_error}[where] - provider = _subprocess(monkeypatch, drive=drive, finish=crash if where == "finish" else None) - - async def scenario(): - async with _SessionServer() as server: - md = _metadata(max_seq_len=8000) - # S24 T-TB2: a worker or verifier failure is infrastructure (no reward), not a crash - result = await subprocess_agent.run(f"{server.url}/sessions/S0", "p", {}, md) - return server, result - - try: - server, result = _run(scenario()) - finally: - subprocess_agent.configure(provider=None) - assert provider.destroyed == 1 - if where == "harness_error": # metadata returned -> the generate wrapper drains them - assert result[cb.SESSIONS_METADATA_KEY] == ["seg1", "seg2", "seg3"] and server.deleted == [] - else: - assert "worker crashed" in result["tbench_infrastructure_error"] - assert cb.SESSIONS_METADATA_KEY not in result - assert sorted(server.deleted) == ["seg1", "seg2", "seg3"] and server.live == set() - - -# --- R-D5a chain-count judge ------------------------------------------------ - -def _seg(i, n, traj="t", **kw): - return {"trajectory_id": traj, "chain_index": i, "segment_index": i, "chains_total": n, - "num_segments": n, "compactions": n - 1, - "chain_break_reason": "compaction_window" if i else None, **kw} - - -def test_chain_judge_off_keeps_chains_total_one(): - ok = [{"trajectory_id": "a", "chains_total": 1}, {"trajectory_id": "b", "chains_total": 1}] - assert chain_judge.judge_chain_counts(ok, compaction_enabled=False) == [] - assert chain_judge.judge_chain_counts([_seg(0, 2), _seg(1, 2)], compaction_enabled=False) - assert chain_judge.judge_chain_counts([{"trajectory_id": "a", "chains_total": 2}], compaction_enabled=False) - - -def test_chain_judge_on_counts_segments(): - good = [_seg(0, 3), _seg(1, 3), _seg(2, 3), _seg(0, 1, traj="u"), - {"trajectory_id": "aborted", "chains_total": 1}] - assert chain_judge.judge_chain_counts(good, compaction_enabled=True) == [] - assert chain_judge.judge_chain_counts([_seg(0, 3), _seg(1, 3)], compaction_enabled=True) # missing segment - bad_comp = [_seg(0, 2), {**_seg(1, 2), "compactions": 2}] - assert any("compactions" in p for p in chain_judge.judge_chain_counts(bad_comp, compaction_enabled=True)) - bad_reason = [_seg(0, 2), {**_seg(1, 2), "chain_break_reason": "retry_fork"}] - assert chain_judge.judge_chain_counts(bad_reason, compaction_enabled=True) - bad_index = [_seg(0, 2), {**_seg(1, 2), "chain_index": 0}] - assert chain_judge.judge_chain_counts(bad_index, compaction_enabled=True) - - -def test_generate_wrapper_registers_upstream_agent_flags(monkeypatch): - """Miles parse_args calls the generate function's add_arguments; the wrapper - must register upstream agentic_tool_call's flags (s19-compaction-g1-20261010c).""" - import argparse - - from yeto.rl.harness.codex import codex_openenv_generate as gen - - def upstream(): - pass - - upstream.add_arguments = lambda p: (p.add_argument("--custom-agent-function-path"), - p.add_argument("--max-seq-len", type=int)) - monkeypatch.setattr(gen, "_load_upstream", lambda: upstream) - parser = argparse.ArgumentParser() - gen.generate.add_arguments(parser) - ns = parser.parse_args(["--custom-agent-function-path", "a.b", "--max-seq-len", "16384"]) - assert ns.custom_agent_function_path == "a.b" and ns.max_seq_len == 16384 diff --git a/tests/test_codex_end_kind.py b/tests/test_codex_end_kind.py deleted file mode 100644 index c420d371a..000000000 --- a/tests/test_codex_end_kind.py +++ /dev/null @@ -1,83 +0,0 @@ -"""S17 C9: Codex end reason split + per-turn lengths reach the trajectory event.""" - -from __future__ import annotations - -import asyncio - -import pytest - -from yeto.rl.engine import timeline -from yeto.rl.adapters.miles.rollout_meta_hook import trajectory_diagnostics -from yeto.rl.harness.codex import agent as legacy -from yeto.rl.harness.codex import codex_harness_agent as H -from yeto.rl.harness.codex import codex_openenv_agent_function as F - - -def _driver(exc=None, submit_calls=0): - metrics = legacy.AgentMetrics(submit_calls=submit_calls) - d = H._AppServerDriver.__new__(H._AppServerDriver) - d._metrics = metrics - - async def proto(): - if exc is not None: - raise exc - return "completed" - - d._drive_protocol = proto - return d, metrics - - -@pytest.mark.parametrize("exc,status,kind", [ - (H.CodexModelFailure("DSV4 must return exactly one tool call"), "max_turns", "protocol_error"), - (H.CodexTurnLimit("Codex reached the signed turn budget"), "max_turns", "turn_limit"), - (H.CodexSequenceLimit("Miles returned a truncated sample"), "max_seq_len", "response_truncated"), - (H.CodexSequenceLimit("Miles sequence limit reached"), "max_seq_len", "context_limit"), - (None, "completed", "completed_no_submit"), -]) -def test_drive_splits_end_kind(exc, status, kind): - d, m = _driver(exc) - assert asyncio.run(d.drive()) == status # signed status unchanged - assert m.end_kind == kind and kind in H.END_KINDS - - -def test_drive_submit(): - d, m = _driver(None, submit_calls=1) - assert asyncio.run(d.drive()) == "completed" and m.end_kind == "submit" - - -def test_end_kind_for_rejects_infra_errors(): - with pytest.raises(TypeError): - H.end_kind_for(H.CodexHarnessError("x")) - - -def test_turn_lengths_recorded_and_capped(): - m = legacy.AgentMetrics() - choice = {"finish_reason": "tool_calls", "message": {"content": "", "reasoning_content": "r"}} - H._note_last_completion(m, choice, {"usage": {"completion_tokens": 120, "total_tokens": 900}}) - H._note_last_completion(m, choice, {}) - H._note_turn_tool_output(m, "héllo") - assert m.turn_completion_tokens == [120, None] - assert m.turn_context_tokens == [900, None] - assert m.turn_tool_output_bytes == [6] - for _ in range(H.TURN_LENGTHS_MAX + 5): - H._note_turn_tool_output(m, "x") - assert len(m.turn_tool_output_bytes) == H.TURN_LENGTHS_MAX - d = F._metrics_dict(m) - assert d["turn_completion_tokens"] == [120, None] and "end_kind" not in d - m.end_kind = "protocol_error" - assert F._metrics_dict(m)["end_kind"] == "protocol_error" - - -def test_event_carries_fields_and_validates(): - meta = {"exit_status": "max_turns", - "agent_metrics": {"turns": 1, "parse_failures": 1, "end_kind": "protocol_error", - "turn_completion_tokens": [4100, True, None] + [1] * 300, - "turn_context_tokens": [5000], "turn_tool_output_bytes": []}} - out = trajectory_diagnostics(meta) - assert out["exit_status"] == "max_turns" and out["end_kind"] == "protocol_error" - assert out["turn_completion_tokens"][:3] == [4100, None, None] - assert len(out["turn_completion_tokens"]) == 256 - record = {"rollout_id": 0, "policy_version": 0, "sample_index": 0, "group_index": 0, "task_id": "t", - "trajectory_id": "x", "reward": 0.0, "success": False, "aborted": False, **out} - assert timeline.validate_trajectory_reward(record) == [] - assert timeline.validate_trajectory_reward({**record, "end_kind": 3}) diff --git a/tests/test_codex_evaluate_time.py b/tests/test_codex_evaluate_time.py deleted file mode 100644 index 7f44eee93..000000000 --- a/tests/test_codex_evaluate_time.py +++ /dev/null @@ -1,52 +0,0 @@ -"""S17 N13: per-trajectory grading time (``evaluate_time``) reaches the event.""" - -from __future__ import annotations - -import asyncio - -from yeto.rl.engine import timeline -from yeto.rl.adapters.miles.rollout_meta_hook import trajectory_diagnostics -from yeto.rl.harness.codex import codex_openenv_agent_function as F - - -class _Verifier: - def __init__(self, delay: float) -> None: - self.delay = delay - self.calls = 0 - - async def evaluate(self, episode_id): - self.calls += 1 - await asyncio.sleep(self.delay) - return {"passed": True, "testsh_rc": 0, "log": "ok"} - - -def _finish(status, verifier): - untrusted = {"status": status, "episode_id": "ep", "metrics": {"turns": 2, "evaluate_time": 0.0}} - return asyncio.run(F.finish_trusted(untrusted, verifier, task_id="t", sample_id="s")) - - -def test_finish_trusted_records_evaluate_time(): - v = _Verifier(0.05) - meta = _finish("completed", v) - et = meta["agent_metrics"]["evaluate_time"] - assert v.calls == 1 and isinstance(et, float) and 0.04 <= et < 5 - assert meta["agent_metrics"]["turns"] == 2 - out = trajectory_diagnostics(meta) - assert out["evaluate_time"] == round(et, 3) - record = {"rollout_id": 0, "policy_version": 0, "sample_index": 0, "group_index": 0, "task_id": "t", - "trajectory_id": "x", "reward": 1.0, "success": True, "aborted": False, **out} - assert timeline.validate_trajectory_reward(record) == [] - assert timeline.validate_trajectory_reward({**record, "evaluate_time": "1"}) - - -def test_timeout_skips_grading_and_drops_worker_value(): - v = _Verifier(0.0) - meta = _finish("timeout", v) - assert v.calls == 0 and "evaluate_time" not in meta["agent_metrics"] - assert "evaluate_time" not in trajectory_diagnostics(meta) - - -def test_hook_ignores_bad_evaluate_time(): - for bad in (True, -1.0, "3", None): - assert "evaluate_time" not in trajectory_diagnostics({"agent_metrics": {"evaluate_time": bad}}) - assert trajectory_diagnostics({"agent_metrics": {"evaluate_time": 2}})["evaluate_time"] == 2.0 diff --git a/tests/test_codex_head_staging.py b/tests/test_codex_head_staging.py deleted file mode 100644 index 89f18f6b6..000000000 --- a/tests/test_codex_head_staging.py +++ /dev/null @@ -1,40 +0,0 @@ -"""S17 M1: a signed Codex run under --controller head stages its bundle and harness env on the head.""" - -from types import SimpleNamespace - -from yeto import launcher - - -def test_non_codex_run_stages_nothing(monkeypatch): - monkeypatch.setattr("yeto.launch.codex.codex_harness_launch", lambda args, environ=None: None) - assert launcher.codex_head_staging(SimpleNamespace(), {"HF_TOKEN": "k"}) == ({}, {}) - - -def test_codex_run_mounts_bundle_and_forwards_harness_env_as_secrets(monkeypatch, tmp_path): - monkeypatch.setattr("yeto.launch.codex.codex_harness_launch", lambda args, environ=None: ("", {}, {})) - env = { - launcher.CODEX_BUNDLE_DIR_ENV: str(tmp_path), - launcher.HARNESS_ENVIRONMENT_PROVIDER_ENV: "yeto.rl.harness.codex.tb2_provider:modal_provider", - "MODAL_TOKEN_ID": "id", - "MODAL_TOKEN_SECRET": "sec", - "YETO_SANDBOX_MODAL_TOKEN_ID": "sid", - "YETO_SANDBOX_MODAL_TOKEN_SECRET": "ssec", - "OPENENV_RUN_ID": "run", - "YETO_HARNESS_TB2_TASKS_DIR": "/opt/yeto/codex/tb2-tasks", - "YETO_HARNESS_TB2_MODAL_APP": "yeto-tbench2", - "UNRELATED": "x", - "TBENCH_REWARD_HMAC_KEY": "retired", - "SECRLENV_MAX_TURNS": "", - } - mounts, secrets = launcher.codex_head_staging(SimpleNamespace(), env) - assert mounts == {launcher.HEAD_CODEX_BUNDLE_PATH: str(tmp_path.resolve())} - assert secrets == {k: env[k] for k in ( - launcher.HARNESS_ENVIRONMENT_PROVIDER_ENV, - "YETO_SANDBOX_MODAL_TOKEN_ID", "YETO_SANDBOX_MODAL_TOKEN_SECRET", "OPENENV_RUN_ID", "YETO_HARNESS_TB2_TASKS_DIR", "YETO_HARNESS_TB2_MODAL_APP")} - - -def test_args_bundle_dir_wins_over_env(monkeypatch, tmp_path): - monkeypatch.setattr("yeto.launch.codex.codex_harness_launch", lambda args, environ=None: ("", {}, {})) - mounts, _ = launcher.codex_head_staging(SimpleNamespace(codex_bundle_dir=str(tmp_path)), - {launcher.CODEX_BUNDLE_DIR_ENV: "/nope"}) - assert list(mounts.values()) == [str(tmp_path.resolve())] diff --git a/tests/test_harness_codex_exec_replay.py b/tests/test_harness_codex_exec_replay.py deleted file mode 100644 index c86d68701..000000000 --- a/tests/test_harness_codex_exec_replay.py +++ /dev/null @@ -1,89 +0,0 @@ -"""G8 / task 10.3 (reduced): replay the black-box Codex 0.145.0 binary on CPU. - -Real stock Codex binary (``YETO_CODEX_BINARY_PATH``) -> in-process -``_ResponsesBridge`` -> fake Miles ``/v1/chat/completions`` endpoint, with the -real TB2 relay (``tb2_provider.LocalProcessBackend``) and trusted verifier. -Skipped when the 0.145.0 binary is not available locally. The request-shape -summary is written to ``/tmp/t3-g8-codex-replay.json``. -""" - -from __future__ import annotations - -import asyncio -import hashlib -import json -from pathlib import Path - -import test_harness_tb2_provider as tb2_tests -import test_secrlenv_codex_harness as legacy_tests -from yeto.rl import tbench_outcome -from yeto.rl.engine.tool_wait import HarnessBoard -from yeto.rl.harness.codex import codex_harness_agent as harness -from yeto.rl.harness.codex import codex_openenv_subprocess_agent_function as subprocess_agent - -SUMMARY = Path("/tmp/t3-g8-codex-replay.json") - - -def _stock_env(monkeypatch, binary: Path) -> None: - for name, value in harness._IDENTITY_ENV.items(): - monkeypatch.setenv(name, value) - monkeypatch.setenv("YETO_CODEX_BINARY_PATH", str(binary)) - monkeypatch.setenv("YETO_CODEX_BINARY_SHA256", hashlib.sha256(binary.read_bytes()).hexdigest()) - monkeypatch.setenv("YETO_CODEX_BINARY_SIZE_BYTES", str(binary.stat().st_size)) - monkeypatch.setenv("YETO_CODEX_VERSION", harness.CODEX_CLI_VERSION) - monkeypatch.setenv("YETO_CODEX_BACKEND_MAX_TOKENS", "4096") - monkeypatch.delenv("YETO_CODEX_COMPACTION_ENABLED", raising=False) - - -def test_stock_codex_0145_single_chain_replay_through_tb2_relay(monkeypatch, tmp_path): - binary = legacy_tests._stock_codex_binary() - _stock_env(monkeypatch, binary) - hb = HarnessBoard() - provider = tb2_tests._provider(tmp_path) - tb2_tests._configure(monkeypatch, provider, harness_board=hb) - monkeypatch.delenv("YETO_CODEX_OPENENV_ALLOW_SCRIPTED_DRIVER", raising=False) - - async def scenario(): - runner, miles_url, requests = await legacy_tests._fake_miles( - [ - legacy_tests._completion("terminal.exec", '{"command":"echo ok > fixed.txt && ls"}', "call-1", "fix it", 2_000), - legacy_tests._completion("terminal.exec", '{"command":"cat fixed.txt"}', "call-2", "check", 2_500), - legacy_tests._completion("submit", '{"evidence":"fixed.txt written"}', "call-3", "done", 3_000), - ] - ) - try: - metadata = { - "task_id": "fix-git", "trajectory_id": "traj-g8", "prompt": "write ok into fixed.txt", - "expected_policy_version": "pv-g8", - } - result = await asyncio.wait_for(subprocess_agent.run(miles_url, "p", {}, metadata), timeout=120) - finally: - await runner.cleanup() - return result, requests - - result, requests = asyncio.run(scenario()) - outcome, reward = tbench_outcome.verified_outcome(result) - assert outcome["status"] == "completed", result - assert reward == 1.0 - assert result["chains_total"] == 1 and result["chain_break_reason"] is None - assert provider.destroyed == 1 and provider.live == {} - snap = hb.snapshot() - assert snap.env_live == 0 and snap.in_flight == 0 - assert snap.tito_session_mismatch == 0 - breaks = snap.tito_chain_breaks - assert not any(breaks.values()) if isinstance(breaks, dict) else breaks == 0 - # black-box Codex request shape: exactly 3 samples, one chain, tool set fixed, history grows by tool turns - assert len(requests) == 3 - assert [t["function"]["name"] for t in requests[0]["tools"]] == ["terminal.exec", "submit"] - roles = [[m["role"] for m in r["messages"]] for r in requests] - assert roles[1][-1] == "tool" and roles[2][-1] == "tool" - assert all(len(roles[i]) < len(roles[i + 1]) for i in range(2)) - SUMMARY.write_text(json.dumps({ - "codex_version": harness.CODEX_CLI_VERSION, - "binary_sha256": hashlib.sha256(binary.read_bytes()).hexdigest(), - "requests": len(requests), - "tools": [t["function"]["name"] for t in requests[0]["tools"]], - "message_roles_per_request": roles, - "status": outcome["status"], "reward": reward, "chains_total": result["chains_total"], - "tito_session_mismatch": snap.tito_session_mismatch, - }, indent=2) + "\n") diff --git a/tests/test_harness_codex_openenv.py b/tests/test_harness_codex_openenv.py deleted file mode 100644 index 9d2104efc..000000000 --- a/tests/test_harness_codex_openenv.py +++ /dev/null @@ -1,935 +0,0 @@ -"""Ports-path acceptance for the Codex Terminal-Bench harness (change rl-codex-harness-rollout). - -Covers the six contracts listed in CODEX-PROGRESS §1.3 on top of the moved -legacy suite (tests/test_secrlenv_codex_harness.py): -1. Responses tool call -> TITO append-only prefix reuse on the Miles wire; -2. rollout-level cancellation cleanup (worker process group, lease, tool-wait); -3. reward verification at all three points; -4. mask / token / logprob alignment; -5. sibling segments share one reward and count once; -6. compaction enabled -> preflight fails closed. - -Stage 4 (IR-1..IR-4 landed, INFRA 8347351): the agent entry is driven through -the real ``tool_wait.HarnessBoard`` (leases / sessions / admission), -``rollout_meta_hook.expected_policy_version`` (IR-3), ``entry.preflight_stage`` -with the ``(miles_args, launch)`` hook (IR-1) and the 1.7 schema names (IR-4). -""" - -from __future__ import annotations - -import asyncio -import copy -import json -from pathlib import Path -from types import SimpleNamespace -from typing import Any - -import aiohttp -import pytest - -import test_secrlenv_codex_harness as legacy_tests -from yeto.rl import tbench_outcome -from yeto.rl.engine.algorithm import AlgorithmSpec -from yeto.rl.adapters.miles import entry, rollout_meta_hook -from yeto.rl.adapters.miles.config import MilesConfigError -from yeto.rl.engine.timeline import HARNESS_METRIC_KEYS, LOAD_SAMPLE_SCHEMA -from yeto.rl.engine.tool_wait import HARNESS_ZERO, HarnessBoard, HarnessSnapshot, ToolWaitBoard, drain_blockers -from yeto.rl.harness.codex import ( - alignment, - codex_harness_agent as harness, - codex_openenv_agent_function as adapter, - codex_openenv_generate as generate_wrapper, - codex_openenv_subprocess_agent_function as subprocess_agent, - preflight, - tbench_reward, -) -from yeto.rl.harness.codex.environment import FakeTerminalEnvironment, serve_environment - - - -class _Board: - def __init__(self) -> None: - self.events: list[tuple[str, str]] = [] - - def enter(self, trajectory_id: str) -> None: - self.events.append(("enter", trajectory_id)) - - def exit(self, trajectory_id: str) -> None: - self.events.append(("exit", trajectory_id)) - - @property - def open(self) -> int: - return sum(1 for e, _ in self.events if e == "enter") - sum(1 for e, _ in self.events if e == "exit") - - -class _Provider: - def __init__(self, env: FakeTerminalEnvironment, *, deadline: float = 30.0, fail: bool = False) -> None: - self.env = env - self.deadline = deadline - self.fail = fail - self.destroyed = 0 - self.runner = None - - async def acquire(self, task_id: str, trajectory_id: str) -> subprocess_agent.EnvironmentLease: - if self.fail: - raise OSError("sandbox quota exhausted") - self.runner, url = await serve_environment(self.env, "tok") - - async def destroy() -> None: - self.destroyed += 1 - await self.runner.cleanup() - - async def describe() -> str: - return "gone" if self.destroyed else "live" - - return subprocess_agent.EnvironmentLease( - env_url=url, env_token="tok", verifier=self.env, destroy=destroy, describe=describe, - deadline_seconds=self.deadline, - ) - - -def _run(coro): - return asyncio.run(coro) - - -# ---------------------------------------------------------------- identity / preflight (6) - -def test_identity_matches_stock_harness_and_profile(): - identity = adapter.codex_openenv_harness_identity() - assert set(identity) == { - "base_instructions_sha256", "terminal_exec_tool_schema_sha256", - "submit_tool_schema_sha256", "dynamic_tools_schema_sha256", - } - assert adapter._OPENENV_IDENTITY_ENV["YETO_CODEX_OPENENV_BACKEND_PROFILE"] == "qwen35_08b" - assert adapter.stock.BACKEND_MODEL == "qwen35" - assert callable(adapter.stock._attest_runtime) - - -def _good_env() -> dict[str, str]: - return dict(adapter._OPENENV_IDENTITY_ENV) - - -def test_preflight_fails_closed_on_compaction_scripted_driver_identity_and_version(tmp_path): - ok = lambda: tmp_path / "codex" # noqa: E731 - report = preflight.preflight_codex_openenv(_good_env(), attest_runtime=ok) - assert report["compaction"] == "disabled" and report["codex_version"] == "codex-cli 0.145.0" - for env, match in ( - ({**_good_env(), "YETO_CODEX_COMPACTION_ENABLED": "1"}, "incompatible"), - ({**_good_env(), "YETO_CODEX_MAX_COMPACTIONS": "3"}, "compaction tuning"), - ({**_good_env(), "YETO_CODEX_OPENENV_ALLOW_SCRIPTED_DRIVER": "1"}, "must not be set"), - ({k: v for k, v in _good_env().items() if "MODEL_REVISION" not in k}, "identity drifted"), - ): - with pytest.raises(preflight.PreflightError, match=match): - preflight.preflight_codex_openenv(env, attest_runtime=ok) - - def bad_version() -> Path: - raise harness.CodexHarnessError("pinned Codex version contract drifted") - - with pytest.raises(preflight.PreflightError, match="version contract drifted"): - preflight.preflight_codex_openenv(_good_env(), attest_runtime=bad_version) - - -# ---------------------------------------------------------------- Responses -> TITO prefix (1) - -def test_responses_tool_call_round_trip_is_append_only_on_the_miles_wire(monkeypatch): - legacy_tests._set_bridge_env(monkeypatch) - - async def scenario() -> None: - runner, miles_url, miles_requests = await legacy_tests._fake_miles([ - legacy_tests._completion("terminal.exec", '{ "command": "git status" }', "call-1", "look", 25), - legacy_tests._completion("submit", '{"evidence":"done"}', "call-2", "finish", 40), - ]) - metrics = harness.legacy.AgentMetrics() - try: - async with harness._ResponsesBridge( - miles_url, "fix the repo", {"temperature": 0.7, "top_p": 0.9, "max_tokens": 128, "stream": True, "n": 1}, - metrics, max_seq_len=512, - ) as bridge: - headers = {"Authorization": f"Bearer {bridge.token}"} - initial = [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "fix the repo"}]}] - async with aiohttp.ClientSession() as session: - async with session.post(f"{bridge.url}/v1/responses", json=legacy_tests._codex_body(initial), headers=headers) as r: - assert r.status == 200 - legacy_tests._sse_events(await r.text()) - bridge.expect_tool_output("call-1", '{"exit_code":0,"output":"clean","timed_out":false,"truncated":false}') - second = copy.deepcopy(bridge._expected_input) - async with session.post(f"{bridge.url}/v1/responses", json=legacy_tests._codex_body(second), headers=headers) as r: - assert r.status == 200 - # Append-only: request 2's messages == request 1's + [assistant, tool] - m1, m2 = miles_requests[0]["messages"], miles_requests[1]["messages"] - assert m2[: len(m1)] == m1 - appended = m2[len(m1):] - assert [m["role"] for m in appended] == ["assistant", "tool"] - assert appended[0]["tool_calls"][0]["id"] == "call-1" - assert appended[1]["tool_call_id"] == "call-1" - # Mutated history is refused (no third Miles sample is taken) - mutated = copy.deepcopy(second) - mutated[0]["content"][0]["text"] = "fix the repo!" - async with session.post(f"{bridge.url}/v1/responses", json=legacy_tests._codex_body(mutated), headers=headers) as r: - assert r.status == 400 - assert len(miles_requests) == 2 - finally: - await runner.cleanup() - - _run(scenario()) - - -# ---------------------------------------------------------------- subprocess entry: outcome, tool-wait, cleanup (2) - -def _configure(monkeypatch, env: FakeTerminalEnvironment, harness_board: Any = None, **kw) -> tuple[_Provider, _Board]: - monkeypatch.setenv("YETO_CODEX_OPENENV_ALLOW_SCRIPTED_DRIVER", "1") - provider = _Provider(env, **kw) - board = _Board() - subprocess_agent.configure(provider=provider, tool_wait_board=board, harness_board=harness_board, member="m0") - return provider, board - - -def _metadata(**extra) -> dict[str, Any]: - # IR-3: the driver's target token travels in the prompt metadata (or the sink). - return {"task_id": "fix-git", "trajectory_id": "traj-1", "prompt": "fix the repo", - "codex_openenv_driver": "scripted", "script": ["git status", "git checkout -- ."], - "expected_policy_version": "pv-7", **extra} - - -def test_subprocess_run_relays_tool_wait_and_builds_outcome(monkeypatch): - env = FakeTerminalEnvironment(passed=True) - provider, board = _configure(monkeypatch, env) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata(expected_policy_version="pv-7"))) - assert result is not None and tbench_outcome.OUTCOME_KEY in result - outcome, reward = tbench_outcome.verified_outcome(result) - assert reward == 1.0 and outcome["status"] == "completed" and outcome["task_id"] == "fix-git" - assert env.evaluations == 1 and [c for _, c in env.calls] == ["git status", "git checkout -- ."] - assert board.events == [("enter", "traj-1"), ("exit", "traj-1")] * 3 and board.open == 0 - assert provider.destroyed == 1 - assert result["chains_total"] == 1 and result["reward_scope"] == "trajectory" - assert result["expected_policy_version"] == "pv-7" - - -def test_subprocess_negative_reward_and_timeout_are_policy_boundaries(monkeypatch): - env = FakeTerminalEnvironment(passed=False) - _configure(monkeypatch, env) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - outcome, reward = tbench_outcome.verified_outcome(result) - assert reward == 0.0 and outcome["status"] == "completed" - - env = FakeTerminalEnvironment(passed=True) - provider, board = _configure(monkeypatch, env, deadline=0.5) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata(hang_seconds=20))) - outcome, reward = tbench_outcome.verified_outcome(result) - assert reward == 0.0 and outcome["status"] == "timeout" and outcome["verifier"] == tbench_outcome.TIMEOUT_VERIFIER - assert env.evaluations == 0 and provider.destroyed == 1 and board.open == 0 - - -def test_subprocess_infrastructure_failures_have_no_outcome_and_abort(monkeypatch): - env = FakeTerminalEnvironment(passed=True) - _configure(monkeypatch, env, fail=True) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - assert tbench_reward.INFRASTRUCTURE_KEY in result and tbench_outcome.OUTCOME_KEY not in result - sample = SimpleNamespace(metadata=result, status=None) - assert _run(tbench_reward.reward_func(None, sample)) == 0.0 and sample.status == "ABORTED" - - -def test_rollout_cancellation_kills_worker_and_releases_everything(monkeypatch): - env = FakeTerminalEnvironment(passed=True) - provider, board = _configure(monkeypatch, env) - terminated: list[int | None] = [] - original = subprocess_agent._terminate - - async def recording(process): - await original(process) - terminated.append(process.returncode) - - monkeypatch.setattr(subprocess_agent, "_terminate", recording) - - async def scenario() -> None: - task = asyncio.create_task(subprocess_agent.run("http://miles", "p", {}, _metadata(hang_seconds=30))) - while not board.events: - await asyncio.sleep(0.05) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - _run(scenario()) - assert terminated and terminated[0] is not None # worker reaped - assert board.open == 0 and provider.destroyed == 1 and env.evaluations == 0 - - -def test_abort_hook_cancels_surplus_workers_and_releases_everything(monkeypatch): - """S17 N17: Miles' call_agent_abort_hook finds ``.abort``; it cancels the - in-flight Codex workers so an over-sampled rollout does not wait for surplus - trajectories. The aborted trajectory is infrastructure (never a reward) and - its worker, lease and board entries are released.""" - env = FakeTerminalEnvironment(passed=True) - provider, board = _configure(monkeypatch, env) - - async def scenario(): - task = asyncio.create_task(subprocess_agent.run("http://miles", "p", {}, _metadata(hang_seconds=30))) - while not board.events: - await asyncio.sleep(0.05) - assert await subprocess_agent.abort(None) == 1 - return await asyncio.wait_for(task, 20) - - result = _run(scenario()) - assert tbench_reward.INFRASTRUCTURE_KEY in result and tbench_outcome.OUTCOME_KEY not in result - assert subprocess_agent.ABORTED_STATUS in json.dumps(result) - assert board.open == 0 and provider.destroyed == 1 and env.evaluations == 0 - assert subprocess_agent._INFLIGHT == {} and subprocess_agent._ABORTED_IDS == set() - assert _run(subprocess_agent.abort(None)) == 0 # nothing in flight: no-op - - -def test_abort_during_acquire_waits_for_the_sandbox_then_destroys_it(monkeypatch): - """S17 N17: acquire is never cancelled (a sandbox made by a cancelled acquire would - leak); the trajectory is flagged and stops as soon as the sandbox exists.""" - env = FakeTerminalEnvironment(passed=True) - provider, board = _configure(monkeypatch, env) - gate = asyncio.Event() - real_acquire = provider.acquire - - async def slow_acquire(task_id, trajectory_id): - await gate.wait() - return await real_acquire(task_id, trajectory_id) - - provider.acquire = slow_acquire - - async def scenario(): - task = asyncio.create_task(subprocess_agent.run("http://miles", "p", {}, _metadata(hang_seconds=30))) - while "traj-1" not in subprocess_agent._INFLIGHT: - await asyncio.sleep(0.01) - assert subprocess_agent._INFLIGHT["traj-1"]["stage"] == "acquire" - assert await subprocess_agent.abort(None) == 1 - gate.set() - return await asyncio.wait_for(task, 20) - - result = _run(scenario()) - assert subprocess_agent.ABORTED_STATUS in json.dumps(result) and tbench_outcome.OUTCOME_KEY not in result - assert provider.destroyed == 1 and env.calls == [] and env.evaluations == 0 - assert subprocess_agent._INFLIGHT == {} and subprocess_agent._ABORTED_IDS == set() - - -def test_abort_during_verification_cancels_it_and_destroys_the_sandbox(monkeypatch): - env = FakeTerminalEnvironment(passed=True) - provider, board = _configure(monkeypatch, env) - started = asyncio.Event() - - async def slow_finish(untrusted, verifier, **kw): - started.set() - await asyncio.sleep(60) # a judge that hangs (detached sandbox) - - monkeypatch.setattr(subprocess_agent.adapter, "finish_trusted", slow_finish) - - async def scenario(): - task = asyncio.create_task(subprocess_agent.run("http://miles", "p", {}, _metadata())) - await asyncio.wait_for(started.wait(), 20) - assert subprocess_agent._INFLIGHT["traj-1"]["stage"] == "verify" - assert await subprocess_agent.abort(None) == 1 - return await asyncio.wait_for(task, 20) - - result = _run(scenario()) - assert subprocess_agent.ABORTED_STATUS in json.dumps(result) and tbench_outcome.OUTCOME_KEY not in result - assert provider.destroyed == 1 and board.open == 0 - assert subprocess_agent._INFLIGHT == {} and subprocess_agent._ABORTED_IDS == set() - - -def test_lease_force_released_by_the_board_is_infrastructure_not_a_crash(monkeypatch): - """S17 N17 A/B arm a: the board force-released an expired lease; the late - lease_released raised ToolWaitError("... is not live") out of run() and the - reward function then failed the whole rollout. Now: infrastructure sample.""" - from yeto.rl.engine.tool_wait import HarnessBoard - - hb = HarnessBoard() - env = FakeTerminalEnvironment(passed=True) - provider, _board = _configure(monkeypatch, env, harness_board=hb, deadline=30.0) - real_destroy_calls = [] - - orig_acquire = provider.acquire - - async def acquire(task_id, trajectory_id): - lease = await orig_acquire(task_id, trajectory_id) - orig_destroy = lease.destroy - - async def destroy(): - hb._leases.pop(trajectory_id, None) # the board expired it meanwhile - real_destroy_calls.append(trajectory_id) - await orig_destroy() - - return subprocess_agent.EnvironmentLease( - env_url=lease.env_url, env_token=lease.env_token, verifier=lease.verifier, - destroy=destroy, describe=lease.describe, deadline_seconds=lease.deadline_seconds) - - provider.acquire = acquire - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - assert tbench_reward.INFRASTRUCTURE_KEY in result and tbench_outcome.OUTCOME_KEY not in result - assert subprocess_agent.LEASE_EXPIRED_ON_BOARD in json.dumps(result) - assert real_destroy_calls == ["traj-1"] and provider.destroyed == 1 - assert hb.snapshot().in_flight == 0 # the session was still closed - sample = SimpleNamespace(metadata=result, status=None) - assert _run(tbench_reward.reward_func(None, sample)) == 0.0 and sample.status == "ABORTED" - - -# ---------------------------------------------------------------- IR-2/IR-3: HarnessBoard, admission, policy token source - -def test_subprocess_run_counts_sessions_and_leases_on_the_harness_board_and_drains_to_zero(monkeypatch): - import time - - hb = HarnessBoard() # monotonic clock, as the island actor - env = FakeTerminalEnvironment(passed=True) - provider, _board = _configure(monkeypatch, env, harness_board=hb, deadline=30.0) - seen: list[HarnessSnapshot] = [] - - async def scenario() -> dict[str, Any]: - task = asyncio.create_task(subprocess_agent.run("http://miles", "p", {}, _metadata(hang_seconds=0.5))) - while not (hb.snapshot().env_live and hb.snapshot().in_flight): - await asyncio.sleep(0.02) - seen.append((hb.snapshot(), time.monotonic())) - return await task - - result = _run(scenario()) - busy, at = seen[0] - # the lease carries its hard deadline (acquire time + provider deadline): drain waits are bounded - assert busy.in_flight == 1 and busy.env_live == 1 and busy.latest_lease_deadline == pytest.approx(at + 30.0, abs=5.0) - # drain probe (IR-2, three-arg): a live sandbox / open session blocks even with router=0, tool_wait=0 - idle_tools = ToolWaitBoard().snapshot() - assert drain_blockers(0, idle_tools, busy) == ["1 harness sessions in flight", "1 sandboxes live"] - assert drain_blockers(0, idle_tools, None) == ["harness counts unknown"] - assert drain_blockers(0, idle_tools, HARNESS_ZERO) == [] - done = hb.snapshot() - assert done.in_flight == 0 and done.env_live == 0 and done.leases_expired_total == 0 - assert drain_blockers(0, idle_tools, done) == [] and provider.destroyed == 1 - assert tbench_outcome.OUTCOME_KEY in result and result["expected_policy_version"] == "pv-7" - - -def test_subprocess_run_refuses_new_session_when_admission_is_closed(monkeypatch): - hb = HarnessBoard() - env = FakeTerminalEnvironment(passed=True) - provider, _board = _configure(monkeypatch, env, harness_board=hb) - hb.close_admission(["m0"]) # drain of our member, first step of MilesRolloutPool.drain - assert hb.allow_new_session("m0") is False - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - assert result[tbench_reward.INFRASTRUCTURE_KEY] == subprocess_agent.ADMISSION_CLOSED - assert tbench_outcome.OUTCOME_KEY not in result and provider.runner is None # nothing acquired - assert hb.snapshot() == HarnessSnapshot(0, 0, hb.snapshot().generation) - hb.open_admission(["m0"]) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - assert tbench_outcome.OUTCOME_KEY in result and hb.snapshot().in_flight == 0 and hb.snapshot().env_live == 0 - - -def test_subprocess_cancellation_releases_board_lease_and_session(monkeypatch): - hb = HarnessBoard() - env = FakeTerminalEnvironment(passed=True) - provider, board = _configure(monkeypatch, env, harness_board=hb) - - async def scenario() -> None: - task = asyncio.create_task(subprocess_agent.run("http://miles", "p", {}, _metadata(hang_seconds=30))) - while not board.events: - await asyncio.sleep(0.05) - assert hb.snapshot().in_flight == 1 and hb.snapshot().env_live == 1 - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - _run(scenario()) - snap = hb.snapshot() - assert snap.in_flight == 0 and snap.env_live == 0 and provider.destroyed == 1 and board.open == 0 - - -def test_subprocess_policy_token_comes_from_metadata_or_driver_sink_and_missing_refuses(monkeypatch, tmp_path): - env = FakeTerminalEnvironment(passed=True) - _configure(monkeypatch, env) - monkeypatch.setenv(rollout_meta_hook.META_SINK_ENV, f"dir:{tmp_path}") - md = _metadata() - del md["expected_policy_version"] - # nothing published by the driver -> refuse before any environment is acquired - with pytest.raises(subprocess_agent.PolicyVersionMissing, match="expected_policy_version missing"): - _run(subprocess_agent.run("http://miles", "p", {}, dict(md))) - assert env.evaluations == 0 - rollout_meta_hook.put_policy_token("yeto:3:abc") # what MilesRolloutPool.generate publishes - result = _run(subprocess_agent.run("http://miles", "p", {}, dict(md))) - assert result["expected_policy_version"] == "yeto:3:abc" - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) # metadata wins over the sink - assert result["expected_policy_version"] == "pv-7" - - -def test_generate_wrapper_refuses_without_policy_token_and_reads_sink(monkeypatch, tmp_path): - monkeypatch.setenv(rollout_meta_hook.META_SINK_ENV, f"dir:{tmp_path}") - calls = [] - - async def upstream(_input): - calls.append(1) - return SimpleNamespace(samples=[_sample([1, 2, 3], [0, 1, 1], [-0.1, -0.2], 2, [("yeto:3:abc", 1, 3)])]) - - inp = SimpleNamespace(sample=SimpleNamespace(group_index=1, index=2, metadata={})) - with pytest.raises(generate_wrapper.PolicyVersionMissing): - _run(generate_wrapper.generate(inp, upstream=upstream)) - assert calls == [] - rollout_meta_hook.put_policy_token("yeto:3:abc") - out = _run(generate_wrapper.generate(inp, upstream=upstream)) - s = out.samples[0] - assert s.status is None and s.metadata["expected_policy_version"] == "yeto:3:abc" - assert rollout_meta_hook.harness_counters([[s]]) == {} - rollout_meta_hook.put_policy_token("yeto:4:def") - out = _run(generate_wrapper.generate(inp, upstream=upstream)) - s = out.samples[0] - assert s.status == "ABORTED" and rollout_meta_hook.harness_counters([[s]]) == {"policy_age_violation": 1} - - -# ---------------------------------------------------------------- IR-1: preflight hook before any allocation; 7.1 reward_scope - -def _island_args(**over) -> SimpleNamespace: - base = dict( - yeto_harness_preflight=preflight.HARNESS_PREFLIGHT_SPEC, - custom_agent_function_path=preflight.EXPECTED_AGENT_FUNCTION, - yeto_harness_reward_scope=None, - yeto_harness_environment_provider="test_harness_codex_openenv:_provider_factory", - yeto_rl_learner_id=0, - ) - return SimpleNamespace(**{**base, **over}) - - -def _provider_factory(miles_args): - return _Provider(FakeTerminalEnvironment(passed=True)) - - -def _stage(monkeypatch): - """Stub the A1 contract pieces so entry.preflight_stage runs without Miles (as test_rl_ir_harness).""" - monkeypatch.setattr(entry, "ports_runtime_fingerprint", lambda launch: "fp") - monkeypatch.setattr(entry, "miles_capabilities", lambda fp, unverified_mechanisms=(): "caps") - monkeypatch.setattr(entry, "with_partitioned_serial", lambda caps: caps) - monkeypatch.setattr(entry, "execution_profile_for", lambda *a, **k: "profile") - monkeypatch.setattr(entry, "expected_algorithm_sha256", lambda a: None) - monkeypatch.setattr(entry, "preflight", lambda profile, algorithm, caps: None) - monkeypatch.setattr(entry, "elastic_wiring_for", lambda a, profile, fingerprint: None) - allocations = [] - monkeypatch.setattr(entry, "connect_island_ray", lambda *a, **k: allocations.append("ray")) - monkeypatch.setattr(entry, "resolve_harness_preflight", - lambda a, environ=None: preflight.harness_preflight if a.yeto_harness_preflight else None) - return allocations - - -def test_entry_preflight_stage_runs_codex_preflight_before_allocation_and_installs_boards(monkeypatch, tmp_path): - allocations = _stage(monkeypatch) - for name, value in _good_env().items(): - monkeypatch.setenv(name, value) - monkeypatch.delenv("YETO_CODEX_OPENENV_ALLOW_SCRIPTED_DRIVER", raising=False) - attested = [] - monkeypatch.setattr(preflight, "preflight_codex_openenv", - lambda env=None, **kw: attested.append(env.get("YETO_CODEX_OPENENV_MODEL_REVISION")) or {}) - subprocess_agent.configure(provider=None) - # the real spec resolves through entry.resolve_harness_preflight (module:callable) - assert entry.resolve_harness_preflight(SimpleNamespace(yeto_harness_preflight=preflight.HARNESS_PREFLIGHT_SPEC)) is preflight.harness_preflight - - # failure (compaction on) -> no Ray connect / placement / allocation, no provider installed - monkeypatch.setenv("YETO_CODEX_COMPACTION_ENABLED", "1") - monkeypatch.setattr(preflight, "preflight_codex_openenv", lambda env=None, **kw: preflight.assert_compaction_disabled(env)) - with pytest.raises(preflight.PreflightError, match="incompatible"): - entry.preflight_stage(_island_args(), "launch", AlgorithmSpec(), yeto_policy_sync=False) - assert allocations == [] and subprocess_agent.configured()["provider"] is None - monkeypatch.delenv("YETO_CODEX_COMPACTION_ENABLED") - - # 7.1: reward_scope=segment fails at startup through the same hook (and in validate_parsed_args) - with pytest.raises(MilesConfigError, match="reward_scope='segment'"): - entry.preflight_stage(_island_args(yeto_harness_reward_scope="segment"), "launch", AlgorithmSpec(), yeto_policy_sync=False) - assert allocations == [] and subprocess_agent.configured()["provider"] is None - - # wrong agent function / no provider -> fail closed - with pytest.raises(preflight.PreflightError, match="expects"): - entry.preflight_stage(_island_args(custom_agent_function_path="x.run"), "launch", AlgorithmSpec(), yeto_policy_sync=False) - with pytest.raises(preflight.PreflightError, match="no environment provider"): - preflight.harness_preflight(_island_args(yeto_harness_environment_provider=None), "launch", env=_good_env()) - - # success: provider + island boards installed (lazy actors named like entry.harness_source uses) - monkeypatch.setattr(preflight, "preflight_codex_openenv", lambda env=None, **kw: attested.append("ok") or {}) - assert entry.preflight_stage(_island_args(yeto_rl_learner_id=3), "launch", AlgorithmSpec(), yeto_policy_sync=False) == ("fp", "caps", "profile", None) - installed = subprocess_agent.configured() - assert isinstance(installed["provider"], _Provider) and attested[-1] == "ok" - assert installed["tool_wait_board"].learner_id == 3 and installed["harness_board"].learner_id == 3 - from yeto.rl.engine.tool_wait import harness_board_actor - - assert installed["harness_board"]._factory is harness_board_actor and installed["tool_wait_board"]._factory is None - assert allocations == [] # preflight_stage never connects; run_ports_island does after it returns - subprocess_agent.configure(provider=None) - - -def test_ir4_schema_names_match_the_harness_payload_keys(): - # IR-4 as implemented: counters are registered without a *_total suffix (kind=counter carries the semantics). - assert set(HARNESS_METRIC_KEYS) == {"harness_in_flight", "env_live", "tito_session_mismatch", "tito_chain_breaks", "policy_age_violation"} - assert all(not k.endswith("_total") for k in LOAD_SAMPLE_SCHEMA) - assert LOAD_SAMPLE_SCHEMA["policy_age_violation"][0] == "counter" and LOAD_SAMPLE_SCHEMA["env_live"][0] == "gauge" - assert generate_wrapper.POLICY_AGE_KEY == "policy_age_violation" == rollout_meta_hook.POLICY_AGE_VIOLATION_KEY - - -# ---------------------------------------------------------------- reward validation at three points (3) - -def _outcome_metadata(reward: float = 1.0, **over) -> dict[str, Any]: - md = tbench_outcome.build_outcome_metadata( - task_id="fix-git", sample_id="s1", episode_id="ep1", status="completed", reward=reward, - verifier=tbench_outcome.TEST_SH_VERIFIER, testsh_rc=0 if reward else 1, - ) - md.update(over) - return md - - -def test_reward_is_validated_at_reward_function_parent_and_evidence(): - good = SimpleNamespace(metadata=_outcome_metadata(), status=None, rollout_id=1) - assert _run(tbench_reward.reward_func(None, [good])) == [1.0] - tampered = copy.deepcopy(good.metadata) - tampered[tbench_outcome.OUTCOME_KEY]["reward"] = 0.0 # pass bit now disagrees - for bad in (tampered, {**_outcome_metadata(), "extra": 1} | {tbench_outcome.OUTCOME_KEY: {**_outcome_metadata()[tbench_outcome.OUTCOME_KEY], "x": 1}}): - with pytest.raises(tbench_outcome.UntrustedTBenchOutcome): - _run(tbench_reward.reward_func(None, SimpleNamespace(metadata=bad, status=None))) - with pytest.raises(tbench_outcome.UntrustedTBenchOutcome): - tbench_outcome.verified_outcome(bad) # parent-side verification point - from yeto.rl import trajectory_evidence - - # evidence point: a legacy MAC field is ignored, structure is still checked - legacy = {**good.metadata, "tbench_trusted_outcome_hmac": "0" * 64} - assert trajectory_evidence._verified_outcome(legacy, "terminal-bench-2.1").reward == 1.0 - with pytest.raises(tbench_outcome.UntrustedTBenchOutcome): - trajectory_evidence._verified_outcome(tampered, "terminal-bench-2.1") - - -# ---------------------------------------------------------------- siblings share one reward (5) - -def test_sibling_segments_share_reward_and_count_once(monkeypatch): - a = SimpleNamespace(metadata=_outcome_metadata(1.0, trajectory_id="t1"), status=None, rollout_id=3) - b = SimpleNamespace(metadata=_outcome_metadata(1.0, trajectory_id="t1"), status=None, rollout_id=3) - c = SimpleNamespace(metadata=_outcome_metadata(0.0, trajectory_id="t2"), status=None, rollout_id=4) - rewards = tbench_reward.check_siblings([a, b, c]) - assert rewards == {("rollout", 3): 1.0, ("rollout", 4): 0.0} # two trajectories, not three samples - assert tbench_reward.check_group(None, [[a, b], c]).keep is True - assert tbench_reward.check_group(None, [[a, b]]).keep is False # zero std over trajectories - b_bad = SimpleNamespace(metadata=_outcome_metadata(0.0, trajectory_id="t1"), status=None, rollout_id=3) - with pytest.raises(tbench_outcome.UntrustedTBenchOutcome, match="different outcomes"): - tbench_reward.check_group(None, [[a, b_bad], c]) - infra = SimpleNamespace(metadata={tbench_reward.INFRASTRUCTURE_KEY: "x"}, status=None) - out = tbench_reward.check_group(None, [a, infra]) - assert out.keep is False and out.reason == "group_has_aborted" and infra.status == "ABORTED" - - -def _sample(tokens, mask, logprobs, response_length, spans=None, **extra): - wv = [SimpleNamespace(spans=[SimpleNamespace(version=v, abs_start=s, abs_end=e) for v, s, e in spans])] if spans else [] - return SimpleNamespace(tokens=tokens, loss_mask=mask, rollout_log_probs=logprobs, response_length=response_length, - weight_versions=wv, metadata={}, status=None, **extra) - - -def test_generate_wrapper_sets_sibling_keys_checks_policy_version_and_alignment(): - inp = SimpleNamespace(sample=SimpleNamespace(group_index=9, index=42, metadata={"expected_policy_version": "pv"})) - s1 = _sample([1, 2, 3, 4], [0, 0, 1, 1], [-0.1, -0.2], 2, [("pv", 2, 4)]) - s2 = _sample([1, 5, 6], [0, 1, 1], [-0.3, -0.4], 2, [("pv", 1, 3)]) - - async def upstream(_input): - return SimpleNamespace(samples=[s1, s2]) - - out = _run(generate_wrapper.generate(inp, upstream=upstream)) - assert [s.group_index for s in out.samples] == [9, 9] and [s.rollout_id for s in out.samples] == [42, 42] - assert [s.metadata["chain_index"] for s in out.samples] == [0, 1] and s1.metadata["chains_total"] == 2 - assert s1.status is None and s1.metadata["policy_versions_actual"] == ["pv"] - - stale = _sample([1, 2, 3], [0, 1, 1], [-0.1, -0.2], 2, [("pv-old", 1, 3)]) - _run(generate_wrapper.generate(inp, upstream=lambda _i: _coro(SimpleNamespace(samples=[stale])))) - assert stale.status == "ABORTED" and stale.metadata["policy_age_violation"] == 1 - assert tbench_outcome.OUTCOME_KEY not in stale.metadata and "policy_age" in stale.metadata[tbench_reward.INFRASTRUCTURE_KEY] - - -async def _coro(value): - return value - - -# ---------------------------------------------------------------- mask / token / logprob alignment (4) - -def test_alignment_contract_fails_closed(): - good = _sample([7, 8, 9, 10, 11], [0, 0, 1, 0, 1], [-0.5, 0.0, -0.1], 3, [("v", 2, 3), ("v", 4, 5)]) - alignment.assert_sample_alignment(good) - cases = { - "neither": _sample([7, 8, 9], [1, 1], [-0.1, -0.1, -0.1], 3), - "rollout_log_probs": _sample([7, 8, 9], [0, 1, 1], None, 2), - "prompt position": _sample([7, 8, 9], [1, 0, 1], [-0.1, -0.1], 2), - "finite rollout logprob": _sample([7, 8, 9], [0, 1, 1], [None, -0.1], 2), - "outside generated": _sample([7, 8, 9, 10], [0, 1, 1, 1], [-0.1, -0.1, -0.1], 3, [("v", 1, 3)]), - } - for match, sample in cases.items(): - with pytest.raises(alignment.AlignmentError, match=match): - alignment.assert_sample_alignment(sample) - - -# ---------------------------------------------------------------- 2.2 legacy forwarding / 6.3 boundary - -def test_legacy_preflight_forwarder_matches_legacy_failure_classes(monkeypatch): - env = dict(adapter._OPENENV_IDENTITY_ENV) - preflight.forward_legacy_openenv_preflight(None, "qwen35_08b", env) - with pytest.raises(ValueError, match="requires backend profile"): - preflight.forward_legacy_openenv_preflight(None, "deepseekv4", env) - with pytest.raises(ValueError, match="environment drifted"): - preflight.forward_legacy_openenv_preflight(None, "qwen35_08b", {k: v for k, v in env.items() if "MODEL_REVISION" not in k}) - pins = preflight.required_pin_updates() - assert pins["CODEX_HARNESS_AGENT_SHA256"] == "459855e8845ab3b3cc1a4e2f47e5a9e8cc6cad79197208e43aa1d822a7c300ad" - assert pins["CODEX_OPENENV_AGENT"].endswith("codex_openenv_subprocess_agent_function.run") - - -# ---------------------------------------------------------------- A-T3-4: rollout workers configure themselves - -def _worker_provider_factory(miles_args): - assert miles_args.yeto_rl_learner_id == 5 - return _Provider(FakeTerminalEnvironment()) - - -def test_run_configures_the_rollout_worker_from_its_environment(monkeypatch): - """codex-smoke-20261003-6: every trajectory failed with "configure(provider=...) was not - called" because upstream Miles calls the agent function inside RolloutExecutor actors.""" - subprocess_agent.configure(provider=None) - monkeypatch.delenv(preflight.ENVIRONMENT_PROVIDER_ENV, raising=False) - # no provider in the environment -> still fails closed, nothing installed - with pytest.raises(RuntimeError, match="configure\\(provider=...\\) was not called"): - asyncio.run(subprocess_agent.run("http://x", "p", metadata={"task_id": "fix-git", "trajectory_id": "t"})) - assert subprocess_agent.configured()["provider"] is None - - monkeypatch.setenv(preflight.ENVIRONMENT_PROVIDER_ENV, f"{__name__}:_worker_provider_factory") - monkeypatch.setenv(preflight.LEARNER_ID_ENV, "5") - monkeypatch.setenv(preflight.MEMBER_CELL_ENV, "cell9") - assert preflight.configure_rollout_worker() is True - installed = subprocess_agent.configured() - assert isinstance(installed["provider"], _Provider) - assert installed["tool_wait_board"].learner_id == 5 and installed["harness_board"].learner_id == 5 - assert installed["member"] == "engine:cell9" - - # run() itself performs that configure when the module is still bare (and then - # refuses for IR-3 on the missing policy token, i.e. it got past the provider check) - subprocess_agent.configure(provider=None) - monkeypatch.setattr(subprocess_agent, "resolve_expected_policy_version", lambda metadata: None) - with pytest.raises(subprocess_agent.PolicyVersionMissing): - asyncio.run(subprocess_agent.run("http://x", "p", metadata={"task_id": "fix-git", "trajectory_id": "t"})) - assert isinstance(subprocess_agent.configured()["provider"], _Provider) - subprocess_agent.configure(provider=None) - - -def test_repeated_provider_failures_tear_the_rollout_down_instead_of_resampling(monkeypatch): - """A-T3-5: Miles has no abort cap (-7 looped 17k ABORTED samples); after the - threshold the worker raises a BaseException that escapes Miles' except Exception.""" - env = FakeTerminalEnvironment(passed=True) - provider, _ = _configure(monkeypatch, env, fail=True) - monkeypatch.setenv(subprocess_agent.PROVIDER_OUTAGE_THRESHOLD_ENV, "3") - for _ in range(2): - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - assert tbench_reward.INFRASTRUCTURE_KEY in result - with pytest.raises(subprocess_agent.EnvironmentProviderOutage, match="3 consecutive"): - _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - assert not isinstance(subprocess_agent.EnvironmentProviderOutage("x"), Exception) - - # a success resets the counter; injected create_fail faults never count - provider.fail = False - assert tbench_outcome.OUTCOME_KEY in _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - from yeto.rl.harness.codex.tb2_provider import InjectedCreateFailure - - async def injected(task_id, trajectory_id): - raise InjectedCreateFailure("injected") - - provider.acquire = injected - for _ in range(5): - assert tbench_reward.INFRASTRUCTURE_KEY in _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - subprocess_agent.configure(provider=None) - - -def test_tool_wait_board_calls_are_awaited_in_order_and_never_fail_the_trajectory(monkeypatch): - """A-T3-7 (codex-smoke-20261003-11): the boards are max_concurrency actors, so - fire-and-forget exit/enter of successive tool calls raced; the agent now awaits each.""" - env = FakeTerminalEnvironment(passed=True) - - class _Ref: # stands in for ray.ObjectRef - def __init__(self, value): - self.value = value - - _Ref.__name__ = "ObjectRef" - order: list[tuple[str, str]] = [] - - class _RemoteBoard: - def __init__(self): - self.enter = SimpleNamespace(remote=lambda tid: order.append(("enter", tid)) or _Ref(1)) - self.exit = SimpleNamespace(remote=lambda tid: order.append(("exit", tid)) or _Ref(2)) - - resolved: list[Any] = [] - monkeypatch.setattr(subprocess_agent, "_board_resolve", lambda ref: resolved.append(ref) or ref.value) - monkeypatch.setenv("YETO_CODEX_OPENENV_ALLOW_SCRIPTED_DRIVER", "1") - subprocess_agent.configure(provider=_Provider(env), tool_wait_board=_RemoteBoard(), harness_board=None, member="m0") - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - assert tbench_outcome.OUTCOME_KEY in result - kinds = [k for k, _ in order] - assert kinds and kinds == ["enter", "exit"] * (len(kinds) // 2) # strictly alternating, same trajectory - assert len(resolved) == len(order) # every board call was awaited - - class _FailingBoard: - def __init__(self): - self.enter = SimpleNamespace(remote=lambda tid: (_ for _ in ()).throw(RuntimeError("board down"))) - self.exit = SimpleNamespace(remote=lambda tid: (_ for _ in ()).throw(RuntimeError("board down"))) - - subprocess_agent.configure(provider=_Provider(env), tool_wait_board=_FailingBoard(), harness_board=None, member="m0") - assert tbench_outcome.OUTCOME_KEY in _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - subprocess_agent.configure(provider=None) - - -# ---------------------------------------------------------------- rl-fn-codex-rollout 1.0: runtime backend profile - -def test_openenv_profile_pin_relaxed_to_allowlist(monkeypatch): - from yeto.rl import codex_backend as cb - from yeto.rl.harness.codex import pins - - assert adapter.IMAGE_BACKEND_PROFILE_NAME == pins.OPENENV_BACKEND_PROFILE == "qwen35_08b" - assert set(pins.OPENENV_BACKEND_PROFILES) == {"qwen35_08b", "qwen35", "qwen38", "qwen38_next", "qwen38_next_4layer", "qwen3_30b_a3b"} # rlf-codex-rl 5.6; S24 27B - # the identity derives from the selected profile (codex_backend allowlist), per profile - assert adapter.profile_identity("qwen35_08b") == (cb.QWEN35_08B_MODEL, cb.QWEN35_08B_REVISION) - assert adapter.profile_identity("qwen38_next_4layer") == (cb.QWEN38_NEXT_4LAYER_MODEL, cb.QWEN38_NEXT_4LAYER_REVISION) - assert adapter.profile_identity("qwen38_next") == (cb.QWEN38_NEXT_MODEL, cb.QWEN38_NEXT_REVISION) - assert adapter.profile_identity("qwen35") == (cb.QWEN35_MODEL, cb.QWEN35_REVISION) # S17 M1: 4B allowlisted - assert adapter.profile_identity("qwen38") == (cb.QWEN38_MODEL, cb.QWEN38_REVISION) # S24: 27B allowlisted - with pytest.raises(ValueError, match="requires backend profile"): - adapter.profile_identity("deepseekv4") - # drift of an allowlisted profile is still refused - drifted = dict(cb.stock_codex_backend_profile("qwen38_next_4layer"), model_revision="0" * 40) - monkeypatch.setattr(adapter, "stock_codex_backend_profile", lambda name: drifted) - with pytest.raises(RuntimeError, match="profile drifted"): - adapter.profile_identity("qwen38_next_4layer") - - -def test_openenv_runtime_profile_resolution(): - from yeto.rl.codex_backend import stock_codex_backend_contract - - # the launch declaration the learner publishes (YETO_CODEX_CHAT_TEMPLATE) names the profile - for name in adapter.OPENENV_BACKEND_PROFILES: - assert stock_codex_backend_contract(name, 1)["chat_template"] == name - assert adapter.resolve_backend_profile(env={adapter.RUNTIME_PROFILE_ENV: name}) == name - assert adapter.resolve_backend_profile(name, env={}) == name - assert adapter.resolve_backend_profile(env={}) == "qwen35_08b" - assert adapter.resolve_backend_profile(env={adapter.RUNTIME_PROFILE_ENV: "deepseekv4"}) == "qwen35_08b" - with pytest.raises(ValueError, match="requires backend profile"): - adapter.resolve_backend_profile("deepseekv4", env={}) - # the image record is independent of the runtime profile - assert adapter._OPENENV_IDENTITY_ENV["YETO_CODEX_OPENENV_BACKEND_PROFILE"] == "qwen35_08b" - assert adapter._OPENENV_IDENTITY_ENV["YETO_CODEX_OPENENV_MODEL_ID"] == "Qwen/Qwen3.5-0.8B" - - -def test_legacy_preflight_accepts_fn_profiles_against_image_env(): - env = dict(adapter._OPENENV_IDENTITY_ENV) # the container env the image line bakes (qwen35_08b record) - for name in ("qwen38_next_4layer", "qwen38_next", "qwen35_08b"): - preflight.forward_legacy_openenv_preflight(None, name, env) - with pytest.raises(ValueError, match="requires backend profile"): - preflight.forward_legacy_openenv_preflight(None, "deepseekv4", env) - # the image record is still enforced for the FN profile - with pytest.raises(ValueError, match="environment drifted"): - preflight.forward_legacy_openenv_preflight( - None, "qwen38_next_4layer", {**env, "YETO_CODEX_OPENENV_BACKEND_PROFILE": "qwen38_next_4layer"}) - with pytest.raises(ValueError, match="environment drifted"): - preflight.forward_legacy_openenv_preflight(None, "qwen38_next_4layer", {k: v for k, v in env.items() if "MODEL_ID" not in k}) - - -def test_subprocess_run_records_phase_timing(monkeypatch): - """agentic-rollout-utilization 1.3: per-turn generation/tool seconds sum to the - worker wall time; judge = evaluate_time; sandbox cold start and wall start/end.""" - from yeto.rl.adapters.miles.rollout_meta_hook import trajectory_diagnostics - from yeto.rl.engine.trajectory_timing import phase_totals - - env = FakeTerminalEnvironment(passed=True) - _configure(monkeypatch, env) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata())) - record = trajectory_diagnostics(result) - assert len(record["turn_tool_seconds"]) == 3 # two commands + submit - assert len(record["turn_generation_seconds"]) == 4 - totals = phase_totals(record) - assert totals["generation_seconds"] + totals["tool_seconds"] == pytest.approx( - record["worker_seconds"], abs=0.01) - assert totals["judge_seconds"] is not None and totals["judge_seconds"] >= 0 - assert record["sandbox_start_seconds"] >= 0 - assert record["trajectory_started_at"] <= record["trajectory_ended_at"] - - -# ---------------------------------------------------------------- S24 T-TB2: infrastructure never reaches reward or advantage - -class ConflictError(Exception): - """Stand-in for ``modal.exception.ConflictError`` ("Modal Sandbox is shutting down").""" - - -class NotFoundError(Exception): - """Stand-in for ``modal.exception.NotFoundError`` (sandbox already shut down).""" - - -def _assert_infrastructure(result: dict[str, Any], reason: str) -> None: - assert tbench_reward.INFRASTRUCTURE_KEY in result and tbench_outcome.OUTCOME_KEY not in result - assert reason in result[tbench_reward.INFRASTRUCTURE_KEY] - sample = SimpleNamespace(metadata=result, status=None) - assert _run(tbench_reward.reward_func(None, sample)) == 0.0 and sample.status == "ABORTED" - assert "success" not in result - - -def test_sandbox_shutdown_at_acquire_is_infrastructure(monkeypatch): - provider, _ = _configure(monkeypatch, FakeTerminalEnvironment(passed=True)) - - async def gone(task_id, trajectory_id): - raise ConflictError("Modal Sandbox is shutting down.") - - monkeypatch.setattr(provider, "acquire", gone) - _assert_infrastructure(_run(subprocess_agent.run("http://miles", "p", {}, _metadata())), "ConflictError") - - -@pytest.mark.parametrize("exc,reason", [ - (lambda: subprocess_agent.adapter.harness.CodexHarnessError( - "terminal_exec output exceeded its signed boundary"), "exceeded its signed boundary"), - (lambda: NotFoundError("Modal Sandbox with container ID ta-x not found"), "worker: NotFoundError"), -]) -def test_worker_failures_are_infrastructure_and_release_the_sandbox(monkeypatch, exc, reason): - provider, board = _configure(monkeypatch, FakeTerminalEnvironment(passed=True)) - - async def fail(*_a, **_k): - raise exc() - - monkeypatch.setattr(subprocess_agent, "_drive_worker", fail) - _assert_infrastructure(_run(subprocess_agent.run("http://miles", "p", {}, _metadata())), reason) - assert provider.destroyed == 1 and board.open == 0 - - -def test_verifier_sandbox_gone_is_infrastructure_not_a_crash(monkeypatch): - env = FakeTerminalEnvironment(passed=True) - provider, board = _configure(monkeypatch, env) - - async def gone(*_a, **_k): - raise NotFoundError("Modal Sandbox with container ID ta-x not found") - - monkeypatch.setattr(subprocess_agent.adapter, "finish_trusted", gone) - _assert_infrastructure(_run(subprocess_agent.run("http://miles", "p", {}, _metadata())), "verify: NotFoundError") - assert provider.destroyed == 1 and board.open == 0 - - -def test_a_group_with_an_infrastructure_member_never_reaches_the_advantage(): - """The TB2 group filter (``--dynamic-sampling-filter-path ...tbench_reward.check_group``) drops the - whole group before Miles computes rewards-to-advantages; Miles itself does not mask ABORTED samples.""" - from yeto.rl.harness.codex.codex_openenv_agent_function import infrastructure_metadata - - def passed(i): - env = FakeTerminalEnvironment(passed=bool(i % 2)) - return dict(_run(subprocess_agent.adapter.finish_trusted( - {"status": "completed", "metrics": {}, "episode_id": f"e{i}"}, env, task_id="fix-git", - sample_id=f"s{i}"))) - - good = [SimpleNamespace(metadata={**passed(i), "trajectory_id": f"t{i}"}, status=None, rollout_id=None) - for i in range(3)] - infra = SimpleNamespace(metadata={**infrastructure_metadata("verify: NotFoundError: gone", episode_id="e9"), - "trajectory_id": "t9"}, status=None, rollout_id=None) - assert tbench_reward.check_group(None, good).keep is True # mixed rewards: kept - out = tbench_reward.check_group(None, [*good, infra]) - assert out.keep is False and out.reason == "group_has_aborted" and infra.status == "ABORTED" - - -def test_tb2_group_check_is_a_bound_neutral_filter(): - import json as _json - - from yeto.rl.adapters.miles import binding - from yeto.rl.engine.algorithm import TBENCH_GROUP_CHECK_FILTER, AlgorithmSpec, normalize_filter_name - - assert normalize_filter_name("yeto.rl.harness.codex.tbench_reward.check_group") == TBENCH_GROUP_CHECK_FILTER - assert binding.filter_argv(TBENCH_GROUP_CHECK_FILTER) == [ - "--dynamic-sampling-filter-path", "yeto.rl.harness.codex.tbench_reward.check_group"] - from importlib import import_module - - module, _, name = binding.filter_path(TBENCH_GROUP_CHECK_FILTER).rpartition(".") - assert getattr(import_module(module), name) is tbench_reward.check_group - raw = _json.loads((Path(__file__).resolve().parents[1] - / "examples/rl_algorithms/decoupled-ppo-c2-agentic-tb2.json").read_text()) - spec = AlgorithmSpec.from_json(_json.dumps(raw)) if hasattr(AlgorithmSpec, "from_json") else AlgorithmSpec.from_dict(raw) - assert spec.sampling.filter == TBENCH_GROUP_CHECK_FILTER and spec.sampling.over_sampling_batch_size == 8 - assert spec.correction.upper == 2.0 and spec.execution.max_policy_staleness == 4 diff --git a/tests/test_harness_codex_preflight_member.py b/tests/test_harness_codex_preflight_member.py deleted file mode 100644 index b1b2c3ae2..000000000 --- a/tests/test_harness_codex_preflight_member.py +++ /dev/null @@ -1,106 +0,0 @@ -"""T3 G4: the Codex harness preflight derives the INFRA rollout member key.""" - -from __future__ import annotations - -from types import SimpleNamespace - -from yeto.rl.adapters.miles.rollout import member_id -from yeto.rl.adapters.miles import entry -from yeto.rl.harness.codex import preflight - - -def test_resolve_member_accepts_member_key_cell_attr_or_env(): - # explicit full member key - assert preflight.resolve_member(SimpleNamespace(yeto_rl_member_id="engine:cell-7"), {}) == "engine:cell-7" - # bare cell id in yeto_rl_member_id -> INFRA key - assert preflight.resolve_member(SimpleNamespace(yeto_rl_member_id="cell-7"), {}) == member_id("cell-7") - # cell id attr set by entry/INFRA - assert preflight.resolve_member(SimpleNamespace(yeto_rl_cell_id=3), {}) == "engine:3" - # environment fallback - assert preflight.resolve_member(SimpleNamespace(), {preflight.MEMBER_CELL_ENV: "c0"}) == "engine:c0" - # nothing -> global admission key (single-island smoke) - assert preflight.resolve_member(SimpleNamespace(), {}) is None - assert preflight.resolve_member(SimpleNamespace(yeto_rl_member_id=""), {}) is None - - -def test_harness_preflight_configures_the_member(monkeypatch): - seen = {} - from yeto.rl.harness.codex import codex_openenv_subprocess_agent_function as agent - - monkeypatch.setattr(agent, "configure", lambda **kw: seen.update(kw)) - monkeypatch.setattr(preflight, "preflight_codex_openenv", lambda env: {}) - monkeypatch.setattr(preflight, "resolve_environment_provider", lambda a, e: object()) - monkeypatch.setattr(preflight, "island_boards", lambda a: (None, None)) - args = SimpleNamespace(custom_agent_function_path=preflight.EXPECTED_AGENT_FUNCTION, - yeto_harness_reward_scope="trajectory", yeto_rl_cell_id="cell-1") - preflight.harness_preflight(args, None, env={}) - assert seen["member"] == "engine:cell-1" - - -# --- rl-fn-codex-rollout 0.6 (A16 / D5): the island entry publishes its own cell --- - -def test_entry_publishes_cell_from_learner_id_single_island_is_engine_0(): - env = {} - args = SimpleNamespace(yeto_rl_learner_id=0) - assert entry.publish_member_cell(args, 0, env) == "0" - assert args.yeto_rl_cell_id == "0" and env[preflight.MEMBER_CELL_ENV] == "0" - assert preflight.resolve_member(args, {}) == "engine:0" # field, no env needed - assert preflight.resolve_member(SimpleNamespace(), env) == "engine:0" # env fallback - assert preflight.worker_runtime_env(args, env)[preflight.MEMBER_CELL_ENV] == "0" - - -def test_entry_publishes_distinct_cells_for_two_islands_and_admission_is_per_member(): - from yeto.rl.engine.tool_wait import HarnessBoard - - envs = [{}, {}] - islands = [SimpleNamespace(yeto_rl_learner_id=i) for i in range(2)] - members = [preflight.resolve_member(a, e) for a, e in zip(islands, envs) - for _ in [entry.publish_member_cell(a, i, e) for i in [a.yeto_rl_learner_id]]] - assert members == ["engine:0", "engine:1"] - assert [e[preflight.MEMBER_CELL_ENV] for e in envs] == ["0", "1"] - hb = HarnessBoard() # the shared fake board: closing one island's member leaves the other open - hb.close_admission([members[1]]) - assert hb.allow_new_session(members[1]) is False - assert hb.allow_new_session(members[0]) is True - - -def test_entry_keeps_an_explicit_cell_and_falls_back_to_learner_field(): - env = {} - args = SimpleNamespace(yeto_rl_cell_id="c7", yeto_rl_learner_id=3) - assert entry.publish_member_cell(args, 3, env) == "c7" and env[preflight.MEMBER_CELL_ENV] == "c7" - args = SimpleNamespace(yeto_rl_learner_id=2) - assert entry.publish_member_cell(args, environ=env) == "2" - assert preflight.resolve_member(args, {}) == "engine:2" - - -def test_configure_rollout_worker_reads_cell_from_env(monkeypatch): - seen = {} - from yeto.rl.harness.codex import codex_openenv_subprocess_agent_function as agent - - monkeypatch.setattr(agent, "configure", lambda **kw: seen.update(kw)) - monkeypatch.setattr(preflight, "resolve_environment_provider", lambda a, e: object()) - monkeypatch.setattr(preflight, "island_boards", lambda a: (None, None)) - assert preflight.configure_rollout_worker({preflight.ENVIRONMENT_PROVIDER_ENV: "x:y", - preflight.MEMBER_CELL_ENV: "1"}) is True - assert seen["member"] == "engine:1" - assert preflight.configure_rollout_worker({}) is False - - -def test_preflight_stage_publishes_engine_0_on_a_single_island(monkeypatch): - from yeto.rl.engine.algorithm import AlgorithmSpec - - for name, val in [("ports_runtime_fingerprint", lambda launch: "fp"), - ("miles_capabilities", lambda fp, unverified_mechanisms=(): "caps"), - ("with_partitioned_serial", lambda caps: caps), - ("execution_profile_for", lambda *a, **k: "profile"), - ("expected_algorithm_sha256", lambda a: None), - ("preflight", lambda profile, algorithm, caps: None), - ("elastic_wiring_for", lambda a, profile, fingerprint: None), - ("resolve_harness_preflight", lambda a, environ=None: None)]: - monkeypatch.setattr(entry, name, val) - monkeypatch.delenv(preflight.MEMBER_CELL_ENV, raising=False) - args = SimpleNamespace(yeto_rl_learner_id=0) - entry.preflight_stage(args, "launch", AlgorithmSpec(), yeto_policy_sync=False) - import os - assert args.yeto_rl_cell_id == "0" and os.environ[preflight.MEMBER_CELL_ENV] == "0" - assert preflight.resolve_member(args) == "engine:0" diff --git a/tests/test_harness_session_protocol.py b/tests/test_harness_session_protocol.py deleted file mode 100644 index e9e112806..000000000 --- a/tests/test_harness_session_protocol.py +++ /dev/null @@ -1,88 +0,0 @@ -"""codex harness against the session-service protocol test double (decoupling 3.8). - -The trusted-layer session code (``codex_openenv_agent_function``) talks real -HTTP to ``tests/session_service_double.py`` instead of a Miles session server; -the protocol is ``yeto/rl/harness/session_protocol.md``. CPU only, localhost. -""" - -from __future__ import annotations - -import asyncio - -import aiohttp -import pytest - -from session_service_double import SessionServiceDouble -from yeto.rl.harness.codex import codex_openenv_agent_function as agent_fn -from yeto.rl.harness.codex import compaction_bridge - - -def _run(coro): - return asyncio.run(coro) - - -async def _main_session(double: SessionServiceDouble) -> str: - async with aiohttp.ClientSession() as http: - async with http.post(f"{double.router}/sessions", json={}) as reply: - assert reply.status == 200 - return double.session_url((await reply.json())["session_id"]) - - -def test_create_chat_collect_delete_through_the_protocol(): - async def scenario(): - async with SessionServiceDouble() as double: - base = await _main_session(double) - ids, urls = await agent_fn.create_segment_sessions( - base, {"temperature": 0.7, "top_p": 0.9, "top_k": 20, "max_tokens": 5}, 2) - assert urls == [double.session_url(i) for i in ids] - for sid in ids: # sampling defaults went into CreateSessionRequest, typed - assert double.sessions[sid].params == {"evaluation": False, "temperature": 0.7, - "top_p": 0.9, "top_k": 20} - async with aiohttp.ClientSession() as http: - async with http.post(f"{urls[0]}/v1/chat/completions", - json={"messages": [{"role": "user", "content": "hi"}]}) as r: - body = await r.json() - async with http.post(f"{urls[0]}/samples", json={"max_seq_len": 64}) as r: - samples = await r.json() - choice = body["choices"][0] - assert [t[1] for t in choice["meta_info"]["output_token_logprobs"]] == [11, 12, 13] - assert double.sessions[ids[0]].records[0]["request"]["logprobs"] is True - assert samples["samples"][0]["tokens"] == [11, 12, 13] - # Unreturned segments are released; a second release sees 404 and is still fine. - segments = {compaction_bridge.SESSIONS_METADATA_KEY: ids} - assert await agent_fn.release_unreturned_segments(base, segments) == [] - assert double.deleted == ids - assert await agent_fn.delete_segment_sessions(base, ids) == [] - _run(scenario()) - - -def test_half_failed_pre_creation_deletes_what_it_made(): - async def scenario(): - async with SessionServiceDouble() as double: - base = await _main_session(double) - calls = {"n": 0} - - async def flaky_post(url, body): - calls["n"] += 1 - if calls["n"] == 2: - raise RuntimeError("router went away") - return await agent_fn._post_json(url, body) - - with pytest.raises(RuntimeError, match="router went away"): - await agent_fn.create_segment_sessions(base, {}, 3, post=flaky_post) - assert double.deleted == ["s2"] # the one segment created before the failure - assert set(double.sessions) == {"s1"} # only the main session remains - _run(scenario()) - - -def test_bad_create_body_is_rejected_by_the_double(): - async def scenario(): - async with SessionServiceDouble() as double: - base = await _main_session(double) - with pytest.raises(RuntimeError, match="HTTP 400"): - await agent_fn._post_json(f"{double.router}/sessions", {"top_k": "20"}) - with pytest.raises(RuntimeError, match="no session_id"): - async def no_id(url, body): - return {} - await agent_fn.create_segment_sessions(base, {}, 1, post=no_id) - _run(scenario()) diff --git a/tests/test_harness_tb2_provider.py b/tests/test_harness_tb2_provider.py deleted file mode 100644 index d5cd82f84..000000000 --- a/tests/test_harness_tb2_provider.py +++ /dev/null @@ -1,707 +0,0 @@ -"""G1 (real TB2 EnvironmentProvider, CPU backend) and G6(a) (tito counter mirroring). - -The ``LocalProcessBackend`` stands in for a Modal sandbox: commands run in a -scratch root on this host, ``/tests`` and ``/logs/verifier`` are redirected -through ``TB2_TESTS_DIR`` / ``TB2_VERIFIER_LOGS_DIR``. Everything above the -backend (relay wire, verifier, lease, faults, trusted-side counters) is the -production code path. -""" - -from __future__ import annotations - -import asyncio -import copy -import json -import os -import textwrap -from pathlib import Path -from types import SimpleNamespace -from typing import Any - -import aiohttp -import pytest - -import test_secrlenv_codex_harness as legacy_tests -from yeto.rl import tbench_outcome -from yeto.rl.engine.tool_wait import HarnessBoard -from yeto.rl.harness.codex import ( - codex_harness_agent as harness, - codex_openenv_agent_function as adapter, - codex_openenv_subprocess_agent_function as subprocess_agent, - preflight, - tb2_provider, - tbench_reward, -) -from yeto.rl.harness.codex.environment import HttpTerminalEnvironment - - -TEST_SH = textwrap.dedent( - """\ - #!/bin/bash - L=${TB2_VERIFIER_LOGS_DIR:-/logs/verifier} - mkdir -p "$L" - if [ -f fixed.txt ] && grep -q ok fixed.txt; then - echo 1 > "$L/reward.txt"; exit 0 - else - echo 0 > "$L/reward.txt"; exit 1 - fi - """ -) - - -def _make_task(root: Path, task_id: str = "fix-git") -> Path: - task_dir = root / task_id - (task_dir / "tests").mkdir(parents=True) - (task_dir / "environment").mkdir() - (task_dir / "task.toml").write_text( - '[environment]\ndocker_image = "alexgshaw/fix-git:20251031"\ncpus = 1\nmemory_mb = 2048\n' - "[agent]\ntimeout_sec = 900.0\n[verifier]\ntimeout_sec = 60.0\n" - ) - (task_dir / "environment" / "Dockerfile").write_text("FROM ubuntu:24.04\nWORKDIR /app\n") - (task_dir / "tests" / "test.sh").write_text(TEST_SH) - (task_dir / "tests" / "test_outputs.py").write_text("# fixture\n") - return root - - -def _run(coro): - return asyncio.run(coro) - - -def _provider(tmp_path: Path, **kw: Any) -> tb2_provider.Tb2EnvironmentProvider: - tasks_dir = _make_task(tmp_path / "tb2") - backend = tb2_provider.LocalProcessBackend(tmp_path / "sandboxes") - return tb2_provider.Tb2EnvironmentProvider(backend, tasks_dir, **kw) - - -def _configure(monkeypatch, provider, harness_board: Any = None) -> None: - monkeypatch.setenv("YETO_CODEX_OPENENV_ALLOW_SCRIPTED_DRIVER", "1") - subprocess_agent.configure(provider=provider, tool_wait_board=None, harness_board=harness_board, member="m0") - - -def _metadata(script: list[str], **extra) -> dict[str, Any]: - return {"task_id": "fix-git", "trajectory_id": "traj-1", "prompt": "fix the repo", - "codex_openenv_driver": "scripted", "script": script, "expected_policy_version": "pv-7", **extra} - - -# ---------------------------------------------------------------- task resolution / faults - - -def test_resolve_task_reads_task_toml_and_dockerfile_workdir(tmp_path): - tasks_dir = _make_task(tmp_path) - task = tb2_provider.resolve_task("fix-git", tasks_dir) - assert task.docker_image == "alexgshaw/fix-git:20251031" and task.workdir == "/app" - assert task.cpus == 1 and task.memory_mb == 2048 and task.verifier_timeout_s == 60.0 - with pytest.raises(FileNotFoundError): - tb2_provider.resolve_task("nope", tasks_dir) - with pytest.raises(ValueError): - tb2_provider.resolve_task("../fix-git", tasks_dir) - - -def test_parse_faults_accepts_matrix_entries_and_rejects_unknown_kinds(): - faults = tb2_provider.parse_faults("create_fail:0, deadline:traj-9=2, max_turns:2=2") - assert [f.kind for f in faults] == ["create_fail", "deadline", "max_turns"] - assert faults[0].matches(0, "x") and not faults[0].matches(1, "x") - assert faults[1].matches(5, "traj-9") and faults[1].value == "2" - assert tb2_provider.parse_faults("") == () - with pytest.raises(ValueError, match="invalid fault"): - tb2_provider.parse_faults("explode:0") - with pytest.raises(ValueError, match="selector"): - tb2_provider.parse_faults("create_fail:") - - -# ---------------------------------------------------------------- lease lifecycle on the wire - - -def test_lease_relays_execute_submit_verifies_and_is_gone_after_destroy(tmp_path): - provider = _provider(tmp_path) - - async def scenario() -> None: - lease = await provider.acquire("fix-git", "traj-1") - assert lease.deadline_seconds == 900.0 and await lease.describe() == "live" - async with HttpTerminalEnvironment(lease.env_url, lease.env_token) as env: - out = await env.execute("ep-1", "echo hello; pwd; exit 3", timeout_seconds=10, output_bytes=4096) - assert out["exit_code"] == 3 and "hello" in out["output"] and out["output"].rstrip().endswith("/app") - assert out["timed_out"] is False and out["truncated"] is False - big = await env.execute("ep-1", "head -c 100 /dev/zero | tr '\\0' x", timeout_seconds=10, output_bytes=16) - assert big["truncated"] is True and len(big["output"]) == 16 - slow = await env.execute("ep-1", "sleep 5", timeout_seconds=1, output_bytes=64) - assert slow["timed_out"] is True - # the verifier is not reachable before evaluate: /tests is empty - probe = await env.execute("ep-1", "ls $TB2_TESTS_DIR 2>/dev/null | wc -l", timeout_seconds=10, output_bytes=64) - assert probe["output"].strip() == "0" - unknown = HttpTerminalEnvironment(lease.env_url, "wrong") - async with unknown: - with pytest.raises(Exception): - await unknown.execute("ep-1", "true", timeout_seconds=1, output_bytes=1) - assert (await env.execute("ep-1", "echo ok > fixed.txt", timeout_seconds=10, output_bytes=64))["exit_code"] == 0 - assert (await env.submit("ep-1", {"evidence": "done"})) == {"accepted": True} - evaluation = await lease.verifier.evaluate("ep-1") - assert {k: v for k, v in evaluation.items() if k != "log"} == {"passed": True, "testsh_rc": 0, "timed_out": False} - assert "YETO_TB2_REWARD=1" in evaluation["log"] - assert lease.environment.commands[0].startswith("echo hello") and lease.environment.submitted - await lease.destroy() - assert await lease.describe() == "gone" and provider.destroyed == 1 and provider.live == {} - await lease.destroy() # idempotent - assert provider.destroyed == 1 - - _run(scenario()) - - -def test_verifier_reports_failure_and_binds_to_its_episode(tmp_path): - provider = _provider(tmp_path) - - async def scenario() -> None: - lease = await provider.acquire("fix-git", "traj-2") - evaluation = await lease.verifier.evaluate("ep-1") # unbound: no command ran (e.g. timeout path) - assert evaluation["passed"] is False and evaluation["testsh_rc"] == 1 - async with HttpTerminalEnvironment(lease.env_url, lease.env_token) as env: - await env.execute("ep-1", "true", timeout_seconds=5, output_bytes=16) - with pytest.raises(Exception): # a second episode on the same sandbox is refused - await env.execute("ep-2", "true", timeout_seconds=5, output_bytes=16) - with pytest.raises(ValueError): - await lease.verifier.evaluate("other-episode") - await lease.destroy() - - _run(scenario()) - - -# ---------------------------------------------------------------- full subprocess path (R-TB matrix) - - -def test_subprocess_run_positive_and_negative_rewards_through_the_real_relay(monkeypatch, tmp_path): - hb = HarnessBoard() - provider = _provider(tmp_path) - _configure(monkeypatch, provider, harness_board=hb) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata(["ls", "echo ok > fixed.txt"]))) - outcome, reward = tbench_outcome.verified_outcome(result) - assert reward == 1.0 and outcome["status"] == "completed" and outcome["verifier"] == tbench_outcome.TEST_SH_VERIFIER - assert result["chains_total"] == 1 and result["chain_break_reason"] is None - assert provider.destroyed == 1 and provider.live == {} - snap = hb.snapshot() - assert snap.env_live == 0 and snap.in_flight == 0 and snap.tito_session_mismatch == 0 - - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata(["ls", "echo nope > fixed.txt"], trajectory_id="traj-neg"))) - outcome, reward = tbench_outcome.verified_outcome(result) - assert reward == 0.0 and outcome["status"] == "completed" and tbench_outcome.OUTCOME_KEY in result - assert provider.destroyed == 2 and hb.snapshot().env_live == 0 - - -def test_create_fail_fault_is_infrastructure_aborted_and_touches_no_sandbox(monkeypatch, tmp_path): - hb = HarnessBoard() - provider = _provider(tmp_path, faults=tb2_provider.parse_faults("create_fail:traj-1")) - _configure(monkeypatch, provider, harness_board=hb) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata(["ls"]))) - assert tbench_reward.INFRASTRUCTURE_KEY in result and tbench_outcome.OUTCOME_KEY not in result - assert "InjectedCreateFailure" in result[tbench_reward.INFRASTRUCTURE_KEY] - sample = SimpleNamespace(metadata=result, status=None) - assert _run(tbench_reward.reward_func(None, sample)) == 0.0 and sample.status == "ABORTED" - assert provider.backend.created == [] and provider.fault_log == [("create_fail", "traj-1", 0)] - assert hb.snapshot().env_live == 0 and hb.snapshot().in_flight == 0 - # ordinal selectors: the second acquire is unaffected - ok = _run(subprocess_agent.run("http://miles", "p", {}, _metadata(["echo ok > fixed.txt"], trajectory_id="traj-2"))) - assert tbench_outcome.verified_outcome(ok)[1] == 1.0 - - -def test_deadline_fault_cancels_the_episode_destroys_and_drains_env_live(monkeypatch, tmp_path): - hb = HarnessBoard() - provider = _provider(tmp_path, faults=tb2_provider.parse_faults("deadline:0=1")) - _configure(monkeypatch, provider, harness_board=hb) - seen: list[int] = [] - - async def scenario() -> dict[str, Any]: - task = asyncio.create_task(subprocess_agent.run("http://miles", "p", {}, _metadata(["sleep 30"]))) - while not hb.snapshot().env_live: - await asyncio.sleep(0.02) - seen.append(hb.snapshot().env_live) - return await task - - result = _run(scenario()) - outcome, reward = tbench_outcome.verified_outcome(result) - assert seen == [1] and reward == 0.0 and outcome["status"] == "timeout" - assert outcome["verifier"] == tbench_outcome.TIMEOUT_VERIFIER - assert provider.destroyed == 1 and not provider.backend.created[0].alive() - snap = hb.snapshot() - assert snap.env_live == 0 and snap.in_flight == 0 and snap.leases_expired_total == 0 - - -def test_max_turns_fault_sets_the_worker_turn_budget(monkeypatch, tmp_path): - provider = _provider(tmp_path, faults=tb2_provider.parse_faults("max_turns:traj-1=2")) - _configure(monkeypatch, provider) - captured: dict[str, Any] = {} - original = asyncio.create_subprocess_exec - - async def spy(*args, **kwargs): - captured["env"] = dict(kwargs["env"]) - return await original(*args, **kwargs) - - monkeypatch.setattr(asyncio, "create_subprocess_exec", spy) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata(["echo ok > fixed.txt"]))) - assert tbench_outcome.verified_outcome(result)[1] == 1.0 - assert captured["env"][tb2_provider.TURN_BUDGET_ENV] == "2" - # the budget is the bridge's signed turn limit - monkeypatch.setenv(tb2_provider.TURN_BUDGET_ENV, "2") - legacy_tests._set_bridge_env(monkeypatch) - bridge = harness._ResponsesBridge("http://miles", "p", {"temperature": 0.7, "top_p": 0.9, "max_tokens": 128, "stream": True, "n": 1}, - harness.legacy.AgentMetrics(), max_seq_len=512) - assert bridge._max_turns == 2 - - -# ---------------------------------------------------------------- factories (module:callable) - - -def test_local_provider_factory_reads_env_and_resolves_through_preflight(monkeypatch, tmp_path): - tasks_dir = _make_task(tmp_path / "tb2") - monkeypatch.delenv(tb2_provider.TASKS_DIR_ENV, raising=False) - with pytest.raises(RuntimeError, match=tb2_provider.TASKS_DIR_ENV): - tb2_provider.local_provider(None) - monkeypatch.setenv(tb2_provider.TASKS_DIR_ENV, str(tasks_dir)) - monkeypatch.setenv(tb2_provider.FAULT_ENV, "create_fail:3") - monkeypatch.setenv(tb2_provider.LEASE_SECONDS_ENV, "120") - provider = preflight.resolve_environment_provider( - SimpleNamespace(yeto_harness_environment_provider=None), - {preflight.ENVIRONMENT_PROVIDER_ENV: "yeto.rl.harness.codex.tb2_provider:local_provider"}, - ) - assert isinstance(provider, tb2_provider.Tb2EnvironmentProvider) - assert provider.lease_seconds == 120.0 and provider.faults[0].kind == "create_fail" - monkeypatch.setattr(tb2_provider, "require_modal_client", lambda: None) # no client in the test venv - modal = tb2_provider.modal_provider(None) - assert isinstance(modal.backend, tb2_provider.ModalSandboxBackend) and modal.backend.app_name == tb2_provider.DEFAULT_MODAL_APP - - -def test_modal_backend_maps_task_to_sandbox_create(monkeypatch, tmp_path): - import sys - import types - - calls: dict[str, Any] = {} - - class _Sandbox: - object_id = "sb-1" - - def __init__(self) -> None: - self.polls = 0 - - def exec(self, *args, **kwargs): - calls["exec"] = (args, kwargs) - return SimpleNamespace(stdout=SimpleNamespace(read=lambda: "out\n"), stderr=SimpleNamespace(read=lambda: ""), wait=lambda: 0) - - def terminate(self): - calls["terminated"] = True - - def poll(self): - return None if "terminated" not in calls else 0 - - @classmethod - def create(cls, *args, **kwargs): - calls["create"] = (args, kwargs) - return cls() - - fake = types.ModuleType("modal") - fake.Sandbox = _Sandbox - fake.Image = SimpleNamespace(from_registry=lambda tag: ("image", tag)) - fake.App = SimpleNamespace(lookup=lambda name, create_if_missing: ("app", name)) - monkeypatch.setitem(sys.modules, "modal", fake) - - task = tb2_provider.resolve_task("fix-git", _make_task(tmp_path)) - backend = tb2_provider.ModalSandboxBackend(app_name="t", ttl_s=600, idle_timeout_s=60, run_id="r1") - handle = backend.create(task, "traj-1") - args, kwargs = calls["create"] - assert args == ("sleep", "infinity") and kwargs["image"] == ("image", "alexgshaw/fix-git:20251031") - assert kwargs["app"] == ("app", "t") and kwargs["timeout"] == 600 and kwargs["idle_timeout"] == 60 - assert kwargs["cpu"] == 1.0 and kwargs["memory"] == 2048 and kwargs["workdir"] == "/app" - assert kwargs["tags"] == {"yeto-tb2-task": "fix-git", "yeto-trajectory": "traj-1", "yeto-run-id": "r1"} - result = handle.exec("echo out", timeout_s=5) - assert result.exit_code == 0 and result.output == "out\n" and calls["exec"][1]["workdir"] == "/app" - assert handle.alive() - handle.terminate() - assert not handle.alive() - - -# ---------------------------------------------------------------- G6(a): tito counters - - -def test_bridge_counts_session_mismatch_and_retry_fork(monkeypatch): - legacy_tests._set_bridge_env(monkeypatch) - - async def scenario() -> None: - runner, miles_url, _requests = await legacy_tests._fake_miles([ - legacy_tests._completion("terminal.exec", '{ "command": "git status" }', "call-1", "look", 25), - ]) - metrics = harness.legacy.AgentMetrics() - try: - async with harness._ResponsesBridge( - miles_url, "fix the repo", {"temperature": 0.7, "top_p": 0.9, "max_tokens": 128, "stream": True, "n": 1}, - metrics, max_seq_len=512, - ) as bridge: - headers = {"Authorization": f"Bearer {bridge.token}"} - initial = [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "fix the repo"}]}] - async with aiohttp.ClientSession() as session: - async with session.post(f"{bridge.url}/v1/responses", json=legacy_tests._codex_body(initial), headers=headers) as r: - assert r.status == 200 - assert harness.tito_counters(metrics) == {"tito_session_mismatch": 0, "tito_chain_breaks": {}} - # byte-identical retry -> retry_fork chain break - async with session.post(f"{bridge.url}/v1/responses", json=legacy_tests._codex_body(initial), headers=headers) as r: - assert r.status == 400 - mutated = copy.deepcopy(initial) - mutated[0]["content"][0]["text"] = "fix the repo!" - async with session.post(f"{bridge.url}/v1/responses", json=legacy_tests._codex_body(mutated), headers=headers) as r: - assert r.status == 400 - finally: - await runner.cleanup() - assert harness.tito_counters(metrics) == {"tito_session_mismatch": 1, "tito_chain_breaks": {"retry_fork": 1}} - assert adapter._metrics_dict(metrics)["tito_chain_breaks"] == {"retry_fork": 1} - - _run(scenario()) - - -def test_mirror_tito_counters_uses_gateway_board_semantics(): - hb = HarnessBoard() - fields = adapter.mirror_tito_counters({"tito_session_mismatch": 2, "tito_chain_breaks": {"retry_fork": 1, "history_rewrite": 0}}, hb) - assert fields == {"chain_break_reason": "retry_fork"} - snap = hb.snapshot() - assert snap.tito_session_mismatch == 2 and snap.tito_chain_breaks == {"retry_fork": 1} - assert adapter.mirror_tito_counters(None, hb) == {"chain_break_reason": None} - assert adapter.mirror_tito_counters({"tito_session_mismatch": 1}, None) == {"chain_break_reason": None} - with pytest.raises(ValueError): - adapter.mirror_tito_counters({"tito_chain_breaks": {"made_up": 1}}, hb) - - -def test_mirror_tito_counters_awaits_remote_board_calls_in_order(monkeypatch, capsys): - """S14/A19 (FINAL-REPORT-S7 §6.1, A-T3-7 leftovers): record_session_mismatch / - record_chain_break on an actor board were fire-and-forget; each call is now - resolved in issue order, and a remote failure never fails the trajectory.""" - - class _Ref: # stands in for ray.ObjectRef - def __init__(self, value): - self.value = value - - _Ref.__name__ = "ObjectRef" - order: list[tuple[str, tuple[Any, ...]]] = [] - resolved: list[Any] = [] - monkeypatch.setattr("yeto.rl.engine.tool_wait._resolve", lambda ref: resolved.append(ref) or ref.value) - - class _RemoteBoard: - def __init__(self): - self.record_session_mismatch = SimpleNamespace( - remote=lambda *a: order.append(("record_session_mismatch", a)) or _Ref(None)) - self.record_chain_break = SimpleNamespace( - remote=lambda *a: order.append(("record_chain_break", a)) or _Ref(None)) - - metrics = {"tito_session_mismatch": 2, "tito_chain_breaks": {"retry_fork": 1, "history_rewrite": 3}} - assert adapter.mirror_tito_counters(metrics, _RemoteBoard()) == {"chain_break_reason": "retry_fork"} - assert order == [("record_session_mismatch", (2,)), ("record_chain_break", ("retry_fork", 1)), - ("record_chain_break", ("history_rewrite", 3))] - assert len(resolved) == len(order) # every remote call was awaited, in issue order - - class _FailingBoard: - def __init__(self): - self.record_session_mismatch = SimpleNamespace( - remote=lambda *a: (_ for _ in ()).throw(RuntimeError("board down"))) - self.record_chain_break = SimpleNamespace( - remote=lambda *a: (_ for _ in ()).throw(RuntimeError("board down"))) - - assert adapter.mirror_tito_counters(metrics, _FailingBoard()) == {"chain_break_reason": "retry_fork"} - assert "harness board record_session_mismatch failed: RuntimeError: board down" in capsys.readouterr().err - - -def test_subprocess_rejection_counters_reach_the_board_and_abort(monkeypatch, tmp_path): - hb = HarnessBoard() - provider = _provider(tmp_path) - _configure(monkeypatch, provider, harness_board=hb) - - async def rejected(job, trajectory_id, board, worker_env=None): - error = harness.CodexHarnessError("Codex truncated or mutated episode history") - error.metrics = {"turns": 1, "tito_session_mismatch": 1, "tito_chain_breaks": {"retry_fork": 1}} - raise error - - monkeypatch.setattr(subprocess_agent, "_drive_worker", rejected) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata(["ls"]))) - assert tbench_reward.INFRASTRUCTURE_KEY in result and tbench_outcome.OUTCOME_KEY not in result - assert result["chains_total"] == 1 and result["chain_break_reason"] == "retry_fork" - assert result["agent_metrics"]["tito_session_mismatch"] == 1 - sample = SimpleNamespace(metadata=result, status=None) - assert _run(tbench_reward.reward_func(None, sample)) == 0.0 and sample.status == "ABORTED" - snap = hb.snapshot() - assert snap.tito_session_mismatch == 1 and snap.tito_chain_breaks == {"retry_fork": 1} - assert snap.env_live == 0 and snap.in_flight == 0 and provider.destroyed == 1 - - -def test_worker_error_event_carries_metrics(monkeypatch): - import subprocess - import sys - - monkeypatch.setenv("YETO_CODEX_OPENENV_ALLOW_SCRIPTED_DRIVER", "1") - - async def boom(job, env, tool_wait=None): - error = harness.CodexHarnessError("Codex retried a Responses sample") - error.metrics = {"tito_chain_breaks": {"retry_fork": 1}} - raise error - - code = ( - "import asyncio, json, sys\n" - "from yeto.rl.harness.codex import codex_openenv_agent_worker as w, codex_openenv_agent_function as a\n" - "async def boom(job, env, tool_wait=None):\n" - " e = a.harness.CodexHarnessError('Codex retried a Responses sample'); e.metrics = {'tito_chain_breaks': {'retry_fork': 1}}; raise e\n" - "a.drive_untrusted = boom\n" - "sys.exit(w.main())\n" - ) - job = {"env_url": "http://127.0.0.1:9", "env_token": "t", "episode_id": "e", "base_url": "b", "prompt": "p", "request_kwargs": {}} - proc = subprocess.run([sys.executable, "-c", code], input=json.dumps(job) + "\n", capture_output=True, text=True, - env={k: v for k, v in os.environ.items() if "HMAC" not in k}) - event = json.loads(proc.stdout.strip().splitlines()[-1]) - assert proc.returncode == 1 and event["event"] == "error" and event["metrics"] == {"tito_chain_breaks": {"retry_fork": 1}} - del boom - - -def test_modal_provider_fails_closed_without_an_importable_modal_client(monkeypatch): - """A-T3-5 (codex-smoke-20261003-7): the ports image has no Modal client; the - driver preflight must refuse instead of every rollout worker failing at acquire.""" - import sys - - monkeypatch.setitem(sys.modules, "modal", None) # import raises ImportError - with pytest.raises(RuntimeError, match="importable `modal` client"): - tb2_provider.modal_provider(None) - assert tb2_provider.InjectedCreateFailure.injected_fault is True - - -# ---------------------------------------------------------------- task statement (S15 stage-2 root cause) - - -SMOKE6_PROMPT = [{"role": "system", "content": "You are an autonomous terminal agent solving a Terminal-Bench task."}] - - -def test_resolve_task_reads_instruction_md(tmp_path): - tasks_dir = _make_task(tmp_path / "tb2") - assert tb2_provider.resolve_task("fix-git", tasks_dir).instruction is None - (tasks_dir / "fix-git" / "instruction.md").write_text("\nFind my lost changes and merge them into master.\n") - assert tb2_provider.resolve_task("fix-git", tasks_dir).instruction == "Find my lost changes and merge them into master." - - -def test_task_prompt_precedence_and_fail_closed(): - lease = SimpleNamespace(task=SimpleNamespace(instruction="TASK STATEMENT")) - no_task = SimpleNamespace(task=None) - assert subprocess_agent.task_prompt({"prompt": "explicit"}, SMOKE6_PROMPT, lease) == "explicit" - # the smoke6 row shape: system-only chat prompt, no metadata prompt -> instruction.md - assert subprocess_agent.task_prompt({"task_id": "fix-git"}, SMOKE6_PROMPT, lease) == "TASK STATEMENT" - assert subprocess_agent.task_prompt({}, "plain", no_task) == "plain" - chat = SMOKE6_PROMPT + [{"role": "user", "content": "do it"}] - assert subprocess_agent.task_prompt({}, chat, no_task) == "do it" - with pytest.raises(subprocess_agent.TaskPromptMissing): - subprocess_agent.task_prompt({"task_id": "fix-git"}, SMOKE6_PROMPT, no_task) - - -def test_subprocess_run_sends_instruction_md_not_the_stringified_chat_prompt(monkeypatch, tmp_path): - provider = _provider(tmp_path) - (provider.tasks_dir / "fix-git" / "instruction.md").write_text("Find my lost changes and merge them into master.\n") - _configure(monkeypatch, provider) - seen: dict[str, Any] = {} - original = subprocess_agent._drive_worker - - async def spy(job, *args, **kwargs): - seen["prompt"] = job["prompt"] - return await original(job, *args, **kwargs) - - monkeypatch.setattr(subprocess_agent, "_drive_worker", spy) - metadata = _metadata(["echo ok > fixed.txt"]) - del metadata["prompt"] - result = _run(subprocess_agent.run("http://miles", SMOKE6_PROMPT, {}, metadata)) - assert seen["prompt"] == "Find my lost changes and merge them into master." - assert "'role'" not in seen["prompt"] - assert tbench_outcome.verified_outcome(result)[1] == 1.0 - assert provider.destroyed == 1 and provider.live == {} - - -# ---------------------------------------------------------------- TB2 system prompt (2026-10-07) / prompt preflight - - -def test_legacy_instructions_hash_unchanged_and_tb2_prompt_pinned(monkeypatch): - import hashlib - import yeto.rl as rl_config - - monkeypatch.delenv(harness.INSTRUCTIONS_FAMILY_ENV, raising=False) - assert harness.base_instructions() is harness.BASE_INSTRUCTIONS - assert harness.BASE_INSTRUCTIONS_SHA256 == rl_config.CODEX_BASE_INSTRUCTIONS_SHA256 == ( - "1c183656ca1319142cba9e76baa199b7ab59f770a51a76660622a087e74ba846") - assert harness.codex_harness_identity()["base_instructions_sha256"] == rl_config.CODEX_BASE_INSTRUCTIONS_SHA256 - assert harness._tb2_surface_sha256() == harness.TB2_BASE_INSTRUCTIONS_SHA256 == rl_config.CODEX_TB2_BASE_INSTRUCTIONS_SHA256 - # input schemas are kept verbatim (optional "flag" property); prompt + descriptions are reworded - tb2_text = harness.TB2_BASE_INSTRUCTIONS + " ".join(t["description"] for t in harness.TB2_DYNAMIC_TOOLS) - for ctf in ("TARGET_URL", "flag", "DEBUG_URL", "security", "attack", "scored target"): - assert ctf not in tb2_text, ctf - # the prompt names the model-facing tools and requires submit - for needle in ("`terminal.exec`", "`submit`", "exactly ONE tool"): - assert needle in " ".join(harness.TB2_BASE_INSTRUCTIONS.split()) - # legacy surface unchanged: legacy tools / Miles tools - assert harness.dynamic_tools() is harness.DYNAMIC_TOOLS and harness.miles_tools() is harness._MILES_TOOLS - assert harness._miles_tools_for(harness.DYNAMIC_TOOLS) == harness._MILES_TOOLS - # compaction head bound was sized for the legacy prompt - assert len(harness.TB2_BASE_INSTRUCTIONS) <= len(harness.BASE_INSTRUCTIONS) - - -def test_tb2_instructions_require_the_signed_hash(monkeypatch): - for name, value in tb2_provider.tb2_instructions_env().items(): - monkeypatch.setenv(name, value) - assert harness.base_instructions() == harness.TB2_BASE_INSTRUCTIONS - # TB2 keeps tool names and input schemas; only descriptions change - tb2_tools = harness.miles_tools() - assert [t["function"]["name"] for t in tb2_tools] == ["terminal.exec", "submit"] - assert [t["function"]["parameters"] for t in tb2_tools] == [t["function"]["parameters"] for t in harness._MILES_TOOLS] - assert harness.dynamic_tools() is harness.TB2_DYNAMIC_TOOLS - assert [t["name"] for t in harness.dynamic_tools()] == [t["name"] for t in harness.DYNAMIC_TOOLS] - with pytest.raises(harness.CodexHarnessError, match="base instructions drifted"): - harness._validate_codex_request({"model": "m", "instructions": harness.BASE_INSTRUCTIONS}, "m") - monkeypatch.setenv(harness.TB2_INSTRUCTIONS_SHA_ENV, "0" * 64) - with pytest.raises(harness.CodexHarnessError, match="signed TB2"): - harness.base_instructions() - monkeypatch.setenv(harness.INSTRUCTIONS_FAMILY_ENV, "ctf2") - with pytest.raises(harness.CodexHarnessError, match="not a signed"): - harness.base_instructions() - - -def test_tb2_lease_selects_tb2_instructions_and_keeps_fault_turn_budget(tmp_path): - provider = _provider(tmp_path, faults=tb2_provider.parse_faults("max_turns:traj-9=3")) - - async def go(): - a = await provider.acquire("fix-git", "traj-1") - b = await provider.acquire("fix-git", "traj-9") - envs = (dict(a.worker_env), dict(b.worker_env)) - await a.destroy(); await b.destroy() - return envs - - a, b = _run(go()) - assert a == tb2_provider.tb2_instructions_env() - assert a[harness.INSTRUCTIONS_FAMILY_ENV] == "tb2" - assert b == {**a, tb2_provider.TURN_BUDGET_ENV: "3"} - - -def test_preflight_task_prompts_catches_the_stage2_dataset(tmp_path): - tasks_dir = _make_task(tmp_path / "tb2") - data = tmp_path / "smoke.jsonl" - data.write_text(json.dumps({"prompt": SMOKE6_PROMPT, "metadata": {"task_id": "fix-git"}}) + "\n") - with pytest.raises(tb2_provider.TaskPromptPreflightError, match="fix-git"): - tb2_provider.preflight_task_prompts(data, tasks_dir) - (tasks_dir / "fix-git" / "instruction.md").write_text("Merge my lost changes into master.\n") - assert tb2_provider.preflight_task_prompts(data, tasks_dir) == [("fix-git", "Merge my lost changes into master.")] - data.write_text(json.dumps({"prompt": SMOKE6_PROMPT, "metadata": {"task_id": "fix-git", "prompt": str(SMOKE6_PROMPT)}}) + "\n") - with pytest.raises(tb2_provider.TaskPromptPreflightError, match="stringified"): - tb2_provider.preflight_task_prompts(data, tasks_dir) - - -def test_harness_preflight_task_prompt_hook(tmp_path): - from yeto.rl.harness.codex import preflight as pf - - tasks_dir = _make_task(tmp_path / "tb2") - data = tmp_path / "smoke.jsonl" - data.write_text(json.dumps({"prompt": SMOKE6_PROMPT, "metadata": {"task_id": "fix-git"}}) + "\n") - provider = SimpleNamespace(tasks_dir=tasks_dir) - with pytest.raises(pf.PreflightError): - pf.assert_task_prompts(SimpleNamespace(prompt_data=str(data)), provider) - pf.assert_task_prompts(SimpleNamespace(prompt_data=str(tmp_path / "absent.jsonl")), provider) # not readable here - (tasks_dir / "fix-git" / "instruction.md").write_text("x\n") - pf.assert_task_prompts(SimpleNamespace(prompt_data=str(data)), provider) - - -def test_last_completion_and_end_reason_reach_the_metrics(monkeypatch): - metrics = harness.legacy.AgentMetrics() - choice = {"finish_reason": "stop", "message": {"role": "assistant", "content": "I will run ls." * 30, - "reasoning_content": "think", "tool_calls": [{"id": "c", "function": {"name": "terminal.exec", "arguments": "{}"}}]}} - harness._note_last_completion(metrics, choice, {"usage": {"completion_tokens": 77}}) - metrics.end_reason = "CodexModelFailure: DSV4 mixed prose with the required tool call" - d = adapter._metrics_dict(metrics) - assert d["end_reason"].startswith("CodexModelFailure") and d["last_completion"]["tool_calls"] == 1 - assert d["last_completion"]["content_chars"] == len("I will run ls." * 30) - assert len(d["last_completion"]["content_head"]) == harness.LAST_COMPLETION_HEAD_CHARS - assert d["last_completion"]["completion_tokens"] == 77 - assert "end_reason" not in adapter._metrics_dict(harness.legacy.AgentMetrics()) - - -def test_verifier_log_tail_reaches_the_trajectory_metadata_and_tape(monkeypatch, tmp_path): - from yeto.rl.adapters.miles import rollout_meta_hook as hook - - provider = _provider(tmp_path) - _configure(monkeypatch, provider) - result = _run(subprocess_agent.run("http://miles", "p", {}, _metadata(["echo nope > fixed.txt"]))) - assert tbench_outcome.verified_outcome(result)[1] == 0.0 - assert "YETO_TB2_REWARD=0" in result["verifier_log"] - diag = hook.trajectory_diagnostics(result) - assert diag["verifier_log"] == result["verifier_log"] and diag["testsh_rc"] == 1 - # S17 M1: observe-only timing line (verifier wall time, agent phase times, per-turn shapes) - timing = json.loads(result["verifier_log"].rsplit("YETO_TIMING ", 1)[1]) - assert timing["verify_s"] >= 0 - long = "x" * 5000 + "TAIL" - excerpt = tb2_provider.verifier_log_excerpt(long) - assert len(excerpt) == tb2_provider.VERIFIER_LOG_CHARS and excerpt.endswith("TAIL") - - -# --- S17 G3: tests/ larger than Modal's 64 KiB exec-argv limit --------------------- - -class _ArgvLimited: - """Wraps a sandbox handle and refuses any command longer than Modal's ARG_MAX.""" - - def __init__(self, inner): - self.inner = inner - self.lengths: list[int] = [] - - def exec(self, command, *, timeout_s, workdir=None): - self.lengths.append(len(command)) - if len(command) > 65536: - raise RuntimeError("Total length of CMD arguments cannot exceed 65536 bytes (ARG_MAX)") - return self.inner.exec(command, timeout_s=timeout_s, workdir=workdir) - - -@pytest.mark.parametrize("big", [False, True]) -def test_verifier_stages_large_tests_dir(tmp_path, big): - import os as _os - - tasks_dir = _make_task(tmp_path / "tasks") - if big: # ~200 KB of incompressible data -> ~270 KB base64 - (tasks_dir / "fix-git" / "tests" / "blob.bin").write_bytes(_os.urandom(200_000)) - task = tb2_provider.resolve_task("fix-git", tasks_dir) - stage = tb2_provider.verifier_stage_commands(task) - assert bool(stage) == big - backend = tb2_provider.LocalProcessBackend(tmp_path / "sandboxes") - handle = _ArgvLimited(backend.create(task, "t1")) - handle.exec("echo ok > fixed.txt", timeout_s=30) - verdict = _run(tb2_provider.Tb2Verifier(handle, task).evaluate("e")) - assert verdict["passed"] is True and verdict["testsh_rc"] == 0 - assert max(handle.lengths) <= 65536 - if big: - assert len(handle.lengths) >= 1 + len(stage) - - -def test_reward_env_run_judge_stages_large_tests(tmp_path): - import os as _os - - from yeto.rl.harness.reward_env import benchmark as bm - from yeto.rl.harness.reward_env import tb2 as rt - - tasks_dir = _make_task(tmp_path / "tasks") - (tasks_dir / "fix-git" / "tests" / "blob.bin").write_bytes(_os.urandom(120_000)) - adapter = rt.Tb2Benchmark(tasks_dir, version="tb2@test") - assert adapter.judge_setup_commands("fix-git") - task = tb2_provider.resolve_task("fix-git", tasks_dir) - handle = _ArgvLimited(tb2_provider.LocalProcessBackend(tmp_path / "sb").create(task, "t1")) - handle.exec("echo ok > fixed.txt", timeout_s=30) - result = bm.run_judge(adapter, "fix-git", handle, prebaked=False) - assert result.passed and not result.infra_error and max(handle.lengths) <= 65536 - -def test_modal_sandbox_exec_reads_bytes_and_decodes_leniently(): - calls = {} - - class _Proc: - stdout = SimpleNamespace(read=lambda: b"ok \xf0\x9f") # cut mid UTF-8 sequence - stderr = SimpleNamespace(read=lambda: b"") - - @staticmethod - def wait(): - return 0 - - class _Sandbox: - object_id = "sb" - - def exec(self, *args, **kwargs): - calls["kwargs"] = kwargs - return _Proc() - - result = tb2_provider.ModalSandbox(_Sandbox(), "/app").exec("echo", timeout_s=5) - assert calls["kwargs"]["text"] is False - assert result.output.startswith("ok ") and result.exit_code == 0 diff --git a/tests/test_import_boundaries.py b/tests/test_import_boundaries.py index 7994a6840..2a3603f8c 100644 --- a/tests/test_import_boundaries.py +++ b/tests/test_import_boundaries.py @@ -30,7 +30,7 @@ ALLOWLIST = REPO / "tests" / "import_boundary_allowlist.txt" README = "yeto/rl/engine/README.md" # Upper bound of allowlist entries; may only go down (D4). -ALLOWLIST_CAP = 8 # rl-code-quality-cleanup 1.10: legacy engine and launcher->overlay gone +ALLOWLIST_CAP = 5 # S26 E0.5: Codex harness removed; rl-code-quality-cleanup 1.10: legacy engine and launcher->overlay gone FRAMEWORKS = ("miles", "miles_plugins", "megatron", "sglang", "verl", "vllm", "torch_npu") ADAPTERS = ("yeto.rl.adapters",) diff --git a/tests/test_qwen3_cyber_profile.py b/tests/test_qwen3_cyber_profile.py index 1fa97fd7f..becdff303 100644 --- a/tests/test_qwen3_cyber_profile.py +++ b/tests/test_qwen3_cyber_profile.py @@ -8,11 +8,6 @@ from __future__ import annotations -import os -import subprocess -import sys -from pathlib import Path - import pytest from yeto.rl import codex_backend as cb @@ -25,8 +20,6 @@ stock_codex_reasoning_markers, validate_stock_codex_fields, ) -from yeto.rl.harness.codex import codex_openenv_agent_function as adapter -from yeto.rl.harness.codex import pins PROFILE = "qwen3_30b_a3b" @@ -37,17 +30,7 @@ } -def test_profile_is_allowlisted_and_resolvable(): - assert PROFILE in pins.OPENENV_BACKEND_PROFILES - assert adapter.profile_identity(PROFILE) == ( - QWEN3_30B_A3B_MODEL, - QWEN3_30B_A3B_REVISION, - ) - assert adapter.resolve_backend_profile(PROFILE, env={}) == PROFILE - assert ( - adapter.resolve_backend_profile(env={adapter.RUNTIME_PROFILE_ENV: PROFILE}) - == PROFILE - ) +def test_profile_chat_template_is_its_name(): # every runtime profile's chat template equals its profile name assert stock_codex_backend_contract(PROFILE, 1)["chat_template"] == PROFILE @@ -99,36 +82,6 @@ def test_validate_stock_codex_fields_accepts_then_rejects_drift(): _validate_qwen3(rollout_model="Qwen/Qwen3-32B") -def _harness_tool_output_bytes(env_extra: dict[str, str]) -> int: - """Resolve MAX_TOOL_OUTPUT_BYTES in a child interpreter (it is bound at - import time from the profile named by YETO_CODEX_CHAT_TEMPLATE).""" - repo = Path(__file__).resolve().parents[1] - env = dict(os.environ) - env.pop("YETO_CODEX_CHAT_TEMPLATE", None) - env.update(env_extra) - env["PYTHONPATH"] = str(repo) - out = subprocess.run( - [ - sys.executable, - "-c", - "from yeto.rl.harness.codex import codex_harness_agent as h;" - " print(h.MAX_TOOL_OUTPUT_BYTES)", - ], - env=env, - capture_output=True, - text=True, - check=True, - ) - return int(out.stdout.strip()) - - -def test_tool_output_budget_is_profile_driven(): - assert _harness_tool_output_bytes({"YETO_CODEX_CHAT_TEMPLATE": PROFILE}) == 8_192 - # profiles without a declaration keep the signed 32 KiB default - assert _harness_tool_output_bytes({"YETO_CODEX_CHAT_TEMPLATE": "qwen35_08b"}) == 32_768 - assert _harness_tool_output_bytes({}) == 32_768 - - def test_existing_profiles_are_unchanged(): for name in ("qwen35_08b", "qwen35", "qwen38_next", "qwen38_next_4layer"): profile = cb.stock_codex_backend_profile(name) diff --git a/tests/test_reasoning_loss_mask.py b/tests/test_reasoning_loss_mask.py index 9069d99c1..a2ae242f8 100644 --- a/tests/test_reasoning_loss_mask.py +++ b/tests/test_reasoning_loss_mask.py @@ -20,7 +20,6 @@ from yeto.rl import codex_backend as cb from yeto.rl.harness import reasoning_loss_mask as rlm -from yeto.rl.harness.codex.alignment import assert_sample_alignment OPEN, CLOSE = 900, 901 IM_START, ASSISTANT, NL, IM_END, TOOL = 10, 11, 12, 13, 14 @@ -113,7 +112,6 @@ def test_masks_reasoning_keeps_close_answers_and_tool(): assert changed == 3 assert s.metadata == {"k": "v", "reasoning_in_loss": False, "reasoning_tokens_masked": 3} assert len(s.rollout_log_probs) == s.response_length # nothing else changed - assert_sample_alignment(s) def test_generated_open_marker_is_masked_and_truncation_masks_to_end(): @@ -216,25 +214,7 @@ def test_miles_calls_hook_before_returning_train_samples(rel): assert "_convert_samples_to_train_data" in executor or "convert_samples_to_train_data" in executor -# --- flag plumbing: CLI -> launcher -> learner argv; default argv unchanged ---- -from test_rl_launcher_codex_bundle import bundle # noqa: E402,F401 (pytest fixture) - -def test_launcher_forwards_flag_only_when_set(bundle): - import test_rl_launcher_codex_bundle as tb # noqa: PLC0415 - from yeto.gpu_spec import parse_gpu_spec - from yeto.launcher import _prepare_rl_args, make_miles_island_task - - args = tb._codex_args() - spec = parse_gpu_spec(args.gpu)[0] - default_run = make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400").run - assert "--codex-exclude-reasoning-from-loss" not in default_run - - args.codex_exclude_reasoning_from_loss = True - _prepare_rl_args(args) - run = make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400").run - assert " --codex-exclude-reasoning-from-loss" in run - assert run.replace(" --codex-exclude-reasoning-from-loss", "") == default_run - +# --- flag plumbing: CLI -> launcher ---------------------------------------------- def test_launcher_rejects_flag_without_signed_codex_agent(): from test_rl_engine_selection import _cli # noqa: PLC0415 diff --git a/tests/test_reward_env.py b/tests/test_reward_env.py index 345b89d01..29fe7b07f 100644 --- a/tests/test_reward_env.py +++ b/tests/test_reward_env.py @@ -1,367 +1,19 @@ -"""CPU tests for rl-agentic-reward-env (benchmark-neutral reward sandbox, TB2 first).""" +"""CPU tests for the benchmark-neutral reward-env helpers (rl-agentic-reward-env). -from __future__ import annotations +The TB2 adapter and its prebake / judge tests left with the self-written Codex +harness (S26 E0.5). +""" -import json -import sys -import types -from pathlib import Path -from types import SimpleNamespace +from __future__ import annotations import pytest from yeto.rl.harness.reward_env import benchmark as bm -from yeto.rl.harness.reward_env import tb2 as rt - -UVX_TEST_SH = """#!/bin/bash -# Install curl -apt-get update -apt-get install -y curl - -curl -LsSf https://astral.sh/uv/0.9.5/install.sh | sh -source $HOME/.local/bin/env -if [ "$PWD" = "/" ]; then - echo "Error: No working directory set." - exit 1 -fi -uvx \\ - -p 3.13 \\ - -w pytest==8.4.1 \\ - -w pytest-json-ctrf==0.3.5 \\ - pytest --ctrf /logs/verifier/ctrf.json /tests/test_outputs.py -rA -if [ $? -eq 0 ]; then - echo 1 > /logs/verifier/reward.txt -else - echo 0 > /logs/verifier/reward.txt -fi -""" - - -def make_task(root: Path, task_id: str, test_sh: str = UVX_TEST_SH, image: str = "org/img:1", - difficulty: str = "medium") -> Path: - d = root / task_id - (d / "tests").mkdir(parents=True) - (d / "environment").mkdir() - (d / "task.toml").write_text( - f'[metadata]\ndifficulty = "{difficulty}"\n' - f'[environment]\ndocker_image = "{image}"\ncpus = 2\nmemory_mb = 4096\n' - "[agent]\ntimeout_sec = 600\n[verifier]\ntimeout_sec = 300\n" - ) - (d / "tests" / "test.sh").write_text(test_sh) - (d / "environment" / "Dockerfile").write_text("FROM x\nWORKDIR /app\n") - (d / "instruction.md").write_text(f"do {task_id}\n") - return d - - -# --- prebake plan extraction --------------------------------------------------- - -def test_uvx_plan_is_setup_plus_warmup(): - plan = rt.prebake_from_test_sh(UVX_TEST_SH, "img") - assert plan.complete - assert plan.commands == ( - "apt-get update", - "apt-get install -y curl", - "curl -LsSf https://astral.sh/uv/0.9.5/install.sh | sh", - "source $HOME/.local/bin/env", - "uvx -p 3.13 -w pytest==8.4.1 -w pytest-json-ctrf==0.3.5 pytest --version", - ) - script = rt.prebake_script(plan) - assert script.startswith("set -e\n") and script.rstrip().endswith("pytest --version") - - -def test_pip_plan(): - sh = 'pip install pytest==8.4.1 pytest-json-ctrf==0.3.5 --break-system-packages\npython -m pytest /tests -rA\n' - plan = rt.prebake_from_test_sh(sh, "img") - assert plan.complete and plan.commands == ("pip install pytest==8.4.1 pytest-json-ctrf==0.3.5 --break-system-packages",) - - -def test_uv_run_is_test_line_and_venv_setup_kept(): - sh = "uv venv -p 3.12 .tb\nsource .tb/bin/activate\nuv pip install pytest==8.4.1\nuv run pytest /tests -rA\n" - plan = rt.prebake_from_test_sh(sh, "img") - assert plan.complete and plan.commands[-1] == "uv pip install pytest==8.4.1" - - -def test_task_specific_lines_stay_at_judge_time_and_vars_block_warmup(): - sh = ("apt-get update\nX='abc'\ncurl -L -o w.pt https://h/w.pt\ncp /tests/a .\n" - 'uvx -p 3.11 -w "git+https://g@${X}" pytest /tests\n') - plan = rt.prebake_from_test_sh(sh, "img") - assert plan.commands == ("apt-get update",) - assert not plan.complete - assert "curl -L -o w.pt https://h/w.pt" in plan.skipped and "cp /tests/a ." in plan.skipped - assert any(s.startswith("uvx ") for s in plan.skipped) - - -def test_uvx_value_flags(): - line = "uvx --index https://x/cpu --index-strategy unsafe-best-match -p 3.11 -w torch==2.5.1 pytest -q" - assert rt._uvx_warmup(line).endswith("-w torch==2.5.1 pytest --version") - - -def test_no_test_line_means_no_guessing(): - plan = rt.prebake_from_test_sh("apt-get update\nmake test\n", "img") - assert plan.commands == () and not plan.complete -def test_digest_is_content_hash(): - a = rt.prebake_from_test_sh(UVX_TEST_SH, "img") - assert a.digest() == rt.prebake_from_test_sh(UVX_TEST_SH, "img").digest() - assert a.digest() != rt.prebake_from_test_sh(UVX_TEST_SH, "img2").digest() - assert a.digest() != rt.prebake_from_test_sh(UVX_TEST_SH.replace("8.4.1", "8.4.2"), "img").digest() - - -# --- registry + adapter ---------------------------------------------------------- - -def test_registry_lazy_import_and_errors(tmp_path): - make_task(tmp_path, "t1") - adapter = bm.get_adapter("tb2", tasks_dir=tmp_path) - assert adapter.name == "tb2" and "tb2" in bm.registered() +def test_unknown_benchmark_is_refused(): with pytest.raises(KeyError, match="unknown benchmark"): bm.get_adapter("no_such_bench") - with pytest.raises(ValueError, match="already registered"): - bm.register("tb2", lambda **_: None) - - -def test_adapter_spec_and_plans(tmp_path): - make_task(tmp_path, "t1") - make_task(tmp_path, "t2", image="org/other:2") - adapter = rt.Tb2Benchmark(tmp_path) - assert adapter.task_ids() == ["t1", "t2"] - spec = adapter.task_spec("t2") - assert (spec.base_image, spec.cpus, spec.memory_mb, spec.workdir, spec.judge_timeout_s) == ("org/other:2", 2, 4096, "/app", 300.0) - assert spec.instruction.startswith("do t2") - assert set(bm.plans(adapter)) == {"t1", "t2"} - assert "test.sh" in adapter.judge_command("t1") - - -# --- judge runner ---------------------------------------------------------------- - -class FakeHandle: - def __init__(self, output, exit_code=0, timed_out=False): - self.result = SimpleNamespace(output=output, exit_code=exit_code, timed_out=timed_out) - self.commands = [] - - def exec(self, command, *, timeout_s, workdir=None): - self.commands.append((command, timeout_s)) - return self.result - - -@pytest.mark.parametrize("reward,passed", [("1", True), ("0", False), ("", False)]) -def test_run_judge(tmp_path, reward, passed): - make_task(tmp_path, "t1") - adapter = rt.Tb2Benchmark(tmp_path) - handle = FakeHandle(f"pytest ok\nYETO_TB2_TESTSH_RC=0\nYETO_TB2_REWARD={reward}\n") - ticks = iter([10.0, 12.5]) - result = bm.run_judge(adapter, "t1", handle, prebaked=True, clock=lambda: next(ticks)) - assert (result.passed, result.reward, result.seconds, result.prebaked) == (passed, float(passed), 2.5, True) - assert result.extra["testsh_rc"] == 0 and handle.commands[0][1] == 300.0 - json.dumps(result.to_dict()) - - -# --- Modal backend (fake modal module) ------------------------------------------- - -class _FakeImage: - calls: list = [] - - @classmethod - def from_registry(cls, ref): - img = cls() - img.ops = [("from_registry", ref)] - return img - - def run_commands(self, *cmds): - self.ops.append(("run_commands", cmds)) - return self - - -@pytest.fixture -def fake_modal(monkeypatch): - created = {} - - class Sandbox: - @staticmethod - def create(*args, **kwargs): - created.update(kwargs, args=args) - return SimpleNamespace(object_id="sb-1") - - class App: - @staticmethod - def lookup(name, create_if_missing=False): - created["app_name"] = name - return SimpleNamespace(name=name) - - monkeypatch.setitem(sys.modules, "modal", types.SimpleNamespace(Image=_FakeImage, Sandbox=Sandbox, App=App)) - return created - - -@pytest.mark.parametrize("prebake", [True, False]) -def test_prebaked_backend(tmp_path, fake_modal, prebake): - from yeto.rl.harness.codex import tb2_provider as tb2 - - make_task(tmp_path, "t1") - task = tb2.resolve_task("t1", tmp_path) - from yeto.cloud import modal_reward_env - - backend = modal_reward_env.PrebakedModalSandboxBackend(prebake=prebake) - handle = backend.create(task, "traj-1") - assert fake_modal["app_name"] == "yeto-reward-env" and fake_modal["app"].name == "yeto-reward-env" - assert handle.workdir == "/app" - ops = fake_modal["image"].ops - assert ops[0] == ("from_registry", "org/img:1") - if prebake: - assert ops[1][0] == "run_commands" and len(ops[1][1]) == 1 - run = ops[1][1][0] - assert "\n" not in run and run.startswith("echo ") # one Dockerfile RUN line (S17 G3) - import base64 - - script = base64.b64decode(run.split()[1]).decode() - assert script == rt.prebake_script(rt.prebake_from_test_sh(UVX_TEST_SH, "org/img:1")) - assert "pytest --version" in script - assert fake_modal["tags"]["yeto-prebake"] == rt.prebake_from_test_sh(UVX_TEST_SH, "org/img:1").digest()[:16] - else: - assert len(ops) == 1 and fake_modal["tags"]["yeto-prebake"] == "none" - assert fake_modal["cpu"] == 2.0 and fake_modal["memory"] == 4096 - - -def test_build_never_targets_production_app(): - with pytest.raises(ValueError, match="production"): - rt.check_build_app("yeto-tbench2") - assert rt.check_build_app("yeto-reward-env-build") == "yeto-reward-env-build" - - -def test_build_tool_plan_only(tmp_path, capsys): - sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools" / "reward_env")) - import build_images - - make_task(tmp_path, "t1") - out = tmp_path / "plan.json" - assert build_images.main(["--tasks-dir", str(tmp_path), "--out", str(out)]) == 0 - report = json.loads(out.read_text()) - assert report["total"] == 1 and report["tasks"][0]["complete"] - with pytest.raises(SystemExit): - build_images.main(["--tasks-dir", str(tmp_path), "--build"]) # no --app - - -# --- hold-out list (WP3 D6b) + metadata (D6a) ------------------------------------ - -def _tb2_pool(root): - for i in range(4): - make_task(root, f"e{i}", difficulty="easy") - for i in range(20): - make_task(root, f"m{i:02d}", difficulty="medium") - for i in range(12): - make_task(root, f"h{i:02d}", difficulty="hard") - return rt.Tb2Benchmark(root, version="tb2@test") - - -def test_tb2_metadata_shape(tmp_path): - make_task(tmp_path, "t1", difficulty="hard") - meta = rt.Tb2Benchmark(tmp_path, version="tb2@abc1234").task_spec("t1").metadata() - assert meta == {"task_id": "t1", "benchmark": "tb2", "benchmark_version": "tb2@abc1234", - "difficulty": "hard", "difficulty_source": "tb2-task.toml", "eval_bucket": "tb2-hard"} - - -def test_tb2_holdout_stratified_and_reproducible(tmp_path): - adapter = _tb2_pool(tmp_path) - h = rt.build_holdout(adapter) - assert h["schema"] == "yeto-eval-holdout/1" and h["seed"] == 20261008 and h["benchmark_version"] == "tb2@test" - from collections import Counter - - assert Counter(i["eval_bucket"] for i in h["items"]) == {"tb2-easy": 2, "tb2-medium": 18, "tb2-hard": 10} - assert rt.build_holdout(adapter, task_ids=list(reversed(adapter.task_ids()))) == h - assert bm.holdout_sha256(h) == bm.holdout_sha256(rt.build_holdout(adapter)) - ids = bm.holdout_ids(h) - excluded = rt.build_holdout(adapter, exclude=ids[:1] if ids[0].startswith("m") else ids[2:3]) - assert len(excluded["items"]) == 30 - with pytest.raises(ValueError, match="quota"): - rt.build_holdout(adapter, exclude=["e0", "e1", "e2"]) - - -def test_tb2_holdout_excludes_smoke6_by_default(tmp_path): - adapter = _tb2_pool(tmp_path) - for t in rt.SMOKE6_TASK_IDS: - make_task(tmp_path, t, difficulty="medium") - adapter = rt.Tb2Benchmark(tmp_path, version="tb2@test") - h = rt.build_holdout(adapter) - ids = set(bm.holdout_ids(h)) - assert not ids & set(rt.SMOKE6_TASK_IDS) and len(ids) == 30 - assert h["excluded"] == [{"task_id": t, "reason": rt.SMOKE6_REASON} for t in sorted(rt.SMOKE6_TASK_IDS)] - assert "smoke6" in h["rule"] - none = rt.build_holdout(adapter, exclude={}) - assert "excluded" not in none and len(none["items"]) == 30 - - -def test_holdout_cli_exclusions(tmp_path): - sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools" / "reward_env")) - import holdout - - tasks = tmp_path / "tasks" - tasks.mkdir() - _tb2_pool(tasks) - out = tmp_path / "out" - # built-in smoke-6 list, but the checkout lacks them -> fail closed - with pytest.raises(SystemExit, match="not in the TB2 checkout"): - holdout.main(["--tb2-tasks-dir", str(tasks), "--out-dir", str(out)]) - assert holdout.main(["--tb2-tasks-dir", str(tasks), "--tb2-no-exclude", "--out-dir", str(out)]) == 0 - assert "excluded" not in json.loads((out / "tb2-holdout.json").read_text()) - jl = tmp_path / "ex.jsonl" - jl.write_text(json.dumps({"metadata": {"task_id": "m00"}}) + "\n\n") - assert holdout.main(["--tb2-tasks-dir", str(tasks), "--tb2-exclude-jsonl", str(jl), "--tb2-no-unusable", - "--tb2-exclude-reason", "r", "--out-dir", str(out)]) == 0 - h = json.loads((out / "tb2-holdout.json").read_text()) - assert h["excluded"] == [{"task_id": "m00", "reason": "r"}] and "m00" not in bm.holdout_ids(h) - jl.write_text(json.dumps({"prompt": "x"}) + "\n") - with pytest.raises(SystemExit, match="metadata.task_id"): - holdout.main(["--tb2-tasks-dir", str(tasks), "--tb2-exclude-jsonl", str(jl), "--out-dir", str(out)]) - - -def test_tb2_unusable_table_shipped(): - table = rt.load_unusable() - assert len(table) == 13 and "qemu-startup" in table - assert all(r.startswith("unusable/") for r in table.values()) - assert not set(table) & set(rt.SMOKE6_TASK_IDS) - - -def test_tb2_unusable_bad_file(tmp_path): - bad = tmp_path / "u.json" - bad.write_text(json.dumps({"schema": "nope", "benchmark": "tb2", "items": []})) - with pytest.raises(ValueError, match="not a"): - rt.load_unusable(bad) - bad.write_text(json.dumps({"schema": rt.UNUSABLE_SCHEMA, "benchmark": "tb2", "items": [ - {"task_id": "a", "category": "c", "reason": "r"}, {"task_id": "a", "category": "c", "reason": "r"}]})) - with pytest.raises(ValueError, match="duplicate"): - rt.load_unusable(bad) - - -def test_holdout_cli_unusable_leaves_train_and_eval(tmp_path): - sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools" / "reward_env")) - import holdout - - tasks = tmp_path / "tasks" - tasks.mkdir() - _tb2_pool(tasks) - for t in rt.SMOKE6_TASK_IDS: - make_task(tasks, t, difficulty="medium") - table = tmp_path / "u.json" - table.write_text(json.dumps({"schema": rt.UNUSABLE_SCHEMA, "benchmark": "tb2", "items": [ - {"task_id": "m01", "category": "timeout", "reason": "slow"}, - {"task_id": "h03", "category": "env", "reason": "404"}]})) - out = tmp_path / "out" - assert holdout.main(["--tb2-tasks-dir", str(tasks), "--tb2-unusable", str(table), "--out-dir", str(out)]) == 0 - h = json.loads((out / "tb2-holdout.json").read_text()) - tr = json.loads((out / "tb2-train.json").read_text()) - held, train = set(bm.holdout_ids(h)), {i["task_id"] for i in tr["items"]} - assert len(held) == 30 and h["rule"] == rt.HOLDOUT_RULE_UNUSABLE - assert {"m01", "h03"} <= {e["task_id"] for e in h["excluded"]} - assert not {"m01", "h03"} & (held | train) and not held & train - assert set(rt.SMOKE6_TASK_IDS) <= train # smoke-6 only leaves the eval pool - assert len(held) + len(train) + 2 == 36 + 6 - assert tr["holdout_sha256"] == bm.holdout_sha256(h) and tr["schema"] == bm.TRAIN_SPLIT_SCHEMA - assert tr["excluded"] == [{"task_id": "h03", "reason": "unusable/env: 404"}, - {"task_id": "m01", "reason": "unusable/timeout: slow"}] - # an unusable task missing from the checkout fails closed - table.write_text(json.dumps({"schema": rt.UNUSABLE_SCHEMA, "benchmark": "tb2", "items": [ - {"task_id": "zz", "category": "c", "reason": "r"}]})) - with pytest.raises(SystemExit, match="unusable tasks not in the TB2 checkout"): - holdout.main(["--tb2-tasks-dir", str(tasks), "--tb2-unusable", str(table), "--out-dir", str(out)]) def test_rows_split_and_leak_check(): @@ -375,23 +27,3 @@ def test_rows_split_and_leak_check(): bm.assert_disjoint(["a"], ["a"]) with pytest.raises(ValueError, match="duplicate"): bm.holdout_ids({"schema": bm.HOLDOUT_SCHEMA, "items": [{"task_id": "a"}, {"task_id": "a"}]}) - - -def test_run_judge_infra_error(tmp_path): - make_task(tmp_path, "t1") - - class Gone: - def exec(self, command, *, timeout_s, workdir=None): - raise RuntimeError("sandbox terminated") - - result = bm.run_judge(rt.Tb2Benchmark(tmp_path), "t1", Gone(), prebaked=False) - assert result.infra_error and not result.passed and "terminated" in result.log - - -# --- launcher recognises the new provider ----------------------------------------- - -def test_launcher_treats_prebaked_provider_as_modal(): - from yeto import launcher as L - - assert L.PREBAKED_MODAL_SANDBOX_PROVIDER in L.MODAL_SANDBOX_PROVIDERS - assert L.MODAL_SANDBOX_PROVIDER in L.MODAL_SANDBOX_PROVIDERS diff --git a/tests/test_rl_agentic_suspend.py b/tests/test_rl_agentic_suspend.py deleted file mode 100644 index 4dd1c6712..000000000 --- a/tests/test_rl_agentic_suspend.py +++ /dev/null @@ -1,491 +0,0 @@ -"""agentic-rollout-utilization stage 3 (5.1-5.3) and the Miles in-flight cut wiring (3.3). - -CPU only: fake session server (aiohttp), fake Codex client, fake executor. -No Ray, no Miles process, no GPU. -""" - -from __future__ import annotations - -import asyncio -import json -import os -from dataclasses import asdict -from types import SimpleNamespace -from typing import Any - -import aiohttp -import pytest -from aiohttp import web - -from test_codex_bridge_compaction import FakeCodex, _exec, _submit -from test_rl_miles_carry_over import Call, FakePartialEngine, Sample, Span, tok -from test_secrlenv_codex_harness import _fake_miles_sse -from yeto.rl.adapters.miles import carry_over, in_flight -from yeto.rl.adapters.miles import policy_age as miles_policy_age -from yeto.rl.engine.cut import context_problems -from yeto.rl.engine.policy_age import PolicyAgeError -from yeto.rl.engine.trajectory_timing import PhaseClock -from yeto.rl.harness.codex import codex_harness_agent as harness -from yeto.rl.harness.codex import codex_openenv_generate as wrapper -from yeto.rl.harness.codex import codex_openenv_subprocess_agent_function as trusted - -AGENT = "yeto.rl.harness.codex.codex_openenv_subprocess_agent_function.run" - - -# ------------------------------------------------------------ 5.1 the gate - - -class GatedSessionServer: - """Fake Miles session server; ``on_request(n, body)`` may return an HTTP status - to answer instead of the next completion (503 = engine aborted the turn).""" - - def __init__(self, replies, on_request=None): - self.replies = list(replies) - self.bodies: list[dict[str, Any]] = [] - self.on_request = on_request - - async def __aenter__(self): - async def chat(request: web.Request) -> web.Response: - body = await request.json() - self.bodies.append(body) - status = self.on_request(len(self.bodies), body) if self.on_request else None - if status is not None: - return web.json_response({"error": "upstream generation aborted"}, status=status) - return web.Response(body=_fake_miles_sse(self.replies.pop(0)), content_type="text/event-stream") - - app = web.Application() - app.router.add_post("/sessions/{sid}/v1/chat/completions", chat) - self.runner = web.AppRunner(app, access_log=None) - await self.runner.setup() - site = web.TCPSite(self.runner, "127.0.0.1", 0) - await site.start() - self.url = f"http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}/sessions/S0" - return self - - async def __aexit__(self, *_): - await self.runner.cleanup() - - -def _gate_env(monkeypatch, gate, max_s="30"): - monkeypatch.setenv("YETO_CODEX_BACKEND_MAX_TOKENS", "512") - monkeypatch.delenv("YETO_CODEX_COMPACTION_ENABLED", raising=False) - monkeypatch.setenv(harness.SUSPEND_GATE_ENV, str(gate)) - monkeypatch.setenv(harness.SUSPEND_MAX_SECONDS_ENV, max_s) - monkeypatch.setattr(harness, "SUSPEND_POLL_SECONDS", 0.02) - - -def _bridge(url, metrics): - return harness._ResponsesBridge(url, "solve target", {"max_tokens": 512}, metrics, max_seq_len=8000) - - -def test_cutoff_during_a_tool_call_parks_the_next_model_turn_until_resume(monkeypatch, tmp_path): - gate = tmp_path / "t.gate" - _gate_env(monkeypatch, gate) - edges: list[str] = [] - harness.set_suspend_event_sink(edges.append) - - async def scenario(): - metrics = harness.legacy.AgentMetrics() - async with GatedSessionServer([_exec(1, 500), _submit(900)]) as server: - async with _bridge(server.url, metrics) as bridge, aiohttp.ClientSession() as http: - codex = FakeCodex(bridge, http) - assert (await codex.turn())[0] == 200 # model turn 1 -> tool call - gate.write_text("suspended\n") # cut-off while the tool runs - # the tool returns and its result goes into the next request; the - # bridge holds that request (no model call while suspended) - pending = asyncio.ensure_future(codex.turn('{"output":"tool result"}')) - await asyncio.sleep(0.3) - assert not pending.done() and len(server.bodies) == 1 - gate.unlink() # resume - status, resp = await pending - assert status == 200 and resp["output"][1]["name"] == "submit" - return server, metrics - - try: - server, metrics = asyncio.run(scenario()) - finally: - harness.set_suspend_event_sink(None) - # the request after resume carries the tool result written before the suspension - assert any("tool result" in json.dumps(m) for m in server.bodies[1]["messages"]) - assert metrics.suspensions == 1 and metrics.suspended_seconds >= 0.25 - assert metrics.suspend_retried_turns == 0 and metrics.suspend_expired == 0 - assert edges == ["enter", "exit"] - # parked time is not generation time - assert metrics.total_generation_time < metrics.suspended_seconds - - -def test_a_model_turn_aborted_by_the_cutoff_is_redone_after_resume(monkeypatch, tmp_path): - gate = tmp_path / "t.gate" - _gate_env(monkeypatch, gate) - - def on_request(n, _body): - if n == 2: # the cut-off: gate closed first, then the engine aborts this turn - gate.write_text("suspended\n") - asyncio.get_event_loop().call_later(0.2, gate.unlink) - return 503 - return None - - async def scenario(): - metrics = harness.legacy.AgentMetrics() - async with GatedSessionServer([_exec(1, 500), _submit(900)], on_request) as server: - async with _bridge(server.url, metrics) as bridge, aiohttp.ClientSession() as http: - codex = FakeCodex(bridge, http) - assert (await codex.turn())[0] == 200 - status, resp = await codex.turn('{"output":"r"}') - assert status == 200 and resp["output"][1]["name"] == "submit" - return server, metrics - - server, metrics = asyncio.run(scenario()) - assert len(server.bodies) == 3 and server.bodies[1] == server.bodies[2] # same turn redone - assert metrics.suspend_retried_turns == 1 and metrics.suspensions == 1 - - -def test_a_503_without_a_suspension_is_still_an_error(monkeypatch, tmp_path): - _gate_env(monkeypatch, tmp_path / "never.gate") - - async def scenario(): - metrics = harness.legacy.AgentMetrics() - async with GatedSessionServer([_exec(1, 500)], lambda n, _b: 503 if n == 2 else None) as server: - async with _bridge(server.url, metrics) as bridge, aiohttp.ClientSession() as http: - codex = FakeCodex(bridge, http) - assert (await codex.turn())[0] == 200 - status, _ = await codex.turn('{"output":"r"}') - return status, bridge, metrics - - status, bridge, metrics = asyncio.run(scenario()) - assert status == 400 and "HTTP 503" in str(bridge.fatal.result()) - assert metrics.suspensions == 0 - - -def test_a_suspension_past_the_survival_limit_discards_the_trajectory(monkeypatch, tmp_path): - gate = tmp_path / "t.gate" - _gate_env(monkeypatch, gate, max_s="0.2") - - async def scenario(): - metrics = harness.legacy.AgentMetrics() - async with GatedSessionServer([_exec(1, 500), _submit(900)]) as server: - async with _bridge(server.url, metrics) as bridge, aiohttp.ClientSession() as http: - codex = FakeCodex(bridge, http) - assert (await codex.turn())[0] == 200 - gate.write_text("suspended\n") - status, _ = await codex.turn('{"output":"r"}') - return status, bridge, metrics, server - - status, bridge, metrics, server = asyncio.run(scenario()) - assert status == 400 and isinstance(bridge.fatal.result(), harness.CodexSuspendExpired) - assert metrics.suspend_expired == 1 and len(server.bodies) == 1 # never sampled again - - -def test_codex_waits_for_a_held_request_only_when_suspension_is_on(monkeypatch, tmp_path): - monkeypatch.delenv(harness.SUSPEND_GATE_ENV, raising=False) - assert harness._suspend_settings() == [] # default command line unchanged - monkeypatch.setenv(harness.SUSPEND_GATE_ENV, str(tmp_path / "g")) - monkeypatch.setenv(harness.SUSPEND_MAX_SECONDS_ENV, "600") - assert harness._suspend_settings() == ["model_providers.miles.stream_idle_timeout_ms=1200000"] - - -def test_phase_clock_keeps_parked_time_out_of_generation(): - now = [0.0] - clock = PhaseClock(lambda: now[0]) - now[0] = 2.0 - clock.enter_tool() - now[0] = 3.0 - clock.exit_tool() - now[0] = 4.0 - clock.enter_suspend() # parked before model turn 2 - now[0] = 104.0 - clock.exit_suspend() - now[0] = 106.0 - out = clock.finish() - assert out["turn_generation_seconds"] == [2.0, 3.0] and out["turn_tool_seconds"] == [1.0] - assert out["suspended_seconds"] == 100.0 and out["worker_seconds"] == 106.0 - assert "suspended_seconds" not in PhaseClock(lambda: 0.0).finish() # unchanged when never parked - - -# ------------------------------------------------- 5.1/5.3 trusted-side hooks - - -def test_suspend_closes_the_gates_of_trajectories_in_flight_and_resume_opens_them(monkeypatch, tmp_path): - monkeypatch.setattr(trusted, "_SUSPEND_DIR", None) - monkeypatch.setattr(trusted, "_SUSPEND_STATE", {"suspended_at": None, "gates": {}}) - monkeypatch.setattr(trusted, "_INFLIGHT", {"a": {"stage": "worker"}, "b": {"stage": "verify"}, - "c": {"stage": "teardown"}}) - assert trusted._suspend_worker_env("x") == {} # no suspension in use: env unchanged - asyncio.run(trusted.resume()) # Miles calls resume at the start of every rollout - assert trusted._suspend_worker_env("x") == {harness.SUSPEND_GATE_ENV: trusted._gate_path("x")} - assert asyncio.run(trusted.suspend()) == 2 # teardown is not suspended - gates = trusted._SUSPEND_STATE["gates"] - assert set(gates) == {"a", "b"} and all(os.path.exists(p) for p in gates.values()) - # a trajectory started while suspended is parked too - env = trusted._suspend_worker_env("late") - assert os.path.exists(env[harness.SUSPEND_GATE_ENV]) - paths = list(trusted._SUSPEND_STATE["gates"].values()) - assert asyncio.run(trusted.resume()) == 3 - assert not any(os.path.exists(p) for p in paths) and trusted._SUSPEND_STATE["gates"] == {} - - -# ------------------------------------------------------- 5.2 version segments - - -def _traj(versions): - calls = [{"spans": [{"version": v}]} for v in versions] - return SimpleNamespace(metadata={}, weight_versions=calls, status="completed", loss_mask=None) - - -def test_codex_versions_are_recorded_as_segments_within_the_window(monkeypatch): - monkeypatch.setattr(wrapper, "assert_sample_alignment", lambda s: None) - ok = _traj([tok(3), tok(3), tok(4)]) - wrapper.apply_trajectory_bookkeeping(SimpleNamespace(group_index=1, index=9), [ok], - expected_version=tok(3), max_policy_age=1, current_version=tok(4)) - assert ok.status == "completed" and wrapper.POLICY_AGE_KEY not in ok.metadata - assert ok.metadata[wrapper.VERSION_SEGMENTS_KEY] == sorted({tok(3), tok(4)}) - old = _traj([tok(2), tok(4)]) - wrapper.apply_trajectory_bookkeeping(SimpleNamespace(group_index=1, index=9), [old], - expected_version=tok(2), max_policy_age=1, current_version=tok(4)) - reason = old.metadata[wrapper.INFRASTRUCTURE_KEY] - assert reason.startswith(wrapper.POLICY_AGE_EXCEEDED) and wrapper.POLICY_AGE_KEY not in old.metadata - # limit 0: drift is still a policy-identity violation (unchanged) - drift = _traj([tok(3), tok(4)]) - wrapper.apply_trajectory_bookkeeping(SimpleNamespace(group_index=1, index=9), [drift], - expected_version=tok(3)) - assert drift.metadata[wrapper.POLICY_AGE_KEY] == 1 - - -def test_window_problem_rules(): - assert wrapper.window_problem([tok(3)], tok(4), 1) is None - assert "newer" in wrapper.window_problem([tok(5)], tok(4), 1) - assert "unparsable" in wrapper.window_problem(["default"], tok(4), 1) - assert "no generation" in wrapper.window_problem([], tok(4), 1) - - -# ----------------------------------------------------- policy age (launcher) - - -def test_miles_stage_three_accepts_the_suspendable_codex_agent_only(): - assert miles_policy_age.SUPPORT.stage == 3 - miles_policy_age.check_task(1, custom_generate=miles_policy_age.AGENTIC_GENERATE, custom_agent=AGENT) - with pytest.raises(PolicyAgeError, match="suspended between turns"): - miles_policy_age.check_task(1, custom_generate=miles_policy_age.AGENTIC_GENERATE, - custom_agent="some.other.agent.run") - with pytest.raises(PolicyAgeError, match="stage 3"): - miles_policy_age.check_task(1, custom_generate="miles.rollout.generate_hub.multi_turn.generate") - assert miles_policy_age.policy_age_argv(1, agentic=True) == ( - "--agentic-suspend-between-turns", "--agentic-suspend-max-rounds", "1") - assert miles_policy_age.policy_age_argv(0, agentic=True) == () - - -# ------------------------------------------- carry-over metrics (5.3 / 5.5) - - -def test_agentic_tool_tokens_are_not_cross_version_and_ratios_are_summarised(): - s = Sample(index=0, group_index=0) - FakePartialEngine().generate(s, 3, 2, finish=False) - s.tokens = s.tokens + [9, 9] # tool output (no version span) - s.response_length += 2 - s.rollout_log_probs = s.rollout_log_probs + [0.0, 0.0] - FakePartialEngine().generate(s, 4, 1, finish=True) - s.loss_mask = [1, 1, 0, 0, 1] - assert carry_over.response_token_versions(s) == [3, 3, None, None, 4] - assert carry_over.trained_token_versions(s, 4) == [3, 3, 4, 4, 4] - out = carry_over.estimate_cross_version_truncation( - SimpleNamespace(tis_clip_low=0.0, tis_clip=2.0), [s], 4, lambda _s: [-0.5, 0.5, 0.0, 0.0, -0.5]) - assert out["cross_version_scored_tokens"] == 2 and out["cross_version_unscored_samples"] == 0 - assert out["cross_version_truncated_fraction"] == 0.5 # ratios 1 and e (>2: truncated) - assert out["cross_version_ratio_max"] == pytest.approx(2.718281828, rel=1e-6) - - -def test_suspend_fields_report_the_fork_stats_and_the_survival_cost(): - args = SimpleNamespace(rollout_suspend_stats={"suspended_groups": 3, "over_age_cancelled_groups": 1}, - rollout_resume_stats={"resumed_groups": 2, "resumed_done_groups": 0}) - expired = SimpleNamespace(metadata={"tbench_infrastructure_error": "CodexSuspendExpired: x", - "agent_metrics": {"suspensions": 1, "suspended_seconds": 600.0}}) - parked = SimpleNamespace(metadata={"agent_metrics": {"suspensions": 1, "suspended_seconds": 40.5, - "suspend_retried_turns": 1}}) - plain = SimpleNamespace(metadata={}) - out = carry_over.suspend_fields(args, [[expired, parked], [plain]]) - assert out["suspended_groups"] == 3 and out["resumed_groups"] == 2 - assert out["suspend_expired_trajectories"] == 1 and out["suspended_trajectories"] == 2 - assert out["suspended_env_seconds"] == 640.5 and out["suspend_retried_turns"] == 1 - assert args.rollout_suspend_stats is None # consumed - assert carry_over.suspend_fields(args, []) == {} - - -# ----------------------------------------- 3.3 Miles in-flight cut wiring - - -class MilesLikeSample(Sample): - def to_dict(self): - d = asdict(self) - d["status"] = self.status.value - return d - - -def _buffered_group(gi, version=3): - group = [MilesLikeSample(index=gi * 10 + i, group_index=gi) for i in range(2)] - for s in group: - FakePartialEngine().generate(s, version, 3, finish=False) - s.metadata["task_id"] = f"task-{gi}" - return group - - -def test_export_and_restore_buffered_groups_and_agentic_references(monkeypatch): - executor = SimpleNamespace(data_source=SimpleNamespace(buffer=[_buffered_group(1), _buffered_group(2)])) - monkeypatch.setattr(in_flight, "_codex_suspension", - lambda: ({"t7": "/tmp/t7.gate"}, {"t7": {"stage": "worker"}})) - exported = in_flight.export_in_flight(executor) - assert exported["buffer_groups"] == 2 and len(exported["entries"]) == 5 - assert in_flight.check_export(exported) == [] - agentic = [e for e in exported["entries"] if e.get("session_ref")] - assert agentic[0]["session_ref"] == "codex-suspended:t7" and agentic[0]["stage"] == "worker" - # the cut accepts the buffer because the in-flight section carries it - base = dict(cut_id="r000004-x", progress=SimpleNamespace(problems=lambda: []), - algorithm=SimpleNamespace(algorithm_spec_sha256="a" * 64), - data={"sample_offset": 0, "epoch_id": 0, "sample_group_index": 3, "sample_index": 6, - "buffer_length": 2}, - outer={"settled": True}, runtime={}) - entries = tuple(exported["entries"]) - ok_ledger = {"carried_over": 5, "ready_unconsumed": 0, "max_policy_age": 1} - assert not [p for p in context_problems(**base, ledger=ok_ledger, in_flight=entries, max_policy_age=1) - if p.startswith(("data", "in_flight", "ledger"))] - zero = context_problems(**base, ledger={"carried_over": 0, "ready_unconsumed": 0}) - assert any("not carried by a cut" in p for p in zero) # limit 0: unchanged refusal - short = context_problems(**base, ledger=ok_ledger, in_flight=entries[:2], max_policy_age=1) - assert any("carries 1" in p for p in short) - - # restore at version 4 with limit 1: group 2 already trained, agentic lost - target = SimpleNamespace(data_source=SimpleNamespace(buffer=[], args=None)) - target.data_source.add_samples = lambda groups: target.data_source.buffer.extend(groups) - report = in_flight.import_in_flight(target, exported["entries"], max_policy_age=1, current_version=4, - completed_group_ids=["g2"], sample_from_dict=dict) - assert report["resumed_groups"] == 1 and report["resumed"] == 2 - assert report["resumed_group_ids"] == ["g1"] - assert report["reasons"] == {"group_already_completed": 2, in_flight.AGENTIC_SESSION_LOST: 1} - restored = target.data_source.buffer[0] - assert [m["index"] for m in restored] == [10, 11] and restored[0]["metadata"]["task_id"] == "task-1" - # too old at version 5 (limit 1): discarded - target.data_source.buffer.clear() - report = in_flight.import_in_flight(target, exported["entries"], max_policy_age=1, current_version=5, - sample_from_dict=dict) - assert report["resumed_groups"] == 0 and report["reasons"]["policy_age_exceeded"] == 4 - assert report["resumed_group_ids"] == [] - # restored under limit 0: everything discarded and reported - report = in_flight.import_in_flight(target, exported["entries"], max_policy_age=0, current_version=4, - sample_from_dict=dict) - assert report["discarded"] == 5 and target.data_source.buffer == [] - - -def test_export_refuses_a_buffered_token_of_unknown_version(): - group = _buffered_group(1) - group[0].weight_versions.append(Call([Span("default", 0, len(group[0].tokens))])) - group[0].tokens = group[0].tokens + [1] - group[0].response_length += 1 - group[0].rollout_log_probs.append(-0.1) - group[0].weight_versions = [Call([Span("default", len(group[0].tokens) - 4, len(group[0].tokens))])] - with pytest.raises(in_flight.InFlightExportError, match="unknown"): - in_flight.buffered_entries([group]) - - -def test_cut_source_writes_the_in_flight_section_under_a_limit(): - from yeto.rl.adapters.miles.rebuild_wiring import CutSource - - exported = {"entries": in_flight.buffered_entries([_buffered_group(1)]), "buffer_groups": 1} - state = SimpleNamespace(policy_version=4, policy_tensor_hash=lambda: "h" * 64) - driver = SimpleNamespace(published_state=state, published_version=4, local_step=4, - at_safe_point=True, expected_token=tok(4), - profile=SimpleNamespace(max_policy_age=1)) - rollout = SimpleNamespace(data_cursor=lambda: {"sample_offset": 0, "epoch_id": 0, - "sample_group_index": 2, "sample_index": 4}, - export_in_flight=lambda: exported) - ledger = SimpleNamespace(cut_summary=lambda: {"carried_over": 0, "ready_unconsumed": 0, - "engine_buffer_length": 1}) - source = CutSource.__new__(CutSource) - source.driver_ref, source.rollout, source.ledger = driver, rollout, ledger - source.cut_root, source.backend_fingerprint, source.global_batch_size = "/tmp", "f", 4 - source.identity, source.shared_filesystem = SimpleNamespace(), True - context = source.context("r000004-x") - assert len(context.in_flight) == 2 and context.data["buffer_length"] == 1 - assert context.ledger["carried_over"] == 2 and context.ledger["max_policy_age"] == 1 - driver.profile.max_policy_age = 0 - assert source.context("r000004-y").in_flight == () # limit 0: section absent - - -def test_translate_emits_the_suspend_flags_for_an_agentic_run(): - import dataclasses - - from tests.test_rl_miles_adapter_config import make_config # noqa: PLC0415 - from yeto.rl.adapters.miles.config import translate_run_config - from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec, ExecutionSpec - - tolerant = AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=1), - correction=CorrectionSpec(method="tis", tis_clip=2.0, tis_clip_low=0.0)) - config = make_config() - agent = dataclasses.replace(config.agent, custom_generate_function_path=miles_policy_age.AGENTIC_GENERATE, - custom_agent_function_path=AGENT, use_session_server=True) - config = dataclasses.replace(config, agent=agent) - zero = translate_run_config(config, tolerant).argv - one = translate_run_config(config, tolerant, max_policy_age=1).argv - assert [a for a in one if a not in zero] == ["--agentic-suspend-between-turns", "--agentic-suspend-max-rounds"] - assert "--partial-rollout" not in one - i = one.index("--agentic-suspend-max-rounds") - assert one[i + 1] == "1" - - -def test_estimate_reports_why_samples_are_unscored_and_skips_unknown_trained_tokens(): - s = Sample(index=0, group_index=0) - FakePartialEngine().generate(s, 3, 2, finish=False) - FakePartialEngine().generate(s, 4, 1, finish=True) - s.weight_versions = s.weight_versions[1:] # the first turn's span is missing - s.weight_versions.insert(0, Call([Span(tok(3), len(s.tokens) - 3, len(s.tokens) - 2)])) - args = SimpleNamespace(tis_clip_low=0.0, tis_clip=2.0) - out = carry_over.estimate_cross_version_truncation(args, [s], 4, lambda _s: [-0.5, -0.5, -0.5]) - assert out["cross_version_unknown_version_tokens"] == 1 and out["cross_version_scored_tokens"] == 1 - - def failing(_s): - failing.last_error = "ValueError: alignment" - return None - - out = carry_over.estimate_cross_version_truncation(args, [s], 4, failing) - assert out["cross_version_unscored_samples"] == 1 - assert out["cross_version_unscored_reasons"] == {"ValueError: alignment": 1} - - -def _fake_prefill_module(monkeypatch): - import sys - import types - - mod = types.ModuleType("miles.rollout.generate_utils.prefill_logprobs") - mod._build_prefill_scoring_payload = lambda args, sample, params: {"input_ids": list(sample.tokens)} - mod._extract_response_logprobs = lambda sample, meta: meta["input_token_logprobs"] - for name in ("miles", "miles.rollout", "miles.rollout.generate_utils"): - monkeypatch.setitem(sys.modules, name, sys.modules.get(name) or types.ModuleType(name)) - monkeypatch.setitem(sys.modules, mod.__name__, mod) - - -def test_scorer_sends_routing_key_under_manual_router_policy(monkeypatch): - """S18 ARU-3 M1: the session server sets the router policy to ``manual``, - which rejects /generate without X-SMG-Routing-Key (one router client error - per unscored sample). The scorer must send the key.""" - import io - - _fake_prefill_module(monkeypatch) - seen: list[dict[str, str]] = [] - - def fake_urlopen(request, timeout): - seen.append({k.lower(): v for k, v in request.header_items()}) - return io.BytesIO(json.dumps({"meta_info": {"input_token_logprobs": [-0.1, -0.2]}}).encode()) - - monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) - args = SimpleNamespace(sglang_router_ip="r", sglang_router_port=1, sglang_router_policy="manual") - score = carry_over.router_prefill_scorer(args) - - keyed = SimpleNamespace(tokens=[1, 2, 3], response_length=2, routing_key="sess-7") - assert score(keyed) == [-0.1, -0.2] and score.last_error is None - assert seen[-1]["x-smg-routing-key"] == "sess-7" - - unkeyed = SimpleNamespace(tokens=[1, 2, 3], response_length=2, routing_key=None) - assert score(unkeyed) == [-0.1, -0.2] - assert seen[-1]["x-smg-routing-key"] == carry_over.SCORING_ROUTING_KEY - - args.sglang_router_policy = "cache_aware" - assert score(unkeyed) == [-0.1, -0.2] - assert "x-smg-routing-key" not in seen[-1] diff --git a/tests/test_rl_agentic_tool_call_args.py b/tests/test_rl_agentic_tool_call_args.py index 54b07e6d8..93a1e0abf 100644 --- a/tests/test_rl_agentic_tool_call_args.py +++ b/tests/test_rl_agentic_tool_call_args.py @@ -79,7 +79,7 @@ def upstream_add_arguments(monkeypatch): def _agent_argv() -> list[str]: c = sub(make_config(), "agent", custom_generate_function_path=cfg.AGENTIC_TOOL_CALL_GENERATE, - custom_agent_function_path="yeto.rl.harness.codex.codex_openenv_subprocess_agent_function.run", + custom_agent_function_path="yeto.rl.adapters.miles.harbor_agent.run", agent_max_seq_len=8192, use_session_server=True, tito_model="qwen35") return list(cfg.translate_run_config(c, AlgorithmSpec(), tito_parser_resolver=lambda m: ("qwen3", "qwen")).argv) @@ -93,7 +93,7 @@ def test_pinned_add_arguments_recognises_the_ir1_flags(upstream_add_arguments): agent_flags = argv[i:i + 2] + argv[j:j + 2] ns, unknown = parser.parse_known_args(agent_flags) assert unknown == [] - assert ns.custom_agent_function_path == "yeto.rl.harness.codex.codex_openenv_subprocess_agent_function.run" + assert ns.custom_agent_function_path == "yeto.rl.adapters.miles.harbor_agent.run" assert ns.max_seq_len == 8192 # The launcher-level spelling (--agent-max-seq-len) is NOT an upstream flag: # config.py must keep mapping agent_max_seq_len to --max-seq-len. diff --git a/tests/test_rl_algorithm_capabilities.py b/tests/test_rl_algorithm_capabilities.py index 84a6e29b7..df4de8b3c 100644 --- a/tests/test_rl_algorithm_capabilities.py +++ b/tests/test_rl_algorithm_capabilities.py @@ -136,10 +136,6 @@ def test_old_declaration_reads_as_r0_mechanisms(): f"dynamic_sampling_filters:{alg.BOUNDED_NONZERO_STD_FILTER}", f"dynamic_sampling_filters:{alg.STOCK_NONZERO_STD_FILTER}", } -# Moved from the removed legacy engine (P0-1.7b.4-Q1); declared with the R0 -# filters, it needs SecRLEnv episode evidence, so no CPU candidate exercises it. -SECRLENV_FILTER_P0 = {f"dynamic_sampling_filters:{alg.SECRLENV_GROUP_CHECK_FILTER}", - f"dynamic_sampling_filters:{alg.TBENCH_GROUP_CHECK_FILTER}"} # S24 T-TB2 # P0 built-ins without a candidate, with the reason. EXEMPT = {"advantage_estimators:ppo": "critic: refused by the rejection matrix"} @@ -147,7 +143,7 @@ def test_old_declaration_reads_as_r0_mechanisms(): def test_candidates_cover_the_p0_mechanism_registry(): registered = alg.mechanism_names() assert set(CANDIDATES) <= registered - p0_builtin = set(CANDIDATES) | set(EXEMPT) | R0_MECHANISMS_P0 | SECRLENV_FILTER_P0 + p0_builtin = set(CANDIDATES) | set(EXEMPT) | R0_MECHANISMS_P0 # every P0 built-in beyond R0 has a candidate (extension mechanisms may add more) builtin = { f"{m.dimension}:{m.name}" for m in alg.registered_mechanisms() @@ -382,7 +378,7 @@ def test_miles_and_fake_declarations(): miles = miles_capabilities(FP) assert set(MILES_DECLARED) == EXPECTED_MILES_DECLARED assert miles.declared_mechanisms() == ( - R0_MECHANISMS_P0 | SECRLENV_FILTER_P0 | EXPECTED_MILES_DECLARED + R0_MECHANISMS_P0 | EXPECTED_MILES_DECLARED ) assert all(MILES_DECLARED[m] for m in MILES_DECLARED) # every entry cites evidence # the fake declares every correction for CPU tests (and nothing else extra) diff --git a/tests/test_rl_binding_check.py b/tests/test_rl_binding_check.py index 5b8c41c02..e43220a05 100644 --- a/tests/test_rl_binding_check.py +++ b/tests/test_rl_binding_check.py @@ -5,7 +5,6 @@ import pytest -from yeto.rl import SECRLENV_GROUP_FILTER from yeto.rl.adapters.miles import binding from yeto.rl.engine import algorithm as alg from yeto.rl.engine.algorithm import AlgorithmSpec, SamplingSpec @@ -33,8 +32,7 @@ def test_every_legacy_name_normalizes_to_a_neutral_name(): - assert set(alg.LEGACY_FILTER_NAMES) == {B_OLD, S_OLD, SECRLENV_GROUP_FILTER, - "yeto.rl.harness.codex.tbench_reward.check_group"} + assert set(alg.LEGACY_FILTER_NAMES) == {B_OLD, S_OLD} for old, new in alg.LEGACY_FILTER_NAMES.items(): assert alg.normalize_filter_name(old) == new assert alg.normalize_filter_name(new) == new diff --git a/tests/test_rl_codex_schema.py b/tests/test_rl_codex_schema.py deleted file mode 100644 index 9425fd421..000000000 --- a/tests/test_rl_codex_schema.py +++ /dev/null @@ -1,116 +0,0 @@ -"""Regression tests for live stock Codex app-server schema attestation.""" - -from __future__ import annotations - -import hashlib -import importlib -from types import SimpleNamespace - -import pytest - -from yeto.rl import ( - CODEX_OPENENV_AGENT_MODULES, - CODEX_OPENENV_IDENTITY_ENV, -) -from yeto.rl.adapters.miles import island_entry as learner - - -def test_live_codex_schema_accepts_reordered_identical_json(tmp_path): - pinned = tmp_path / "pinned.json" - generated = tmp_path / "generated.json" - pinned.write_text( - '{"methods":{"initialize":{"params":{"b":2,"a":1}}},"revision":2}', - encoding="utf-8", - ) - generated.write_text( - '{\n "revision": 2,\n "methods": {"initialize": {"params": ' - '{"a": 1, "b": 2}}}\n}\n', - encoding="utf-8", - ) - assert ( - hashlib.sha256(pinned.read_bytes()).digest() - != hashlib.sha256(generated.read_bytes()).digest() - ) - - learner._verify_live_codex_app_server_schema(pinned, generated) - - -def test_live_codex_schema_rejects_semantic_drift(tmp_path): - pinned = tmp_path / "pinned.json" - generated = tmp_path / "generated.json" - pinned.write_text('{"revision":2,"methods":{}}', encoding="utf-8") - generated.write_text('{"methods":{},"revision":3}', encoding="utf-8") - - with pytest.raises(ValueError, match="schema drifted"): - learner._verify_live_codex_app_server_schema(pinned, generated) - - -def test_live_codex_schema_rejects_exact_numeric_drift(tmp_path): - pinned = tmp_path / "pinned.json" - generated = tmp_path / "generated.json" - pinned.write_text('{"revision":9007199254740992.0}', encoding="utf-8") - generated.write_text('{"revision":9007199254740993.0}', encoding="utf-8") - - with pytest.raises(ValueError, match="schema drifted"): - learner._verify_live_codex_app_server_schema(pinned, generated) - - -@pytest.mark.parametrize( - "malformed", - [ - '{"revision":', - '{"revision":2,"revision":2}', - '{"revision":NaN}', - ], -) -def test_live_codex_schema_rejects_malformed_json(tmp_path, malformed): - pinned = tmp_path / "pinned.json" - generated = tmp_path / "generated.json" - pinned.write_text('{"revision":2}', encoding="utf-8") - generated.write_text(malformed, encoding="utf-8") - - with pytest.raises(ValueError, match="schema drifted"): - learner._verify_live_codex_app_server_schema(pinned, generated) - - -def test_codex_openenv_preflight_attests_pinned_adapter_and_environment( - tmp_path, monkeypatch -): - adapter_dir = tmp_path / "examples" / "experimental" / "openenv" - adapter_dir.mkdir(parents=True) - for name in CODEX_OPENENV_AGENT_MODULES: - (adapter_dir / name).write_text("# pinned adapter\n", encoding="utf-8") - identity = { - name.removeprefix("YETO_CODEX_OPENENV_").lower(): value - for name, value in CODEX_OPENENV_IDENTITY_ENV.items() - if name.endswith("_SHA256") - } - adapter = SimpleNamespace( - __file__=str(adapter_dir / "codex_openenv_agent_function.py"), - _OPENENV_IDENTITY_ENV=dict(CODEX_OPENENV_IDENTITY_ENV), - codex_openenv_harness_identity=lambda: identity, - ) - subprocess_adapter = SimpleNamespace( - __file__=str(adapter_dir / "codex_openenv_subprocess_agent_function.py"), - run=lambda: None, - ) - - def import_module(name): - return { - "codex_openenv_agent_function": adapter, - "codex_openenv_subprocess_agent_function": subprocess_adapter, - }[name] - - monkeypatch.setattr(importlib, "import_module", import_module) - for name, value in CODEX_OPENENV_IDENTITY_ENV.items(): - monkeypatch.setenv(name, value) - - args = SimpleNamespace(miles_root=str(tmp_path)) - learner._preflight_codex_openenv_adapter(args, "qwen35_08b") - - monkeypatch.delenv("YETO_CODEX_OPENENV_MODEL_REVISION") - with pytest.raises(ValueError, match="environment drifted"): - learner._preflight_codex_openenv_adapter(args, "qwen35_08b") - # qwen35 joined OPENENV_BACKEND_PROFILES in S17 M1, so use a profile outside the allowlist. - with pytest.raises(ValueError, match="requires backend profile"): - learner._preflight_codex_openenv_adapter(args, "not_an_openenv_profile") diff --git a/tests/test_rl_compactionrl_counts_receipt.py b/tests/test_rl_compactionrl_counts_receipt.py index eff525f7c..2eb0aba77 100644 --- a/tests/test_rl_compactionrl_counts_receipt.py +++ b/tests/test_rl_compactionrl_counts_receipt.py @@ -6,7 +6,7 @@ from yeto.rl.adapters.miles import rollout_meta_hook as hook from yeto.rl.adapters.miles.trainer import receipt_trajectory_ids -from yeto.rl.harness.codex.codex_openenv_generate import refresh_segment_counts +from yeto.rl.compaction import refresh_segment_counts def _batch(): diff --git a/tests/test_rl_engine_ports.py b/tests/test_rl_engine_ports.py index 8fcfc6aff..e808a7649 100644 --- a/tests/test_rl_engine_ports.py +++ b/tests/test_rl_engine_ports.py @@ -138,35 +138,3 @@ def test_connect_island_ray_forwards_elastic_metadata_env_only_when_on(): assert ELASTIC_METADATA_ENV not in off.init_calls[0][1]["env_vars"] -def test_connect_island_ray_forwards_codex_harness_env_and_learner_id_to_workers(): - """A-T3-4 (codex-smoke-20261003-6): the agent function runs inside RolloutExecutor - actors and must self-configure from the job runtime_env, not the raylet env.""" - from types import SimpleNamespace - - from yeto.rl.adapters.miles.entry import connect_island_ray - - machine = _TwoRayMachine(raylet_env={}) - environ = { - "RAY_ADDRESS": "a:6379", - "YETO_HARNESS_ENVIRONMENT_PROVIDER": "yeto.rl.harness.codex.tb2_provider:modal_provider", - "YETO_HARNESS_TB2_FAULT": "create_fail:2", - "YETO_CODEX_OPENENV_MODEL_REVISION": "abc", - "TBENCH_REWARD_HMAC_KEY": "k", - "YETO_SANDBOX_MODAL_TOKEN_SECRET": "s", - "MODAL_TOKEN_SECRET": "mainTok", - "UNRELATED": "x", - } - connect_island_ray(environ=environ, ray_module=machine, - miles_args=SimpleNamespace(yeto_rl_learner_id=3, yeto_rl_cell_id="c7")) - env_vars = machine.init_calls[0][1]["env_vars"] - for key in ("YETO_HARNESS_ENVIRONMENT_PROVIDER", "YETO_HARNESS_TB2_FAULT", "YETO_CODEX_OPENENV_MODEL_REVISION", - "YETO_SANDBOX_MODAL_TOKEN_SECRET"): - assert env_vars[key] == environ[key] - assert "MODAL_TOKEN_SECRET" not in env_vars # secret-handling-hardening - assert "TBENCH_REWARD_HMAC_KEY" not in env_vars # retired reward key - assert env_vars["YETO_RL_LEARNER_ID"] == "3" and env_vars["YETO_RL_CELL_ID"] == "c7" - assert "UNRELATED" not in env_vars - # without miles_args (legacy callers) nothing harness-specific is forwarded - bare = _TwoRayMachine(raylet_env={}) - connect_island_ray(environ=environ, ray_module=bare) - assert "YETO_HARNESS_ENVIRONMENT_PROVIDER" not in bare.init_calls[0][1]["env_vars"] diff --git a/tests/test_rl_eval_infer.py b/tests/test_rl_eval_infer.py deleted file mode 100644 index 184740ea9..000000000 --- a/tests/test_rl_eval_infer.py +++ /dev/null @@ -1,450 +0,0 @@ -"""rl-eval-difficulty-buckets 5.9/5.10/5.1/1.1/1.3: eval inference agent, Modal function, roles, data (CPU only).""" - -from __future__ import annotations - -import asyncio -import hashlib -import json -import re -from pathlib import Path -from types import SimpleNamespace - -import pytest - -from yeto.rl.eval import codex_infer as ci -from yeto.rl.eval.export import StoreEvalExporter -from yeto.rl.eval.island import EvalTask -from yeto.rl.eval.store import EvalStore -from yeto.rl.eval.tb2_attempt import tb2_attempt - -REPO = Path(__file__).resolve().parents[1] -MILES_SESSION_CONFIG = Path("/home/michael/work/s17-infra-fork-m/miles/rollout/session/config.py") - - -# --- PEFT writer pieces --------------------------------------------------------------------- - - -def test_rank_and_targets_from_tensor_shapes(): - shapes = {"base_model.model.layers.0.self_attn.o_proj.lora_A.weight": (16, 1024), - "base_model.model.layers.0.self_attn.o_proj.lora_B.weight": (1024, 16), - "base_model.model.layers.0.linear_attn.out_proj.lora_A.weight": (16, 2048), - "base_model.model.layers.0.linear_attn.out_proj.lora_B.weight": (1024, 16)} - assert ci.lora_rank_and_targets(shapes) == (16, ["o_proj", "out_proj"]) - with pytest.raises(ValueError, match="single LoRA rank"): - ci.lora_rank_and_targets({**shapes, "x.q_proj.lora_A.weight": (8, 4)}) - - -def test_peft_config_is_what_initial_adapter_accepts(): - cfg = ci.peft_adapter_config(rank=16, targets=["o_proj"], base_model="Qwen/Qwen3.5-0.8B", revision="r") - assert cfg["peft_type"] == "LORA" and cfg["task_type"] == "CAUSAL_LM" - assert cfg["lora_alpha"] == cfg["r"] == 16 and cfg["lora_dropout"] == 0.0 and cfg["bias"] == "none" - assert cfg["rank_pattern"] == cfg["alpha_pattern"] == {} and not cfg["use_rslora"] and not cfg["use_dora"] - - -def test_write_peft_dir_checks_tensor_hash(tmp_path): - torch = pytest.importorskip("torch") - pytest.importorskip("safetensors") - from yeto.rl.core import canonical_lora_config_hash, canonical_state - from yeto.rl.engine.trainable_state import TrainableState - from yeto.rl.eval.export import canonical_policy_files - from yeto.rl.eval.store import EvalIntegrityError - - tensors = {"m.layers.0.self_attn.o_proj.lora_A.weight": torch.ones(4, 8), - "m.layers.0.self_attn.o_proj.lora_B.weight": torch.zeros(8, 4)} - state = TrainableState.from_lora(canonical_state( - 3, tensors, base_model_revision="a" * 40, - lora_config_hash=canonical_lora_config_hash(rank=4, target_modules=["o_proj"]))) - store = EvalStore(tmp_path / "s") - m = store.put_version(3, files=canonical_policy_files(state), policy_tensor_hash=state.policy_tensor_hash(), - policy_token="yeto:3:x", sampling={}) - info = ci.write_peft_dir(store.files_dir(3), m, tmp_path / "peft", base_model="B") - assert info["rank"] == 4 and info["targets"] == ["o_proj"] - assert json.loads((tmp_path / "peft/adapter_config.json").read_text())["revision"] == "a" * 40 - with pytest.raises(EvalIntegrityError): - ci.write_peft_dir(store.files_dir(3), {**m, "policy_tensor_hash": "0" * 64}, tmp_path / "p2", base_model="B") - - -# --- SGLang / session server supervision ------------------------------------------------------ - - -def test_sglang_argv_and_session_config_match_training(): - cfg = ci.InferConfig(base_model="/models/q", revision="r", sglang_extra_args=("--x", "1")) - serving = {"reasoning_parser": "qwen3", "tool_call_parser": "qwen3_coder", "chat_template_path": "/t.jinja", - "chat_template_kwargs": {"enable_thinking": True}} - argv = ci.sglang_argv(cfg, rank=16, targets=["o_proj", "out_proj"], python="py", serving=serving) - joined = " ".join(argv) - assert "--reasoning-parser qwen3 --tool-call-parser qwen3_coder --chat-template /t.jinja" in joined - assert "parser" not in " ".join(ci.sglang_argv(cfg, rank=16, targets=["o_proj"], python="py")) - assert "--enable-lora --max-lora-rank 16 --lora-target-modules o_proj out_proj" in joined - assert "--context-length 8192" in joined and argv[-2:] == ["--x", "1"] - sc2 = ci.session_server_config(cfg, rank=16, serving=serving) - assert sc2["chat_template_path"] == "/t.jinja" - assert sc2["apply_chat_template_kwargs"] == {"clear_thinking": False, "enable_thinking": True} - with pytest.raises(ValueError, match="conflicts"): - ci.session_server_config(cfg, rank=16, serving={"chat_template_kwargs": {"clear_thinking": True}}) - assert "--revision" not in argv or not Path("/models/q").exists() - sc = ci.session_server_config(cfg, rank=16) - assert sc["tito_model"] == "qwen35" and sc["apply_chat_template_kwargs"] == {"clear_thinking": False} - assert sc["lora_rank"] == 16 and sc["backend_url"] == "http://127.0.0.1:30000" - assert sc["use_rollout_routing_replay"] is False and sc["session_message_matcher"] == "strict" - if MILES_SESSION_CONFIG.is_file(): # every SessionServerConfig field, nothing else (strict model) - text = MILES_SESSION_CONFIG.read_text().split("def compute_session_server_config")[0] - fields = set(re.findall(r"^ (\w+): ", text, re.M)) - assert set(sc) == fields - - -class _Proc: - def __init__(self): - self.returncode = None - self.terminated = False - - def poll(self): - return self.returncode - - def terminate(self): - self.terminated = True - self.returncode = 0 - - def wait(self, timeout=None): - return 0 - - -class _Http: - def __init__(self, refuse_load=False): - self.calls = [] - self.refuse_load = refuse_load - - def __call__(self, method, url, body=None, timeout=60.0): - self.calls.append((method, url.split(":", 2)[-1], body)) - if url.endswith("/load_lora_adapter") and self.refuse_load: - return 400, {"success": False, "message": "bad"} - if url.endswith("/sessions") and method == "POST": - return 200, {"session_id": "s1"} - return 200, {} - - -def _loader(monkeypatch, tmp_path, *, rank=16, targets=("o_proj",), refuse=False): - spawned = [] - - def popen(argv, **kw): - spawned.append(argv) - return _Proc() - - def fake_write(files_dir, manifest, out_dir, *, base_model): - return {"rank": rank, "targets": list(targets), "tensors": 2, "path": str(out_dir), "seconds": 0.1} - - monkeypatch.setattr(ci, "write_peft_dir", fake_write) - events = [] - http = _Http(refuse_load=refuse) - loader = ci.SglangSessionLoader(ci.InferConfig(base_model="B", work_dir=str(tmp_path)), popen=popen, - http=http, sleep=lambda s: None, - emit=lambda e, **f: events.append((e, f)), - serving=lambda tito: {"reasoning_parser": "qwen3", - "tool_call_parser": "qwen3_coder", - "chat_template_path": "/t.jinja"}) - return loader, spawned, http, events - - -def test_loader_starts_once_then_swaps_adapter(monkeypatch, tmp_path): - loader, spawned, http, events = _loader(monkeypatch, tmp_path) - assert loader.load({"policy_version": 0, "rl/policy_token": "yeto:0:h"}, tmp_path) == "yeto:0:h" - assert len(spawned) == 2 and "sglang.launch_server" in spawned[0] - assert spawned[0][spawned[0].index("--tool-call-parser") + 1] == "qwen3_coder" - assert any(c[1].endswith("/health_generate") for c in http.calls) - assert loader.load({"policy_version": 10, "rl/policy_token": "yeto:10:h"}, tmp_path) == "yeto:10:h" - assert len(spawned) == 2 # no restart - paths = [c[1] for c in http.calls if c[0] == "POST"] - assert paths.count("30000/load_lora_adapter") == 2 and paths.count("30000/unload_lora_adapter") == 1 - load_bodies = [c[2] for c in http.calls if c[1].endswith("/load_lora_adapter")] - assert {b["lora_name"] for b in load_bodies} == {ci.SERVED_LORA_NAME} - assert [e for e, _ in events] == ["rl_eval_policy_load"] * 2 and "sglang_start_s" in events[0][1] - assert events[0][1]["serving"]["tool_call_parser"] == "qwen3_coder" - assert json.loads(spawned[1][-1])["chat_template_path"] == "/t.jinja" - loader.close() - assert all(p is not None for p in loader.procs.values()) - - -def test_loader_start_base_serves_without_lora(monkeypatch, tmp_path): - loader, spawned, http, _ = _loader(monkeypatch, tmp_path) - out = loader.start_base() - assert "sglang_start_s" in out and len(spawned) == 2 - assert not any(a.startswith("--enable-lora") or "lora" in a for a in spawned[0]) - session = json.loads(spawned[1][-1]) - assert session["lora_rank"] == 0 and session["lora_adapter_path"] is None - assert not [c for c in http.calls if "lora" in c[1]] - with pytest.raises(RuntimeError, match="already started"): - loader.start_base() - - -def test_loader_fails_closed(monkeypatch, tmp_path): - loader, *_ = _loader(monkeypatch, tmp_path, refuse=True) - with pytest.raises(RuntimeError, match="refused the adapter"): - loader.load({"policy_version": 0, "rl/policy_token": "t"}, tmp_path) - loader, *_ = _loader(monkeypatch, tmp_path) - loader.load({"policy_version": 0, "rl/policy_token": "t"}, tmp_path) - monkeypatch.setattr(ci, "write_peft_dir", lambda *a, **k: {"rank": 8, "targets": ["o_proj"], "path": "p", - "seconds": 0}) - with pytest.raises(RuntimeError, match="differs from the served"): - loader.load({"policy_version": 1, "rl/policy_token": "t"}, tmp_path) - - -def test_loader_reports_a_dead_server(monkeypatch, tmp_path): - loader, *_ = _loader(monkeypatch, tmp_path) - - def popen(argv, **kw): - p = _Proc() - p.returncode = 1 - return p - - loader.popen = popen - with pytest.raises(RuntimeError, match="sglang exited"): - loader.load({"policy_version": 0, "rl/policy_token": "t"}, tmp_path) - - -# --- the agent port ----------------------------------------------------------------------------- - - -def _lease(tmp_path, instruction="Do the thing."): - task = SimpleNamespace(instruction=instruction) - return SimpleNamespace(env_url="http://127.0.0.1:9", env_token="tok", deadline_seconds=5.0, - worker_env={"W": "1"}, task=task) - - -def _agent(drive): - http = _Http() - return ci.CodexEvalAgent(ci.InferConfig(base_model="B", top_k=20), drive=drive, http=http), http - - -def test_agent_runs_the_training_worker_job(tmp_path): - seen = {} - - async def drive(job, trajectory, board, worker_env): - seen.update(job=job, board=board, worker_env=worker_env) - return {"status": "max_turns", "metrics": {"turns": 12, "max_model_total_tokens": 7000, - "end_reason": "CodexModelFailure: x", "secret": 1}, - "episode_id": job["episode_id"]} - - agent, http = _agent(drive) - out = asyncio.run(agent(_lease(tmp_path), EvalTask("a", "tb2-easy"), 0, policy_token="t")) - job = seen["job"] - assert job["base_url"] == "http://127.0.0.1:30100/sessions/s1" and job["prompt"] == "Do the thing." - assert job["request_kwargs"] == {"temperature": 1.0, "top_p": 1.0, "max_tokens": 4096, "top_k": 20} - assert job["max_seq_len"] == 8192 and job["env_token"] == "tok" and job["driver"] == "stock" - assert seen["worker_env"] == {"W": "1"} and seen["board"] is None - assert out["end_reason"] == "max_turns" and out["turns"] == 12 and out["tokens"] == 7000 - assert out["agent_detail"] == {"end_reason": "CodexModelFailure: x", "turns": 12, "max_model_total_tokens": 7000} - create = [c for c in http.calls if c[1].endswith("/sessions")][0] - assert create[2]["evaluation"] is True and create[2]["top_k"] == 20 - assert http.calls[-1][0] == "DELETE" - - -def test_agent_maps_statuses_and_harness_errors(tmp_path): - from yeto.rl.harness.codex import codex_harness_agent as harness - - async def timeout_drive(job, *a): - return {"status": "timeout", "metrics": {}, "episode_id": job["episode_id"]} - - agent, _ = _agent(timeout_drive) - assert asyncio.run(agent(_lease(tmp_path), EvalTask("a", "x"), 0, policy_token="t"))["end_reason"] == "timed_out" - - async def broken(job, *a): - raise harness.CodexHarnessError("protocol violation") - - agent, http = _agent(broken) - out = asyncio.run(agent(_lease(tmp_path), EvalTask("a", "x"), 0, policy_token="t")) - assert out["end_reason"] == "infra_error" and "protocol" in out["error"] - assert http.calls[-1][0] == "DELETE" # session released on error too - - async def hang(job, *a): - await asyncio.sleep(10) - - agent, _ = _agent(hang) - lease = _lease(tmp_path) - lease.deadline_seconds = 0.05 - assert asyncio.run(agent(lease, EvalTask("a", "x"), 0, policy_token="t"))["end_reason"] == "timed_out" - - -def test_tb2_attempt_records_timing_and_results_hash_ignores_it(tmp_path): - class Verifier: - async def evaluate(self, episode_id): - return {"passed": True} - - class Provider: - async def acquire(self, task_id, tid): - async def destroy(): - return None - return SimpleNamespace(verifier=Verifier(), destroy=destroy) - - async def agent(lease, task, trial, *, policy_token): - return {"episode_id": "e", "end_reason": "completed", "turns": 2, "agent_detail": {"end_reason": "x"}} - - out = tb2_attempt(Provider(), agent)(EvalTask("a", "x"), 0, policy_version=0, policy_token="t") - assert set(out["timing"]) == {"sandbox_s", "agent_s", "judge_s", "destroy_s"} - from yeto.rl.eval.island import EvalIsland - - rec = EvalIsland._record(SimpleNamespace(island_id="e"), EvalTask("a", "tb2-easy"), 0, 0, "t", out, 1.0) - assert rec["timing"] == out["timing"] and rec["agent_detail"] == {"end_reason": "x"} - s1, s2 = EvalStore(tmp_path / "1"), EvalStore(tmp_path / "2") - s1.append_unit(rec) - s2.append_unit({**rec, "timing": {"agent_s": 99.0}, "agent_detail": {"end_reason": "y"}}) - assert s1.results_sha256(0) == s2.results_sha256(0) - - -# --- training-finished marker (5.10) -------------------------------------------------------------- - - -def test_final_export_marks_training_finished(tmp_path): - store = EvalStore(tmp_path) - exp = StoreEvalExporter(store, sampling={}, interval=10, - policy_files=lambda s: {"policy.safetensors": b"w"}) - state = SimpleNamespace(policy_tensor_hash=lambda: "h") - exp(10, state, "yeto:10:h") - assert not store.training_finished() - exp(13, state, "yeto:13:h", final=True) - assert store.training_finished() - assert json.loads((tmp_path / "training-finished.json").read_text()) == {"final_version": 13} - - -# --- Modal function definition (fake modal) ------------------------------------------------------ - - -def test_eval_function_spec_env_and_client(): - from yeto.cloud import modal_eval_island as mei - - spec = mei.EvalFunctionSpec(app_name="a", volume="vol", workdir=".", codex_dir="c", tb2_dir="t", - envs={"YETO_CODEX_VERSION": "x"}) - env = spec.container_env() - assert env["YETO_RL_ISLAND_ROLE"] == "eval" and env["YETO_RL_EVAL_STORE"] == mei.EVAL_STORE_MOUNT - assert env["YETO_RL_EVAL_STORE_MOUNTED"] == "1" and env["YETO_RL_EVAL_VOLUME"] == "vol" - assert env["YETO_HARNESS_TB2_TASKS_DIR"] == mei.TB2_MOUNT and env["YETO_CODEX_VERSION"] == "x" - assert spec.image_ref.startswith("ghcr.io/") and "@sha256:" in spec.image_ref - - class Call: - def get(self): - return {"evaluated": [0]} - - class Fn: - def __init__(self): - self.args = [] - - def spawn(self, cfg): - self.args.append(cfg) - return Call() - - fn = Fn() - events = [] - launcher = mei.EvalIslandLauncher(client=mei.ModalFunctionClient(fn), emit=lambda e, **f: events.append(e), - config={"store": "s"}, gpu="H100!", gpus=1) - assert launcher.run() == {"evaluated": [0]} - assert fn.args == [{"store": "s"}] and events == ["rl_eval_island", "rl_eval_island"] - - -# --- roles: launcher + island ledger (5.1) --------------------------------------------------------- - - -def test_eval_role_never_joins_the_merge_pool(): - from yeto.launcher import ISLAND_ROLES, island_role_env - from yeto.rl.engine.island_ledger import CrossIslandLedger, LedgerError - from yeto.rl.eval.island import EVAL_ISLAND_ROLE - - assert EVAL_ISLAND_ROLE in ISLAND_ROLES and island_role_env("eval") == {"YETO_RL_ISLAND_ROLE": "eval"} - with pytest.raises(ValueError): - island_role_env("judge") - ledger = CrossIslandLedger(mode="elastic") - ledger.join("i0", now=0.0) - with pytest.raises(LedgerError, match="cannot join the merge pool"): - ledger.join("eval-0", now=0.0, role="eval") - assert set(ledger.members) == {"i0"} - - -# --- eval / train jsonl (1.1 second half, 1.3) ------------------------------------------------------- - - -def _fake_tb2(root: Path, spec): - for tid, diff in spec: - d = root / tid - (d / "tests").mkdir(parents=True) - (d / "tests/test.sh").write_text("echo ok\n") - (d / "instruction.md").write_text(f"solve {tid}\n") - (d / "task.toml").write_text(f'[metadata]\ndifficulty = "{diff}"\n[environment]\ndocker_image = "img/{tid}"\n') - - -def test_build_eval_data_split_and_determinism(tmp_path): - import importlib.util - - from yeto.rl.harness.reward_env.benchmark import HOLDOUT_SCHEMA - - mod_spec = importlib.util.spec_from_file_location("build_eval_data", REPO / "tools/build_eval_data.py") - tool = importlib.util.module_from_spec(mod_spec) - mod_spec.loader.exec_module(tool) - _fake_tb2(tmp_path / "tb2", [("a", "easy"), ("b", "medium"), ("c", "hard"), ("d", "medium")]) - hold = {"schema": HOLDOUT_SCHEMA, "benchmark": "tb2", "benchmark_version": "tb2@unknown", "seed": 1, - "rule": "t", "items": [{"task_id": "b", "difficulty": "medium", "eval_bucket": "tb2-medium"}, - {"task_id": "c", "difficulty": "hard", "eval_bucket": "tb2-hard"}]} - out = tool.build(tmp_path / "tb2", hold) - assert out == tool.build(tmp_path / "tb2", hold) # same inputs, same bytes - ev = [json.loads(x) for x in out[tool.EVAL_JSONL].decode().splitlines()] - tr = [json.loads(x) for x in out[tool.TRAIN_JSONL].decode().splitlines()] - assert [r["metadata"]["task_id"] for r in ev] == ["b", "c"] and [r["metadata"]["task_id"] for r in tr] == ["a", "d"] - assert ev[1]["metadata"] == {"task_id": "c", "benchmark": "tb2", "benchmark_version": "tb2@unknown", - "difficulty": "hard", "difficulty_source": "tb2-task.toml", "eval_bucket": "tb2-hard"} - assert "eval_bucket" not in tr[0]["metadata"] and tr[0]["metadata"]["difficulty"] == "easy" - sha = json.loads(out[tool.SHA_JSON]) - assert sha[tool.EVAL_JSONL]["sha256"] == hashlib.sha256(out[tool.EVAL_JSONL]).hexdigest() - with pytest.raises(ValueError, match="checkout"): - tool.build(tmp_path / "tb2", {**hold, "benchmark_version": "tb2@other"}) - # S17 N13: the unusable table leaves training; an unusable task in the hold-out list is refused - tr = [json.loads(x) for x in tool.build(tmp_path / "tb2", hold, unusable={"d": "unusable/x: y"})[ - tool.TRAIN_JSONL].decode().splitlines()] - assert [r["metadata"]["task_id"] for r in tr] == ["a"] - with pytest.raises(ValueError, match="hold-out"): - tool.build(tmp_path / "tb2", hold, unusable={"b": "unusable/x: y"}) - - -def test_shipped_eval_and_train_jsonl_match_their_sha_record(): - sha = json.loads((REPO / "data/eval/tb2-data.sha256.json").read_text()) - for rel, n in (("eval/tb2-holdout-eval.jsonl", 30), ("tb2/tb2-train.jsonl", 46)): - data = (REPO / "data" / rel).read_bytes() - assert sha[rel] == {"rows": n, "sha256": hashlib.sha256(data).hexdigest()} - hold = json.loads((REPO / "data/eval/tb2-holdout.json").read_text()) - ev = {json.loads(x)["metadata"]["task_id"] for x in (REPO / "data/eval/tb2-holdout-eval.jsonl").read_text().splitlines()} - tr = {json.loads(x)["metadata"]["task_id"] for x in (REPO / "data/tb2/tb2-train.jsonl").read_text().splitlines()} - assert ev == {i["task_id"] for i in hold["items"]} and not ev & tr - unusable = {i["task_id"] for i in json.loads((REPO / "data/eval/tb2-unusable.json").read_text())["items"]} - assert not unusable & (ev | tr) # S17 N13: unusable tasks leave both sets - assert {e["task_id"] for e in hold["excluded"]} - unusable <= tr # smoke-6: out of the eval pool, still trainable - split = json.loads((REPO / "data/eval/tb2-train.json").read_text()) - assert tr == {i["task_id"] for i in split["items"]} and sha["unusable_excluded"] == sorted(unusable) - - -def test_eval_island_body_keeps_logs_off_the_volume_until_the_end(tmp_path, monkeypatch): - import subprocess - - from yeto.cloud import modal_eval_island as mei - - import os - import sys - - monkeypatch.setenv("PYTHONPATH", os.environ.get("PYTHONPATH", "")) # restored after the test - monkeypatch.setenv("YETO_MODAL_RUNTIME_PYTHONPATH", "") - monkeypatch.setattr(sys, "path", list(sys.path)) - monkeypatch.setattr(subprocess, "run", lambda *a, **k: SimpleNamespace(stdout="NVIDIA H100 80GB HBM3\n")) - seen = {} - - def fake_main(config): - seen["infer"] = config["infer"] - Path(config["infer"]["log_dir"]).mkdir(parents=True) - (Path(config["infer"]["log_dir"]) / "sglang.log").write_text("up") - return {"evaluated": [0], "loads": []} - - monkeypatch.setattr(mei, "island_main", fake_main) - out = mei.eval_island_body({"store": str(tmp_path / "vol"), "island_id": "e1", "assert_gpu_name": "H100", - "local_run_dir": str(tmp_path / "local"), "infer": {"base_model": "B"}}) - assert out["evaluated"] == [0] and out["gpu_names"] == ["NVIDIA H100 80GB HBM3"] - assert seen["infer"]["log_dir"].startswith(str(tmp_path / "local")) - assert (tmp_path / "vol/runs/e1/logs/sglang.log").read_text() == "up" - assert (tmp_path / "vol/runs/e1/body.json").is_file() - assert os.environ["PYTHONPATH"] == mei.IMAGE_PYTHONPATH - with pytest.raises(RuntimeError, match="asked for A100"): - mei.eval_island_body({"store": str(tmp_path / "vol"), "assert_gpu_name": "A100", - "local_run_dir": str(tmp_path / "l2")}) diff --git a/tests/test_rl_eval_island.py b/tests/test_rl_eval_island.py index e8e0c1c6e..6a242ae75 100644 --- a/tests/test_rl_eval_island.py +++ b/tests/test_rl_eval_island.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import hashlib import json from pathlib import Path @@ -12,9 +11,8 @@ from yeto.rl.eval import guard, stats from yeto.rl.eval.export import StoreEvalExporter, is_eval_version -from yeto.rl.eval.island import EvalIsland, EvalPlan, EvalSettingError, EvalTask, Preempted, jsonl_emitter +from yeto.rl.eval.island import EvalIsland, EvalPlan, EvalSettingError, Preempted, jsonl_emitter from yeto.rl.eval.store import EvalIntegrityError, EvalStore -from yeto.rl.eval.tb2_attempt import tb2_attempt, trajectory_id from yeto.rl.harness.reward_env.benchmark import HOLDOUT_SCHEMA, holdout_sha256 REPO = Path(__file__).resolve().parents[1] @@ -359,52 +357,6 @@ def test_driver_hook_emits_rl_eval_export(tmp_path): assert len(events) == 1 -# --- TB2 attempt over a fake provider ----------------------------------------------------------- - - -class _Verifier: - def __init__(self, passed): - self.passed = passed - - async def evaluate(self, episode_id): - return {"passed": self.passed, "timed_out": False} - - -class _Provider: - def __init__(self, passed=True, fail=False): - self.passed, self.fail, self.destroyed, self.ids = passed, fail, 0, [] - - async def acquire(self, task_id, tid): - if self.fail: - raise RuntimeError("no sandbox") - self.ids.append(tid) - - async def destroy(): - self.destroyed += 1 - return SimpleNamespace(verifier=_Verifier(self.passed), destroy=destroy) - - -def test_tb2_attempt_judges_and_always_destroys(): - async def agent(lease, task, trial, *, policy_token): - return {"episode_id": "e", "end_reason": "max_turns", "turns": 30, "tokens": 900} - - task = EvalTask("a", "tb2-easy") - prov = _Provider(passed=True) - out = tb2_attempt(prov, agent)(task, 1, policy_version=10, policy_token="t") - assert out["success"] and out["reward"] == 1.0 and out["end_reason"] == "max_turns" - assert prov.ids == [trajectory_id(10, "a", 1)] and prov.destroyed == 1 - - async def broken(lease, task, trial, *, policy_token): - raise Preempted("gone") - - prov2 = _Provider() - with pytest.raises(Preempted): - tb2_attempt(prov2, broken)(task, 0, policy_version=10, policy_token="t") - assert prov2.destroyed == 1 - with pytest.raises(RuntimeError, match="no sandbox"): - tb2_attempt(_Provider(fail=True), agent)(task, 0, policy_version=10, policy_token="t") - - # --- Modal launcher / store (fake client) ------------------------------------------------------ @@ -507,15 +459,6 @@ def test_island_main_with_dotted_fakes(tmp_path, monkeypatch): assert tape[0]["event"] == "rl_eval" and tape[0]["eval/units"] == 2 -def test_asyncio_default_runner_is_used(): - # tb2_attempt defaults to asyncio.run: works from sync code - async def agent(lease, task, trial, *, policy_token): - return {"end_reason": "completed"} - - assert tb2_attempt(_Provider(passed=False), agent, run=asyncio.run)( - EvalTask("a", "x"), 0, policy_version=0, policy_token="t")["success"] is False - - def test_canonical_policy_files_round_trip_policy_tensor_hash(tmp_path): torch = pytest.importorskip("torch") pytest.importorskip("safetensors") diff --git a/tests/test_rl_ir_harness.py b/tests/test_rl_ir_harness.py index a30cde8f7..41a3a68d3 100644 --- a/tests/test_rl_ir_harness.py +++ b/tests/test_rl_ir_harness.py @@ -7,7 +7,6 @@ from __future__ import annotations import asyncio -import dataclasses import json import time from types import SimpleNamespace @@ -15,7 +14,6 @@ import pytest import torch -from yeto.rl.engine import run_config as rc from yeto.rl.engine.algorithm import AlgorithmSpec from yeto.rl.engine.bridges import LocalOnlySync from yeto.rl.engine.driver import EventTape, IslandDriver, PolicyIdentityError, policy_token @@ -53,11 +51,11 @@ # --------------------------------------------------------------------------- def test_ir1_agent_flags_pass_through_with_agentic_generate(): c = sub(make_config(), "agent", custom_generate_function_path=AGENTIC, - custom_agent_function_path="yeto.rl.harness.codex.codex_openenv_agent_function", + custom_agent_function_path="yeto.rl.adapters.miles.harbor_agent.run", agent_max_seq_len=8192) argv = cfg.translate_run_config(c, AlgorithmSpec()).argv i = argv.index("--custom-agent-function-path") - assert argv[i + 1] == "yeto.rl.harness.codex.codex_openenv_agent_function" + assert argv[i + 1] == "yeto.rl.adapters.miles.harbor_agent.run" assert argv[argv.index("--max-seq-len") + 1] == "8192" assert argv[argv.index("--custom-generate-function-path") + 1] == AGENTIC @@ -82,20 +80,6 @@ def test_ir1_append_roles_rejected_with_template_reason(): cfg.translate_run_config(c, AlgorithmSpec()) -def test_ir1_append_roles_accepted_for_signed_codex_agent_when_equal_to_profile(): - from yeto.rl import CODEX_OPENENV_AGENT - c = sub(make_config(), "agent", use_session_server=True, tito_model="qwen35", - custom_generate_function_path=AGENTIC, custom_agent_function_path=CODEX_OPENENV_AGENT, - tito_allowed_append_roles=("tool", "user")) - stub = lambda _m: (None, None) # noqa: E731 - no miles checkout in the CPU venv - launch = cfg.translate_run_config(c, AlgorithmSpec(), tito_parser_resolver=stub) - assert "--tito-allowed-append-roles" not in " ".join(map(str, launch.argv)) - assert "--tito-model" in launch.argv - bad = sub(c, "agent", tito_allowed_append_roles=("tool",)) - with pytest.raises(cfg.UnmappedConfigError, match=r"fix the\s+append roles"): - cfg.translate_run_config(bad, AlgorithmSpec(), tito_parser_resolver=stub) - - def test_ir1_session_server_and_partial_rollout_are_mutually_exclusive(): with pytest.raises(cfg.UnmappedConfigError) as err: cfg.check_session_server_partial_rollout(True, True) diff --git a/tests/test_rl_launcher.py b/tests/test_rl_launcher.py index 8dd19f953..ad2c7203b 100644 --- a/tests/test_rl_launcher.py +++ b/tests/test_rl_launcher.py @@ -20,20 +20,14 @@ syncer_command, ) from yeto.rl import ( - CODEX_OPENENV_AGENT, MILES_NEXT_COMMIT, MILES_NEXT_REPOSITORY, MILES_PEFT_VERSION, - SECRLENV_GENERATE, - SECRLENV_GROUP_FILTER, - SECRLENV_REWARD, SGLANG_NEXT_COMMIT, SGLANG_NEXT_REPOSITORY, ) from yeto.rl.adapters.miles import island_entry as rl_learner from yeto.rl.codex_backend import ( - QWEN35_08B_MODEL, - QWEN35_08B_REVISION, QWEN35_MODEL, QWEN35_REVISION, ) @@ -158,48 +152,6 @@ def test_over_sampling_below_the_batch_is_rejected(): _prepare_rl_args(args) -def test_legacy_secrlenv_agent_auto_binds_exact_replacement_contract(): - args = _args( - [ - "--reward-function", - SECRLENV_REWARD, - "--custom-generate-function-path", - SECRLENV_GENERATE, - "--custom-agent-function-path", - "yeto.rl.harness.codex.agent.run", - "--use-session-server", - "--tito-model", - "org/model", - ] - ) - - _prepare_rl_args(args) - - assert args.dynamic_sampling_filter_path == SECRLENV_GROUP_FILTER - assert args.dynamic_sampling_max_replacements == 0 - assert args.secrlenv_max_infrastructure_replacements == 1 - assert args.over_sampling_batch_size == args.rollout_batch_size + 1 - - -def test_secrlenv_agent_rejects_the_unwrapped_miles_generator(): - args = _args( - [ - "--reward-function", - SECRLENV_REWARD, - "--custom-generate-function-path", - "miles.rollout.generate_hub.agentic_tool_call.generate", - "--custom-agent-function-path", - "yeto.rl.harness.codex.agent.run", - "--use-session-server", - "--tito-model", - "org/model", - ] - ) - - with pytest.raises(ValueError, match="signed SecRLEnv generate wrapper"): - _prepare_rl_args(args) - - def test_decoupled_rl_preset_fixes_the_fragment_outer_contract(): args = _args( ( @@ -404,45 +356,6 @@ def test_rl_maps_long_task_and_oversampling_options_to_miles(): assert args.tito_allowed_append_roles == ["tool", "user"] -def test_codex_openenv_uses_exact_08b_profile_with_qwen35_tito_family(): - args = _args( - ( - "--model", - QWEN35_08B_MODEL, - "--model-revision", - QWEN35_08B_REVISION, - "--lora-targets", - "attention", - "--custom-generate-function-path", - "miles.rollout.generate_hub.agentic_tool_call.generate", - "--custom-agent-function-path", - CODEX_OPENENV_AGENT, - "--codex-reasoning-effort", - "xhigh", - "--codex-backend-profile", - "qwen35_08b", - "--use-session-server", - "--tito-model", - "qwen35", - "--tito-allowed-append-roles", - "tool", - "user", - ) - ) - - _prepare_rl_args(args) - - assert args.codex_backend_profile == "qwen35_08b" - assert args.tito_model == "qwen35" - assert args.model == QWEN35_08B_MODEL - assert args.model_revision == QWEN35_08B_REVISION - assert args.apply_chat_template_kwargs == {"clear_thinking": False} - - args.tito_model = "qwen35_08b" - with pytest.raises(ValueError, match="Miles TITO family"): - _prepare_rl_args(args) - - def test_rl_parses_cybergym_worker_configuration(): args = _args( ( @@ -1762,3 +1675,13 @@ def test_run_miles_carries_the_island_task_record(monkeypatch, tmp_path): record = {"task_ids": ["web1"], "data_sha256": "a" * 64} rl_learner.run_miles(args, model_path=model, prompt_path=prompts, island_tasks=record) assert captured["miles_args"].yeto_rl_island_tasks == record + + +@pytest.mark.parametrize("flag, value", [ + ("--codex-reasoning-effort", "xhigh"), + ("--codex-backend-profile", "qwen35_08b"), +]) +def test_removed_codex_agent_flags_are_refused(flag, value): + args = _args((flag, value, "--use-session-server", "--tito-model", "qwen35")) + with pytest.raises(ValueError, match="signed Codex agent"): + _prepare_rl_args(args) diff --git a/tests/test_rl_launcher_codex_bundle.py b/tests/test_rl_launcher_codex_bundle.py deleted file mode 100644 index d4193c518..000000000 --- a/tests/test_rl_launcher_codex_bundle.py +++ /dev/null @@ -1,401 +0,0 @@ -"""Stock Codex run bundle on the ports launch paths (T3 G2; design R-SCOPE). - -CPU only. The contract attestation is exercised against the real bundle when -``YETO_CODEX_BUNDLE_DIR`` points at a ``scripts/fetch_codex_bundle.py`` output -(skipped otherwise); everything else uses a synthetic contract built from the -Yeto pins so the launcher wiring (learner flag, ``YETO_CODEX_*`` env, sky -file_mounts, Modal add_local_dir) is checked without a 310 MB binary. -""" - -from __future__ import annotations - -import json -import os -import shlex -import sys -import types -from pathlib import Path - -import pytest - -from yeto import launcher as L -from yeto import modal_runner as mr -from yeto.gpu_spec import parse_gpu_spec -from yeto.launcher import _prepare_rl_args, build_modal_island_config, make_miles_island_task -from yeto.rl import ( - CODEX_CONTAINER_BINARY_PATH, - CODEX_OPENENV_AGENT, - CODEX_OPENENV_IDENTITY_ENV, -) -from yeto.rl.codex_backend import QWEN35_08B_MODEL, QWEN35_08B_REVISION -from yeto.rl.ssh_harness import _plan_digest - -from test_modal_runner import fake_modal -from test_rl_engine_selection import _cli as _args # ports engine (CLI default) -from test_rl_launcher import _Resources, _Storage, _StorageMode, _Task - -PROVIDER = "yeto.rl.harness.codex.tb2_provider:provider" - - -def _codex_args(): - args = _args( - ( - "--model", QWEN35_08B_MODEL, - "--model-revision", QWEN35_08B_REVISION, - "--lora-targets", "attention", - "--custom-generate-function-path", - "miles.rollout.generate_hub.agentic_tool_call.generate", - "--custom-agent-function-path", CODEX_OPENENV_AGENT, - "--codex-reasoning-effort", "xhigh", - "--codex-backend-profile", "qwen35_08b", - "--use-session-server", - "--tito-model", "qwen35", - "--tito-allowed-append-roles", "tool", "user", - "--agent-max-seq-len", "128", - ) - ) - args.data_revision = "d" * 40 - args.source_sha256 = "e" * 64 - args.reward_sha256 = "f" * 64 - _prepare_rl_args(args) - return args - - -def _synthetic_contract(args) -> dict: - """What ``ssh_harness._codex_harness_contract`` returns for a pin-exact bundle.""" - from yeto.rl import ssh_harness as sh - from yeto.rl.codex_backend import stock_codex_backend_contract - - identity = sh._pinned_codex_adapter_identity() - return { - "agent_function_path": sh.CODEX_HARNESS_AGENT, - "agent_source_sha256": sh.CODEX_HARNESS_AGENT_SHA256, - "controller_binary_path": "/bundle/codex/codex-x86_64-unknown-linux-musl", - "controller_package_manifest_path": "/bundle/codex/codex-package.json", - "controller_app_server_schema_path": "/bundle/codex/codex_app_server_protocol.v2.schemas.json", - "bundle_binary_path": sh.CODEX_REMOTE_BUNDLE_PATH, - "bundle_package_manifest_path": sh.CODEX_REMOTE_MANIFEST_PATH, - "bundle_app_server_schema_path": sh.CODEX_REMOTE_SCHEMA_PATH, - "container_binary_path": sh.CODEX_CONTAINER_BINARY_PATH, - "container_app_server_schema_path": sh.CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH, - "binary_sha256": sh.CODEX_LINUX_BINARY_SHA256, - "binary_size_bytes": sh.CODEX_LINUX_BINARY_SIZE_BYTES, - "cli_version": sh.CODEX_CLI_VERSION, - "npm_package": sh.CODEX_NPM_PACKAGE, - "target": sh.CODEX_LINUX_TARGET, - "npm_tarball_sha256": sh.CODEX_NPM_TARBALL_SHA256, - "package_manifest_sha256": sh.CODEX_PACKAGE_MANIFEST_SHA256, - "app_server_protocol_revision": sh.CODEX_APP_SERVER_PROTOCOL_REVISION, - "app_server_schema_sha256": sh.CODEX_APP_SERVER_SCHEMA_SHA256, - **identity, - "reasoning_effort": "xhigh", - "backend": stock_codex_backend_contract("qwen35_08b", args.rollout_max_response_len), - "openenv_identity_env": dict(CODEX_OPENENV_IDENTITY_ENV), - } - - -@pytest.fixture -def bundle(tmp_path, monkeypatch): - """A fake bundle dir + the contract attestation stubbed to the pin values.""" - # An empty HOME keeps machine state (e.g. a cached HF token, which the - # launcher mounts when present) out of the mount lists these tests pin. - monkeypatch.setenv("HOME", str(tmp_path / "home")) - monkeypatch.setitem(sys.modules, "sky", types.SimpleNamespace( - Task=_Task, Resources=_Resources, Storage=_Storage, StorageMode=_StorageMode)) - root = tmp_path / "codex-bundle" / "codex" - root.mkdir(parents=True) - for name in ("codex-x86_64-unknown-linux-musl", "codex-package.json", - "codex_app_server_protocol.v2.schemas.json"): - (root / name).write_bytes(b"x") - seen = {} - - def fake_contract(namespace, args): - seen["namespace"] = vars(namespace) - return _synthetic_contract(args) - - import yeto.rl.ssh_harness as sh - - monkeypatch.setattr(sh, "_codex_harness_contract", fake_contract) - monkeypatch.setenv(L.CODEX_BUNDLE_DIR_ENV, str(root)) - monkeypatch.setenv(L.HARNESS_ENVIRONMENT_PROVIDER_ENV, PROVIDER) - monkeypatch.delenv(L.CODEX_COMPACTION_ENV, raising=False) - monkeypatch.delenv(L.HARNESS_PREFLIGHT_ENV, raising=False) - return root, seen - - -# The learner's container-side expected_env keys (yeto.rl.adapters.miles.island_entry._preflight_codex_harness) -# plus what codex_harness_agent._attest_runtime / preflight.assert_openenv_identity read. -LEARNER_EXPECTED_ENV = { - "YETO_CODEX_BINARY_PATH", "YETO_CODEX_BINARY_SHA256", "YETO_CODEX_BINARY_SIZE_BYTES", - "YETO_CODEX_VERSION", "YETO_CODEX_APP_SERVER_PROTOCOL_REVISION", - "YETO_CODEX_APP_SERVER_SCHEMA_SHA256", "YETO_CODEX_BASE_INSTRUCTIONS_SHA256", - "YETO_CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256", "YETO_CODEX_SUBMIT_TOOL_SCHEMA_SHA256", - "YETO_CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256", "YETO_CODEX_REASONING_EFFORT", - "YETO_CODEX_BACKEND_MAX_TOKENS", "YETO_CODEX_BACKEND_REASONING_EFFORT", - "YETO_CODEX_BACKEND_THINKING", "YETO_CODEX_CHAT_TEMPLATE", "YETO_CODEX_CHAT_TEMPLATE_KWARGS", - "YETO_CODEX_TITO_ALLOWED_APPEND_ROLES", "YETO_CODEX_HARNESS_CONTRACT_SHA256", -} - - -def test_non_codex_runs_are_untouched(): - args = _args() - args.data_revision = "d" * 40 - args.source_sha256 = "e" * 64 - args.reward_sha256 = "f" * 64 - assert L.codex_harness_launch(args, environ={}) is None - - -def test_codex_on_ports_fails_closed_without_bundle_provider_or_with_compaction(bundle, monkeypatch): - args = _codex_args() - with pytest.raises(ValueError, match=L.CODEX_BUNDLE_DIR_ENV): - L.codex_harness_launch(args, environ={}) - root, _ = bundle - with pytest.raises(ValueError, match=L.HARNESS_ENVIRONMENT_PROVIDER_ENV): - L.codex_harness_launch(args, environ={L.CODEX_BUNDLE_DIR_ENV: str(root)}) - with pytest.raises(ValueError, match=L.CODEX_COMPACTION_ENV): - L.codex_harness_launch(args, environ={ - L.CODEX_BUNDLE_DIR_ENV: str(root), L.HARNESS_ENVIRONMENT_PROVIDER_ENV: PROVIDER, - L.CODEX_COMPACTION_ENV: "1"}) - with pytest.raises(ValueError, match="not a directory"): - L.codex_harness_launch(args, environ={ - L.CODEX_BUNDLE_DIR_ENV: str(root / "missing"), L.HARNESS_ENVIRONMENT_PROVIDER_ENV: PROVIDER}) - - -def test_codex_launch_flags_env_and_mount_come_from_the_contract(bundle): - root, seen = bundle - args = _codex_args() - flags, envs, mounts = L.codex_harness_launch(args) - # The attestation sees the three bundle files, named as the SSH harness CLI does. - assert seen["namespace"] == { - "codex_harness_binary": str(root / "codex-x86_64-unknown-linux-musl"), - "codex_package_manifest": str(root / "codex-package.json"), - "codex_app_server_schema": str(root / "codex_app_server_protocol.v2.schemas.json"), - } - contract = _synthetic_contract(args) - parts = shlex.split(flags) - assert parts[0] == "--codex-harness-contract" and json.loads(parts[1]) == contract - assert mounts == {L.CODEX_CONTAINER_DIR: str(root.resolve())} - assert L.CODEX_CONTAINER_DIR == str(Path(CODEX_CONTAINER_BINARY_PATH).parent) - # Every env the container preflights compare is present, with contract values. - assert LEARNER_EXPECTED_ENV <= set(envs) - assert set(CODEX_OPENENV_IDENTITY_ENV) <= set(envs) - assert envs["YETO_CODEX_BINARY_PATH"] == CODEX_CONTAINER_BINARY_PATH - assert envs["YETO_CODEX_BINARY_SIZE_BYTES"] == "310730800" - assert envs["YETO_CODEX_VERSION"] == "codex-cli 0.145.0" - assert envs["YETO_CODEX_BACKEND_MAX_TOKENS"] == str(contract["backend"]["max_tokens"]) - assert envs["YETO_CODEX_CHAT_TEMPLATE_KWARGS"] == json.dumps( - contract["backend"]["chat_template_kwargs"], sort_keys=True, separators=(",", ":")) - assert envs["YETO_CODEX_HARNESS_CONTRACT_SHA256"] == _plan_digest(contract) - assert envs["YETO_CODEX_REASONING_EFFORT"] == "xhigh" - # Harness hook + provider reach the island; compaction stays unset (R-D5a). - assert envs[L.HARNESS_PREFLIGHT_ENV] == "yeto.rl.harness.codex.preflight:harness_preflight" - assert envs[L.HARNESS_ENVIRONMENT_PROVIDER_ENV] == PROVIDER - assert L.CODEX_COMPACTION_ENV not in envs - - -def test_sky_island_task_mounts_the_bundle_and_exports_the_env(bundle): - root, _ = bundle - args = _codex_args() - task = make_miles_island_task(args, parse_gpu_spec(args.gpu)[0], 0, 1, "127.0.0.1:29400") - assert "--codex-harness-contract " in task.run - assert f"--custom-agent-function-path {CODEX_OPENENV_AGENT}" in task.run - assert "--agent-max-seq-len 128" in task.run and "--use-session-server" in task.run - assert "--tito-model qwen35" in task.run and "--codex-backend-profile qwen35_08b" in task.run - assert "--partial-rollout" not in task.run - assert task.file_mounts[L.CODEX_CONTAINER_DIR] == str(root.resolve()) - assert task.envs["YETO_CODEX_BINARY_PATH"] == CODEX_CONTAINER_BINARY_PATH - assert task.envs[L.HARNESS_PREFLIGHT_ENV].endswith(":harness_preflight") - assert task.envs[L.HARNESS_ENVIRONMENT_PROVIDER_ENV] == PROVIDER - assert "YETO_CODEX_COMPACTION_ENABLED" not in task.envs - - -def test_modal_island_mounts_the_bundle_through_the_image(bundle, monkeypatch): - root, _ = bundle - args = _codex_args() - args.gpu = "modal:1xh100" - args.rl_image = "docker:ghcr.io/x/miles@sha256:" + "c" * 64 - spec = parse_gpu_spec(args.gpu)[0] - task = make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400") - cfg = build_modal_island_config(args, spec, 0, task, "1.2.3.4:29400") - assert cfg.codex_dir == str(root.resolve()) and cfg.codex_mount == L.CODEX_CONTAINER_DIR - assert cfg.envs["YETO_CODEX_BINARY_PATH"] == CODEX_CONTAINER_BINARY_PATH - state = fake_modal(monkeypatch) - monkeypatch.setattr(mr, "registry_credentials", lambda *_a, **_k: None) - mr.ModalOps("yeto-run").define(cfg) - (img,) = state["images"] - names = [c[0] for c in img.calls] - assert names == ["from_registry", "env", "add_local_dir", "add_local_dir"] - assert img.calls[-1][1] == (cfg.codex_dir, L.CODEX_CONTAINER_DIR) and img.calls[-1][2] == {"copy": False} - - -def test_modal_config_requires_an_existing_bundle_dir(tmp_path): - with pytest.raises(ValueError, match="go together"): - mr.ModalIslandConfig(app_name="a", learner_id=0, training_mode="sft", gpu="H100", - gpus_per_node=8, num_nodes=1, run_script="x", codex_dir=str(tmp_path)).validate() - with pytest.raises(ValueError, match="not a directory"): - mr.ModalIslandConfig(app_name="a", learner_id=0, training_mode="sft", gpu="H100", - gpus_per_node=8, num_nodes=1, run_script="x", - codex_dir=str(tmp_path / "nope"), codex_mount="/opt/yeto/codex").validate() - - -@pytest.mark.skipif( - not os.environ.get(L.CODEX_BUNDLE_DIR_ENV) or not os.path.isdir(os.environ.get(L.CODEX_BUNDLE_DIR_ENV, "")), - reason="YETO_CODEX_BUNDLE_DIR not set to a fetched bundle (scripts/fetch_codex_bundle.py)", -) -def test_real_bundle_attests_against_the_pins(): - args = _codex_args() - try: - contract = L.codex_bundle_contract(args, os.environ[L.CODEX_BUNDLE_DIR_ENV]) - except ValueError as exc: - if "app-server schema" in str(exc): - pytest.xfail(f"pinned app-server schema file not reproducible from the npm artifact: {exc}") - raise - assert contract["binary_size_bytes"] == 310_730_800 - assert contract["openenv_identity_env"] == CODEX_OPENENV_IDENTITY_ENV - - -def test_modal_island_mounts_local_data_and_passes_tb2_env(bundle, monkeypatch, tmp_path): - root, _ = bundle - data = tmp_path / "smoke.jsonl" - data.write_text('{"prompt": "x", "metadata": {"task_id": "fix-git"}}\n') - monkeypatch.setenv("MODAL_TOKEN_ID", "ak-test") - monkeypatch.setenv("MODAL_TOKEN_SECRET", "as-test") - monkeypatch.setenv("YETO_SANDBOX_MODAL_TOKEN_ID", "sb-id") - monkeypatch.setenv("YETO_SANDBOX_MODAL_TOKEN_SECRET", "sb-secret") - monkeypatch.setenv("YETO_HARNESS_TB2_TASKS_DIR", "/opt/yeto/codex/tb2-tasks") - monkeypatch.setenv("YETO_HARNESS_TB2_FAULT", "create_fail:2") - args = _codex_args() - args.data = str(data) - args.gpu = "modal:1xl40s" - args.rl_image = "docker:ghcr.io/x/miles@sha256:" + "c" * 64 - spec = parse_gpu_spec(args.gpu)[0] - task = make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400") - cfg = build_modal_island_config(args, spec, 0, task, "1.2.3.4:29400") - assert cfg.extra_mounts == {"/root/yeto-data.jsonl": str(data)} - for name in ("YETO_SANDBOX_MODAL_TOKEN_ID", "YETO_SANDBOX_MODAL_TOKEN_SECRET", - "YETO_HARNESS_TB2_TASKS_DIR", "YETO_HARNESS_TB2_FAULT"): - assert cfg.envs[name] == os.environ[name] - # secret-handling-hardening: the main Modal token never reaches the island. - assert not {"MODAL_TOKEN_ID", "MODAL_TOKEN_SECRET"} & set(cfg.envs) - assert not {"YETO_SANDBOX_MODAL_TOKEN_ID", "YETO_SANDBOX_MODAL_TOKEN_SECRET"} & set(task.envs) - assert mr.ModalIslandConfig.from_json(cfg.to_json()).extra_mounts == cfg.extra_mounts - state = fake_modal(monkeypatch) - monkeypatch.setattr(mr, "registry_credentials", lambda *_a, **_k: None) - mr.ModalOps("yeto-run").define(cfg) - (img,) = state["images"] - assert img.calls[-1][0] == "add_local_file" - assert img.calls[-1][1] == (str(data), "/root/yeto-data.jsonl") and img.calls[-1][2] == {"copy": False} - - -def test_learner_command_forwards_the_codex_reasoning_effort(bundle): - args = _codex_args() - spec = parse_gpu_spec(args.gpu)[0] - task = make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400") - assert " --codex-reasoning-effort xhigh" in task.run - assert " --codex-backend-profile qwen35_08b" in task.run - - -def test_island_setup_installs_the_modal_client_only_for_the_modal_sandbox_provider(bundle, monkeypatch): - """A-T3-5 (codex-smoke-20261003-7): rollout workers create task sandboxes with - the Modal client, which the ports image does not ship; `import modal` there - resolved Modal's dependency-less runtime mount and failed on every acquire.""" - args = _codex_args() - spec = parse_gpu_spec(args.gpu)[0] - plain = make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400") - assert L.MODAL_CLIENT_SETUP not in plain.setup # local/other providers: untouched - - monkeypatch.setenv(L.HARNESS_ENVIRONMENT_PROVIDER_ENV, L.MODAL_SANDBOX_PROVIDER) - with pytest.raises(ValueError, match="YETO_SANDBOX_MODAL_TOKEN_ID"): - make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400") - monkeypatch.setenv("YETO_SANDBOX_MODAL_TOKEN_ID", "sb-id") - monkeypatch.setenv("YETO_SANDBOX_MODAL_TOKEN_SECRET", "sb-secret") - task = make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400") - assert task.envs[L.HARNESS_ENVIRONMENT_PROVIDER_ENV] == L.MODAL_SANDBOX_PROVIDER - assert L.MODAL_CLIENT_SETUP in task.setup - assert "--no-deps --ignore-installed" in L.MODAL_CLIENT_SETUP and "modal==1.5.5" in L.MODAL_CLIENT_SETUP - assert 'sysconfig.get_paths()["purelib"]' in L.MODAL_CLIENT_SETUP # must land in the island python itself - # the check runs with the run script's PYTHONPATH, where Modal's own runtime mount is not enough - assert "$HOME/miles:$HOME/sglang/python:$HOME/sky_workdir" in L.MODAL_CLIENT_SETUP - assert "import modal, grpclib, synchronicity" in L.MODAL_CLIENT_SETUP - # the client install precedes the prefetch and fails the setup when unusable - assert task.setup.index(L.MODAL_CLIENT_SETUP) > task.setup.index("pip install -q --no-deps -e ~/miles") - assert "exit 1" in L.MODAL_CLIENT_SETUP - cfg = build_modal_island_config(args, spec, 0, task, "1.2.3.4:29400") - assert L.MODAL_CLIENT_SETUP in cfg.setup_script - - -def test_modal_sandbox_secret_attaches_a_named_secret_and_checks_in_the_island(bundle, monkeypatch): - """--modal-sandbox-secret: the sandbox token comes from a named Modal Secret - (never through the launcher); the island setup fails closed without it.""" - for name in L.SANDBOX_MODAL_TOKEN_ENVS: - monkeypatch.delenv(name, raising=False) - monkeypatch.setenv("MODAL_TOKEN_ID", "ak-test") - monkeypatch.setenv("MODAL_TOKEN_SECRET", "as-test") - monkeypatch.setenv(L.HARNESS_ENVIRONMENT_PROVIDER_ENV, L.MODAL_SANDBOX_PROVIDER) - args = _codex_args() - args.gpu = "modal:1xl40s" - args.rl_image = "docker:ghcr.io/x/miles@sha256:" + "c" * 64 - spec = parse_gpu_spec(args.gpu)[0] - with pytest.raises(ValueError, match="YETO_SANDBOX_MODAL_TOKEN_ID"): - make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400") # unchanged without the flag - args.modal_sandbox_secret = "yeto-sandbox-modal" - task = make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400") - assert L.SANDBOX_TOKEN_IN_ISLAND_CHECK in task.setup - assert task.setup.index(L.SANDBOX_TOKEN_IN_ISLAND_CHECK) < task.setup.index(L.MODAL_CLIENT_SETUP) - cfg = build_modal_island_config(args, spec, 0, task, "1.2.3.4:29400") - assert cfg.named_secrets == ("yeto-sandbox-modal",) - assert not set(L.SANDBOX_MODAL_TOKEN_ENVS) & set(cfg.envs) - assert L.SANDBOX_TOKEN_IN_ISLAND_CHECK in cfg.setup_script - assert mr.ModalIslandConfig.from_json(cfg.to_json()).named_secrets == ("yeto-sandbox-modal",) - state = fake_modal(monkeypatch) - monkeypatch.setattr(mr, "registry_credentials", lambda *_a, **_k: None) - mr.ModalOps("yeto-run").define(cfg) - (_fn, kwargs), = state["functions"].values() - assert ("named-secret", "yeto-sandbox-modal") in kwargs["secrets"] - - -def test_modal_sandbox_secret_in_island_check_fails_closed_without_values(tmp_path): - import subprocess - - env = {"PATH": os.environ["PATH"], "YETO_SANDBOX_MODAL_TOKEN_ID": "ak-x"} - r = subprocess.run(["bash", "-c", L.SANDBOX_TOKEN_IN_ISLAND_CHECK], env=env, - capture_output=True, text=True) - assert r.returncode == 1 and "YETO_SANDBOX_MODAL_TOKEN_SECRET missing" in r.stderr - assert "ak-x" not in r.stderr + r.stdout - env["YETO_SANDBOX_MODAL_TOKEN_SECRET"] = "as-y" - r = subprocess.run(["bash", "-c", L.SANDBOX_TOKEN_IN_ISLAND_CHECK], env=env, - capture_output=True, text=True) - assert r.returncode == 0 and r.stdout + r.stderr == "" - - -def test_modal_sandbox_secret_is_refused_for_non_modal_islands(bundle, monkeypatch): - args = _codex_args() - args.modal_sandbox_secret = "yeto-sandbox-modal" - spec = parse_gpu_spec(args.gpu)[0] - assert spec.cloud != "modal" - with pytest.raises(ValueError, match="--modal-sandbox-secret"): - make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400") - with pytest.raises(ValueError, match="cannot use it"): - L.check_cloud_prerequisites([spec], project_ids={}, args=args, modal_ok=True) - args.modal_sandbox_secret = "bad name;rm" - with pytest.raises(ValueError, match="not a Modal Secret name"): - L.modal_sandbox_secret_name(args, parse_gpu_spec("modal:1xl40s")[0]) - - -def test_modal_island_ships_real_secret_values_not_sky_masks(bundle, monkeypatch, tmp_path): - """sky.Task keeps secrets as pydantic SecretStr ("**********"); the Modal - island must get the plain value (s19-compaction-g1-20261010a).""" - pydantic = pytest.importorskip("pydantic") - masked = pydantic.SecretStr("k" * 48) # what sky>=0.10 Task.secrets holds - assert str(masked) != "k" * 48 - task = types.SimpleNamespace(envs={"A": "1"}, secrets={"HF_TOKEN": masked}, - setup="", run="echo") - monkeypatch.setenv("MODAL_TOKEN_ID", "ak-test") - monkeypatch.setenv("MODAL_TOKEN_SECRET", "as-test") - args = _codex_args() - args.gpu = "modal:1xl40s" - args.rl_image = "docker:ghcr.io/x/miles@sha256:" + "c" * 64 - spec = parse_gpu_spec(args.gpu)[0] - cfg = build_modal_island_config(args, spec, 0, task, "1.2.3.4:29400") - assert cfg.envs["HF_TOKEN"] == "k" * 48 - assert L._secret_value("plain") == "plain" diff --git a/tests/test_rl_launcher_dense_full.py b/tests/test_rl_launcher_dense_full.py index 1ce3a062a..a246dac3b 100644 --- a/tests/test_rl_launcher_dense_full.py +++ b/tests/test_rl_launcher_dense_full.py @@ -132,27 +132,6 @@ def test_island_entry_accepts_the_dense_island_command(): island_entry._check_dense_full_sweep(parsed, yeto_policy_sync=True) -@pytest.mark.parametrize( - ("reward_function", "kind"), - ( - ("yeto.rl.harness.codex.reward:reward_func", "secrlenv"), - ("yeto.rl.gsm8k_reward:score", "reward-function"), - ("yeto.rl.math_reward:reward_func", "reward-function"), - ), -) -def test_dense_island_evidence_kind_follows_its_reward_function(reward_function, kind): - from types import SimpleNamespace - - from yeto.rl.adapters.miles import island_entry - from yeto.rl.trajectory_evidence import EVIDENCE_KINDS - - chosen = island_entry._trajectory_evidence_kind( - SimpleNamespace(reward_function=reward_function) - ) - assert chosen == kind - assert chosen in EVIDENCE_KINDS - - EXPERT_FULL = ( "--rl-sync-preset", "full-parameter", "--rl-placement", "fixed-partition", "--rl-rollout-gpus", "8", "--rl-parameter-layout-sha256", LAYOUT, "--fragments", "4", diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index 503fe2e28..825e2fd0e 100644 --- a/tests/test_rl_miles_adapter_config.py +++ b/tests/test_rl_miles_adapter_config.py @@ -647,31 +647,18 @@ def test_lr_schedule_flags_are_adapter_owned(): mc.check_extra_argv((flag, "1")) -# ---------------------------------------------------------------- rl-fn-codex-rollout 1.0: append roles by tokenizer family +# ---------------------------------------------------------------- append roles are decided by the --tito-model template -@pytest.mark.parametrize("tito_model,ok", [("qwen35", True), ("qwen4exp", True), ("qwen38small", True), ("deepseekv4", True), - ("nosuchfamily", False)]) -def test_stock_codex_append_roles_resolve_by_tito_family(tito_model, ok): - """``--tito-model`` is a tokenizer family (``qwen4exp`` for Flash-Next), not a - profile name; the s15 4-layer smoke failed with UnmappedConfigError here.""" +def test_explicit_append_roles_are_refused(): + """Since S26 E0.5 no signed Codex agent may state the append roles.""" from types import SimpleNamespace - from yeto.rl import CODEX_OPENENV_AGENT from yeto.rl.adapters.miles.config import LEAF_POLICY check = LEAF_POLICY["agent.tito_allowed_append_roles"] - agent = SimpleNamespace(custom_agent_function_path=CODEX_OPENENV_AGENT, tito_model=tito_model) - cfg = SimpleNamespace(agent=agent) + cfg = SimpleNamespace(agent=SimpleNamespace(custom_agent_function_path="x:y", tito_model="qwen35")) assert check(None, cfg) is None - if ok: - assert check(("tool", "user"), cfg) is None - assert "fix the append roles" in check(("tool",), cfg) - else: - assert "upstream Miles has no --tito-allowed-append-roles" in check(("tool", "user"), cfg) - # an unsigned agent is still refused - cfg.agent.custom_agent_function_path = "x:y" - assert check(("tool", "user"), cfg) is not None - + assert "upstream Miles has no --tito-allowed-append-roles" in check(("tool", "user"), cfg) # --- S17 N16: --rl-lr-schedule ---------------------------------------------------- diff --git a/tests/test_rl_miles_carry_over.py b/tests/test_rl_miles_carry_over.py index e94f3d267..a83ad2a90 100644 --- a/tests/test_rl_miles_carry_over.py +++ b/tests/test_rl_miles_carry_over.py @@ -249,7 +249,7 @@ def test_miles_stage_two_support_and_single_turn_check(): check_task(1) check_task(0, custom_agent="x.y", custom_generate="x.z") with pytest.raises(PolicyAgeError, match="stage 3"): - check_task(1, custom_generate="yeto.rl.harness.codex.codex_openenv_generate.generate") + check_task(1, custom_generate="yeto.rl.adapters.miles.records_generate.generate") with pytest.raises(PolicyAgeError, match="recompute"): check_task(1, recompute_prefill=True) @@ -259,7 +259,7 @@ def test_launcher_refuses_agentic_carry_over_before_launch(): from yeto.launcher import _prepare_rl_args args = _args(["--rl-max-policy-age", "1", "--custom-generate-function-path", - "yeto.rl.harness.codex.codex_openenv_generate.generate"]) + "yeto.rl.adapters.miles.records_generate.generate"]) with pytest.raises(PolicyAgeError, match="stage 3"): _prepare_rl_args(args) diff --git a/tests/test_rl_neutral_rewards.py b/tests/test_rl_neutral_rewards.py index 748a644a7..b150750db 100644 --- a/tests/test_rl_neutral_rewards.py +++ b/tests/test_rl_neutral_rewards.py @@ -139,73 +139,6 @@ def test_generic_miles_wrapper_equals_original_entry_point(name, entry): assert a.snapshot() == b.snapshot(), (name, i) -# -- 3.3 codex: abort as a reward result -------------------------------------------- - -def _outcome_metadata(status: str, reward: float) -> dict: - from yeto.rl.harness.codex import reward as codex_reward - - outcome = {"schema": 1, "status": status, "episode_id": "ep", "task_id": "t", "reward": reward, - "passed": reward == 1.0, "class": None, - codex_reward.INFRASTRUCTURE_503_RETRIES_KEY: 0} - return {codex_reward.OUTCOME_KEY: outcome} - - -@pytest.fixture -def fake_miles_sample(monkeypatch): - names = ("miles", "miles.utils", "miles.utils.types") - for n in names[:-1]: - monkeypatch.setitem(sys.modules, n, types.ModuleType(n)) - mod = types.ModuleType(names[-1]) - mod.Sample = cases.FakeSample - monkeypatch.setitem(sys.modules, names[-1], mod) - - -def test_codex_infrastructure_outcome_aborts_in_neutral_and_miles(monkeypatch, fake_miles_sample): - from yeto.rl.harness.codex import reward as codex_reward - - metadata = _outcome_metadata("infrastructure_error", 0.0) - result = codex_reward.secrlenv_reward(Trajectory(metadata=metadata)) - assert (result.value, result.aborted, result.reason) == (0.0, True, "secrlenv_infrastructure_failure") - sample = cases.FakeSample(metadata=dict(metadata)) - assert _run(codex_reward.reward_func(None, sample)) == 0.0 - # The SecRLEnv entry marks FAILED, not ABORTED: Miles drops an aborted - # group before the SecRLEnv group check could schedule its retry. - assert sample.status is cases.Status.FAILED - generic = cases.FakeSample(metadata=dict(metadata)) - assert _run(miles_rewards.miles_reward(codex_reward.secrlenv_reward)(None, generic)) == 0.0 - assert generic.status is cases.Status.ABORTED - - -def test_codex_verdict_and_invalid_outcome(monkeypatch, fake_miles_sample): - from yeto.rl.harness.codex import reward as codex_reward - - metadata = _outcome_metadata("completed", 1.0) - result = codex_reward.secrlenv_reward(Trajectory(metadata=metadata)) - assert (result.value, result.aborted) == (1.0, False) - sample = cases.FakeSample(metadata=dict(metadata)) - assert _run(codex_reward.reward_func(None, [sample])) == [1.0] - assert sample.status is cases.Status.COMPLETED - legacy = {**metadata, "secrlenv_trusted_outcome_hmac": "0" * 64} - assert codex_reward.secrlenv_reward(Trajectory(metadata=legacy)) == result - tampered = {codex_reward.OUTCOME_KEY: {**metadata[codex_reward.OUTCOME_KEY], "reward": 1.5}} - with pytest.raises(codex_reward.UntrustedOutcome): - codex_reward.secrlenv_reward(Trajectory(metadata=tampered)) - with pytest.raises(codex_reward.UntrustedOutcome): - _run(codex_reward.reward_func(None, cases.FakeSample(metadata=tampered))) - - -def test_tbench_infrastructure_marker_aborts(monkeypatch): - from yeto.rl.harness.codex import tbench_reward - - metadata = {tbench_reward.INFRASTRUCTURE_KEY: "x", "success": True} - result = tbench_reward.tbench_reward(Trajectory(metadata=metadata)) - assert result.aborted and result.value == 0.0 and "success" not in metadata - sample = SimpleNamespace(metadata={tbench_reward.INFRASTRUCTURE_KEY: "x", "success": True}, - status=None) - assert _run(tbench_reward.reward_func(None, sample)) == 0.0 - assert sample.status == "ABORTED" and "success" not in sample.metadata - - # -- 3.10 registry and user rewards ------------------------------------------------ def test_custom_reward_via_miles_equals_direct_call(): diff --git a/tests/test_rl_policy_versions.py b/tests/test_rl_policy_versions.py index f600d50dc..b994c6cbb 100644 --- a/tests/test_rl_policy_versions.py +++ b/tests/test_rl_policy_versions.py @@ -191,19 +191,3 @@ def test_fully_async_trained_groups_hook_reports_the_round(monkeypatch): assert len(reported) == 1 # the all-samples hook reports otherwise -def test_rollout_side_matches_plain_versions_against_the_published_label(): - from yeto.rl.adapters.miles import carry_over - from yeto.rl.adapters.miles.rollout_meta_hook import group_matches_policy - from yeto.rl.harness.codex.codex_openenv_generate import window_problem - - def sample(version): - span = SimpleNamespace(version=version, abs_start=1, abs_end=2) - return SimpleNamespace(status=SimpleNamespace(value="completed"), tokens=[0, 1], - response_length=1, weight_versions=[SimpleNamespace(spans=[span])]) - - label = f"yeto:6:{H6}" - assert group_matches_policy([sample("6"), sample("6")], label, 2) - assert not group_matches_policy([sample("6"), sample("5")], label, 2) - assert carry_over.response_token_versions(sample("5")) == [5] - assert window_problem(["5", "6"], label, 1) is None - assert "behind" in window_problem(["4"], label, 1) diff --git a/tests/test_rl_reward_success.py b/tests/test_rl_reward_success.py index 84ab6b78d..13b556dc2 100644 --- a/tests/test_rl_reward_success.py +++ b/tests/test_rl_reward_success.py @@ -2,7 +2,6 @@ import asyncio import importlib.util -import re import sys import types from pathlib import Path @@ -110,24 +109,3 @@ def _tb_sample(monkeypatch, reward, status="completed", verifier=None, rc=None): return SimpleNamespace(metadata=meta, status=None) -def test_tbench_success_is_outcome_pass_bit(monkeypatch): - from yeto.rl import tbench_outcome - from yeto.rl.harness.codex import tbench_reward - - ok = _tb_sample(monkeypatch, 1.0, verifier=tbench_outcome.TEST_SH_VERIFIER, rc=0) - bad = _tb_sample(monkeypatch, 0.0) - timeout = _tb_sample(monkeypatch, 0.0, status="timeout", verifier=tbench_outcome.TIMEOUT_VERIFIER) - bad.metadata["success"] = True # untrusted pre-set flag is overwritten - assert _run(tbench_reward.reward_func(None, [ok, bad, timeout])) == [1.0, 0.0, 0.0] - assert [s.metadata["success"] for s in (ok, bad, timeout)] == [True, False, False] - assert all(_fork_flag(s.metadata) in (0, 1) for s in (ok, bad, timeout)) - # siblings / group filter still accept the extra key - tbench_reward.check_siblings([ok, bad]) - - -def test_tbench_infrastructure_sample_has_no_success(monkeypatch): - from yeto.rl.harness.codex import tbench_reward - - sample = SimpleNamespace(metadata={tbench_reward.INFRASTRUCTURE_KEY: "boom", "success": True}, status=None) - assert _run(tbench_reward.reward_func(None, sample)) == 0.0 - assert sample.status == "ABORTED" and "success" not in sample.metadata diff --git a/tests/test_rl_sao_streaming_runtime.py b/tests/test_rl_sao_streaming_runtime.py index ced8e08eb..737c786cd 100644 --- a/tests/test_rl_sao_streaming_runtime.py +++ b/tests/test_rl_sao_streaming_runtime.py @@ -595,7 +595,7 @@ def test_context_schema_label_2_still_writes_trajectory_evidence_schema_4(tmp_pa persist_published_round_evidence, ) from yeto.rl.engine.policy_versions import PublishedPolicies - from yeto.rl.harness.codex import reward + from yeto.rl import secrlenv_outcome as reward from yeto.rl.trajectory_evidence import ( read_trajectory_batch_evidence, trajectory_batch_evidence_path, diff --git a/tests/test_rl_secrlenv_ports_retry.py b/tests/test_rl_secrlenv_ports_retry.py deleted file mode 100644 index f64ff17bf..000000000 --- a/tests/test_rl_secrlenv_ports_retry.py +++ /dev/null @@ -1,509 +0,0 @@ -"""P0-1.7b.4-Q1: the SecRLEnv group check and same-task retry on ports. - -CPU only: the algorithm spec accepts the group check, the Miles translation -binds the retry classes, the retry data source keeps the legacy contract, the -terminal cleanup sanitizes errors, and the Miles classes load against stub -Miles modules. -""" - -from __future__ import annotations - -import asyncio -import importlib -import sys -import types -from types import SimpleNamespace - -import pytest - -from test_rl_miles_adapter_config import flag_value, make_config, sub -from yeto.rl import ( - SECRLENV_AGENT, - SECRLENV_GENERATE, - SECRLENV_GROUP_FILTER, -) -from yeto.rl.adapters.miles import config as cfg -from yeto.rl.adapters.miles import secrlenv_retry as retry -from yeto.rl.adapters.miles.binding import check_filter -from yeto.rl.adapters.miles.entry import miles_capabilities -from yeto.rl.engine.algorithm import SECRLENV_GROUP_CHECK_FILTER, AlgorithmSpec -from yeto.rl.harness.codex.reward import ( - SecRLEnvReplacementExhausted, - SecRLEnvRolloutCleanupError, - SecRLEnvUntrustedEvidence, - UntrustedOutcome, -) - -TARGET = 2 - - -# ------------------------------------------------------------------ spec -def _secrlenv_spec() -> AlgorithmSpec: - return AlgorithmSpec( - advantage_estimator="grpo", - dynamic_sampling_filter=SECRLENV_GROUP_FILTER, - dynamic_sampling_max_replacements=0, - ) - - -def test_spec_accepts_the_group_check_under_its_neutral_name(): - spec = _secrlenv_spec() - assert spec.dynamic_sampling_filter == SECRLENV_GROUP_CHECK_FILTER - assert spec.sha256() == AlgorithmSpec( - advantage_estimator="grpo", - dynamic_sampling_filter=SECRLENV_GROUP_CHECK_FILTER, - dynamic_sampling_max_replacements=0, - ).sha256() - legacy_args = SimpleNamespace(dynamic_sampling_filter_path=SECRLENV_GROUP_FILTER, - dynamic_sampling_max_replacements=0, kl_coef=None) - assert AlgorithmSpec.from_legacy_args(legacy_args).sha256() == spec.sha256() - capabilities = miles_capabilities("sha256:" + "0" * 64) - assert SECRLENV_GROUP_CHECK_FILTER in capabilities.dynamic_sampling_filters - assert ("dynamic_sampling_filters", SECRLENV_GROUP_CHECK_FILTER) in spec.required_mechanisms() - - -def test_binding_check_skips_only_the_evidence_call(): - row = check_filter(SECRLENV_GROUP_CHECK_FILTER) - assert (row["mapped"], row["argv"]) == (SECRLENV_GROUP_FILTER, "same") - assert row["identity"]["source_sha256"] - assert row["call"] == "skipped: needs environment evidence" - - -# ----------------------------------------------------------- translation -def _secrlenv_config(**batch): - config = sub(make_config(), "agent", - custom_generate_function_path=SECRLENV_GENERATE, - custom_agent_function_path=SECRLENV_AGENT, - dynamic_sampling_filter_path=SECRLENV_GROUP_FILTER) - groups = make_config().batch.groups_per_round - return sub(config, "batch", **{"over_sampling_batch_size": groups + 1, **batch}) - - -def test_translation_binds_the_retry_classes(): - argv = list(cfg.translate_run_config(_secrlenv_config(), _secrlenv_spec()).argv) - assert flag_value(argv, "--dynamic-sampling-filter-path") == SECRLENV_GROUP_FILTER - assert flag_value(argv, "--data-source-path") == retry.DATA_SOURCE_PATH - assert flag_value(argv, "--rollout-function-path") == retry.ROLLOUT_FUNCTION_PATH - assert flag_value(argv, "--custom-generate-function-path") == SECRLENV_GENERATE - assert flag_value(argv, "--custom-agent-function-path") == SECRLENV_AGENT - - -def test_translation_leaves_other_runs_unchanged(): - argv = cfg.translate_run_config(make_config(), AlgorithmSpec()).argv - assert "--data-source-path" not in argv and "--rollout-function-path" not in argv - - -def test_retry_classes_are_adapter_owned(): - assert {"--data-source-path", "--rollout-function-path"} <= cfg.ADAPTER_OWNED_FLAGS - - -@pytest.mark.parametrize("config, spec, match", [ - (sub(_secrlenv_config(), "agent", dynamic_sampling_filter_path=None), AlgorithmSpec(), - "go together"), - (sub(make_config(), "agent", dynamic_sampling_filter_path=SECRLENV_GROUP_FILTER), - _secrlenv_spec(), "go together"), - (_secrlenv_config(over_sampling_batch_size=make_config().batch.groups_per_round), - _secrlenv_spec(), "exactly one refill slot"), -]) -def test_translation_refuses_a_broken_retry_contract(config, spec, match): - with pytest.raises(cfg.UnmappedConfigError, match=match): - cfg.translate_run_config(config, spec) - - -def test_retry_refuses_a_policy_age_limit(): - with pytest.raises(ValueError, match="policy-age limit"): - retry.rollout_argv(custom_agent=SECRLENV_AGENT, - dynamic_filter=SECRLENV_GROUP_CHECK_FILTER, - rollout_batch_size=TARGET, over_sampling_batch_size=TARGET + 1, - max_policy_age=1) - - -# ------------------------------------------------------- retry data source -class _Sample: - def __init__(self, index, status="completed"): - self.index = index - self.status = status - self.metadata = {"secrlenv_trusted_outcome": {"index": index}} - self.resets = 0 - - def reset_for_retry(self): - self.status = "pending" - self.resets += 1 - - -class _Source: - """Stands in for Miles' RolloutDataSourceWithBuffer.""" - - def __init__(self, args): - self.args = args - self.next_index = 0 - self.buffer = [] - self.draws = [] - - def get_samples(self, num_samples): - self.draws.append(num_samples) - groups = [[_Sample(self.next_index + g * 10 + s) for s in range(2)] - for g in range(num_samples)] - self.next_index += 100 - return groups - - def add_samples(self, samples): - self.buffer.extend(samples) - - -class _RetrySource(retry.SecRLEnvTaskRetry, _Source): - pass - - -def _args(**overrides): - values = dict(rollout_batch_size=TARGET, over_sampling_batch_size=TARGET + 1, - yeto_rl_secrlenv_max_infrastructure_replacements=1) - values.update(overrides) - return SimpleNamespace(**values) - - -def _ledger(group, value=0): - return {sample.index: value for sample in group} - - -def _request(source, group, *, key="evidence", reason="secrlenv_infrastructure_failure", - ledger=None): - source.request_retry(group, evidence_key=key, reason=reason, - infrastructure_503_retries=_ledger(group) if ledger is None else ledger) - - -@pytest.mark.parametrize("overrides", [ - dict(over_sampling_batch_size=TARGET), - dict(rollout_batch_size=0, over_sampling_batch_size=1), - dict(yeto_rl_secrlenv_max_infrastructure_replacements=2), - dict(yeto_rl_secrlenv_max_infrastructure_replacements=None), -]) -def test_data_source_refuses_a_broken_contract(overrides): - with pytest.raises(SecRLEnvReplacementExhausted, match="contract is invalid"): - _RetrySource(_args(**overrides)) - - -def test_round_draws_the_exact_batch_then_only_the_queued_retry(): - args = _args() - source = _RetrySource(args) - assert getattr(args, retry.RETRY_CALLBACK_ATTR) == source.request_retry - groups = source.get_samples(TARGET + 1) - assert len(groups) == TARGET and source.draws == [TARGET] - failed = groups[1] - failed[0].status = "failed" - failed[0].metadata["secrlenv_trusted_outcome"] = {"status": "infrastructure_error"} - _request(source, failed, ledger={failed[0].index: 1, failed[1].index: 0}) - _request(source, failed, ledger={failed[0].index: 1, failed[1].index: 0}) # repeated callback - (rerun,) = source.get_samples(TARGET + 1) - assert [s.index for s in rerun] == [s.index for s in failed] - assert [s.resets for s in rerun] == [1, 1] - assert [s.metadata for s in rerun] == [ - {"secrlenv_infrastructure_503_retries": 1}, - {"secrlenv_infrastructure_503_retries": 0}, - ] - assert source.draws == [TARGET] # no new task was drawn for the retry - assert source.retry_stats() == dict( - replacement_attempts=1, nontrainable_groups=1, infrastructure_groups=1, - pending_groups=0, failed_groups=1) - assert retry.retry_metrics(args._yeto_secrlenv_retry_stats)[ - "rl/secrlenv/replacement_attempts"] == 1.0 - - -def test_refill_without_a_queued_retry_is_terminal(): - source = _RetrySource(_args()) - source.get_samples(TARGET + 1) - with pytest.raises(SecRLEnvReplacementExhausted, match="no authenticated task retry"): - source.get_samples(TARGET + 1) - - -def test_refill_size_is_part_of_the_contract(): - source = _RetrySource(_args()) - with pytest.raises(SecRLEnvReplacementExhausted, match="refill contract"): - source.get_samples(TARGET) - - -def test_second_retry_in_a_round_is_terminal(): - source = _RetrySource(_args()) - groups = source.get_samples(TARGET + 1) - _request(source, groups[0], key="a") - source.get_samples(TARGET + 1) - with pytest.raises(SecRLEnvReplacementExhausted, match="budget was exhausted"): - _request(source, groups[1], key="b") - - -@pytest.mark.parametrize("ledger, error, match", [ - ({}, SecRLEnvUntrustedEvidence, "ledger authentication"), - ("bad", SecRLEnvUntrustedEvidence, "ledger authentication"), - (None, SecRLEnvReplacementExhausted, "identity did not match"), -]) -def test_retry_must_name_a_drawn_task_and_its_ledger(ledger, error, match): - source = _RetrySource(_args()) - groups = source.get_samples(TARGET + 1) - group = groups[0] if ledger is not None else [_Sample(9999), _Sample(9998)] - with pytest.raises(error, match=match): - _request(source, group, ledger=ledger) - - -def test_ledger_values_are_bounded(): - source = _RetrySource(_args()) - group = source.get_samples(TARGET + 1)[0] - with pytest.raises(SecRLEnvUntrustedEvidence): - _request(source, group, ledger=_ledger(group, value=2)) - - -def test_round_end_resets_the_budget_and_the_filter_memo(): - args = _args() - source = _RetrySource(args) - groups = source.get_samples(TARGET + 1) - _request(source, groups[0]) - source.get_samples(TARGET + 1) - args._yeto_secrlenv_filter_state = {"rollout_id": None, "decisions": {"k": (False, "x")}} - source.add_samples([]) # Miles returns the round's aborted groups - next_round = source.get_samples(TARGET + 1) - assert len(next_round) == TARGET and source.draws == [TARGET, TARGET] - assert args._yeto_secrlenv_filter_state is None - assert source.retry_stats()["replacement_attempts"] == 0 - with pytest.raises(SecRLEnvReplacementExhausted, match="identity did not match"): - _request(source, groups[1], key="old-round") - - -def test_each_round_has_a_fresh_budget(): - source = _RetrySource(_args()) - _request(source, source.get_samples(TARGET + 1)[0]) - source.get_samples(TARGET + 1) - source.add_samples([]) - _request(source, source.get_samples(TARGET + 1)[0]) - assert len(source.get_samples(TARGET + 1)) == 1 - - -def test_group_check_queues_the_retry_through_the_data_source(monkeypatch): - """The real group check, fed a signed infrastructure outcome, asks the - data source for the rerun (the callback attribute is the contract).""" - - from yeto.rl.harness.codex import reward - - monkeypatch.setattr(reward, "_verified_outcome", lambda metadata: (metadata["outcome"], 0.0)) - output_module = types.ModuleType("miles.rollout.filter_hub.base_types") - output_module.DynamicFilterOutput = lambda keep, reason: SimpleNamespace(keep=keep, reason=reason) - for name in ("miles", "miles.rollout", "miles.rollout.filter_hub"): - monkeypatch.setitem(sys.modules, name, types.ModuleType(name)) - monkeypatch.setitem(sys.modules, "miles.rollout.filter_hub.base_types", output_module) - - args = _args(yeto_rl_dynamic_sampling_max_replacements=0) - source = _RetrySource(args) - group = source.get_samples(TARGET + 1)[0] - for sample in group: - sample.status = SimpleNamespace(value="failed") - sample.metadata = {"outcome": {"status": reward.INFRASTRUCTURE_STATUS, - reward.INFRASTRUCTURE_503_RETRIES_KEY: 0}} - result = reward.check_group(args, group) - assert (result.keep, result.reason) == (False, "secrlenv_infrastructure_failure") - (rerun,) = source.get_samples(TARGET + 1) - assert [s.index for s in rerun] == [s.index for s in group] - - -# ------------------------------------------------------- terminal cleanup -def _cleanup(error, *, abort=None, require_clean=lambda: None, timeout_s=1.0): - calls = [] - - async def default_abort(): - calls.append("abort") - - asyncio.run(retry.raise_after_terminal_cleanup( - error, abort=abort or default_abort, in_flight=list, require_clean=require_clean, - timeout_s=timeout_s)) - return calls - - -def test_terminal_error_is_raised_after_the_abort(): - error = SecRLEnvReplacementExhausted("budget") - with pytest.raises(SecRLEnvReplacementExhausted) as raised: - _cleanup(error) - assert raised.value is error and raised.value.__cause__ is None - - -def test_untrusted_outcome_is_sanitized(): - with pytest.raises(SecRLEnvUntrustedEvidence, match="authentication failed") as raised: - _cleanup(UntrustedOutcome("payload with episode secrets")) - assert "secrets" not in str(raised.value) and raised.value.__cause__ is None - - -@pytest.mark.parametrize("kind", ["abort", "residue", "timeout"]) -def test_cleanup_failure_replaces_the_error(kind): - async def failing_abort(): - raise RuntimeError("engine said no") - - async def slow_abort(): - await asyncio.sleep(5) - - def residue(): - raise RuntimeError("episode left") - - options = {"abort": dict(abort=failing_abort), "residue": dict(require_clean=residue), - "timeout": dict(abort=slow_abort, timeout_s=0.01)}[kind] - with pytest.raises(SecRLEnvRolloutCleanupError) as raised: - _cleanup(SecRLEnvReplacementExhausted("x"), **options) - assert raised.value.__cause__ is None - - -# ------------------------------------------------------- Miles classes -@pytest.fixture -def stub_miles(monkeypatch): - """Minimal Miles modules the binding module imports.""" - - calls = [] - - class RolloutDataSourceWithBuffer(_Source): - pass - - class InferenceRolloutFn: - def __init__(self, input): - self.state = SimpleNamespace(args=input.args) - self.outcome = input.outcome - - async def _call_train(self, input): - if isinstance(self.outcome, BaseException): - raise self.outcome - if callable(self.outcome): # a driver standing in for generate_rollout_async - return await self.outcome(self.state.args) - return self.outcome - - async def abort(state, pendings, rollout_id): - calls.append(("abort", pendings, rollout_id)) - return [] - - modules = { - "miles": {}, "miles.rollout": {}, "miles.rollout.inference_rollout": {}, - "miles.rollout.data_source": {"RolloutDataSourceWithBuffer": RolloutDataSourceWithBuffer}, - "miles.rollout.inference_rollout.inference_rollout_common": { - "InferenceRolloutFn": InferenceRolloutFn}, - "miles.rollout.inference_rollout.inference_rollout_train": {"abort": abort}, - } - for name, attrs in modules.items(): - module = types.ModuleType(name) - for key, value in attrs.items(): - setattr(module, key, value) - monkeypatch.setitem(sys.modules, name, module) - monkeypatch.delitem(sys.modules, "yeto.rl.adapters.miles.secrlenv_miles", raising=False) - module = importlib.import_module("yeto.rl.adapters.miles.secrlenv_miles") - yield module, calls - sys.modules.pop("yeto.rl.adapters.miles.secrlenv_miles", None) - - -def test_miles_paths_name_the_binding_classes(stub_miles): - module, _ = stub_miles - for path in (retry.DATA_SOURCE_PATH, retry.ROLLOUT_FUNCTION_PATH): - module_name, _, name = path.rpartition(".") - assert module_name == module.__name__ and hasattr(module, name) - source = module.SecRLEnvRetryDataSource(_args()) - assert len(source.get_samples(TARGET + 1)) == TARGET - - -def test_rollout_function_reports_the_retry_counters(stub_miles): - module, _ = stub_miles - args = _args() - module.SecRLEnvRetryDataSource(args) - output = SimpleNamespace(metrics={"rollout/x": 1.0}) - fn = module.SecRLEnvRolloutFn(SimpleNamespace(args=args, outcome=output)) - result = asyncio.run(fn._call_train(SimpleNamespace(rollout_id=3))) - assert result.metrics["rl/secrlenv/replacement_attempts"] == 0.0 - assert result.metrics["rollout/x"] == 1.0 - - -def test_rollout_function_drains_before_a_terminal_error(stub_miles, monkeypatch): - module, calls = stub_miles - monkeypatch.setattr(module, "_require_no_episode_residue", lambda: calls.append("clean")) - fn = module.SecRLEnvRolloutFn(SimpleNamespace( - args=_args(), outcome=SecRLEnvReplacementExhausted("budget"))) - with pytest.raises(SecRLEnvReplacementExhausted): - asyncio.run(fn._call_train(SimpleNamespace(rollout_id=7))) - assert calls == [("abort", set(), 7), "clean"] - - -# ------------------------------------------------------------- launcher -def test_launcher_ports_preflight_accepts_a_secrlenv_run(): - """Before P0-1.7b.4-Q1 the ports spec refused the group check here; the - expected spec hash is the one the island derives from the bound filter.""" - - from test_rl_launcher import _args as launcher_args - - from yeto.launcher import _prepare_rl_args - - args = launcher_args([ - "--reward-function", "yeto.rl.harness.codex.reward:reward_func", - "--custom-generate-function-path", SECRLENV_GENERATE, - "--custom-agent-function-path", SECRLENV_AGENT, - "--use-session-server", "--tito-model", "org/model", - ]) - _prepare_rl_args(args) - assert args.dynamic_sampling_filter_path == SECRLENV_GROUP_FILTER - assert args.rl_expected_algorithm_sha256 == _secrlenv_spec().sha256() - # Regression: computed before the contract bound the filter, the hash was - # that of the filterless spec and the island refused the run. - assert args.rl_expected_algorithm_sha256 != AlgorithmSpec(advantage_estimator="grpo").sha256() - assert args.over_sampling_batch_size == args.rollout_batch_size + 1 - - -def _miles_like_round(source, *, stuck): - """Stands in for Miles' generate_rollout_async: group tasks run as - ``generate_and_rm_group``; an exception inside a task is swallowed (the - group only counts as failed) and a dropped group triggers a refill.""" - - from yeto.rl.harness.codex.reward import remember_terminal_error - - async def generate_and_rm_group(args, group, fail): - if fail: - error = UntrustedOutcome("forged outcome with episode payload") - remember_terminal_error(args, error) # what reward_func/generate do - raise error - await asyncio.sleep(3600) - - async def driver(args): - groups = source.get_samples(TARGET + 1) - pendings = {asyncio.create_task(generate_and_rm_group(args, g, i == 0)) - for i, g in enumerate(groups)} - stuck.extend(pendings) - done, pendings = await asyncio.wait(pendings, return_when=asyncio.FIRST_COMPLETED) - for task in done: - try: - task.result() - except Exception: # noqa: BLE001 - Miles: logged, group counted as failed - continue - source.get_samples(TARGET + 1) # refill for the dropped group: no retry queued - - return driver - - -def test_rollout_function_raises_the_first_terminal_error_and_stops_group_tasks( - stub_miles, monkeypatch): - module, calls = stub_miles - monkeypatch.setattr(module, "_require_no_episode_residue", lambda: calls.append("clean")) - args = _args() - source = module.SecRLEnvRetryDataSource(args) - stuck = [] - fn = module.SecRLEnvRolloutFn(SimpleNamespace( - args=args, outcome=_miles_like_round(source, stuck=stuck))) - - async def run(): - with pytest.raises(SecRLEnvUntrustedEvidence) as raised: - await fn._call_train(SimpleNamespace(rollout_id=4)) - # checked inside the loop: asyncio.run would cancel leftovers on exit - assert len(stuck) == TARGET and all(task.done() for task in stuck) - return raised.value - - error = asyncio.run(run()) - assert "payload" not in str(error) # sanitized first error, not the refill error - assert calls == [("abort", set(), 4), "clean"] - assert sum(task.cancelled() for task in stuck) == TARGET - 1 - - -def test_round_start_forgets_the_previous_terminal_error(): - from yeto.rl.harness.codex.reward import TERMINAL_ERROR_ATTR - - args = _args() - source = _RetrySource(args) - setattr(args, TERMINAL_ERROR_ATTR, UntrustedOutcome("old")) - source.get_samples(TARGET + 1) - assert getattr(args, TERMINAL_ERROR_ATTR) is None - fallback = SecRLEnvReplacementExhausted("x") - assert retry.first_terminal_error(args, fallback) is fallback diff --git a/tests/test_rl_ssh_harness.py b/tests/test_rl_ssh_harness.py index 1d15010ce..091bbe386 100644 --- a/tests/test_rl_ssh_harness.py +++ b/tests/test_rl_ssh_harness.py @@ -1,4 +1,3 @@ -import copy import hashlib import json import os @@ -10,8 +9,6 @@ import pytest import torch -import yeto.rl as rl_config -from yeto.provenance import python_spec_sha256 from yeto.rl import ( MILES_NEXT_COMMIT, MILES_NEXT_REPOSITORY, @@ -21,13 +18,6 @@ ssh_harness, ) from yeto.rl.bridge import _write_round_audit -from yeto.rl.codex_backend import ( - QWEN35_08B_MODEL, - QWEN35_08B_REVISION, - QWEN38_MODEL, - QWEN38_REVISION, - stock_codex_backend_contract, -) from yeto.rl.core import ( CanonicalTensorSpec, canonical_layout_hash, @@ -151,38 +141,16 @@ def _plan(): } -def _secrlenv_daemon_contract(run_id="acceptance"): - return { - "source_root": "/data/yeto-rl/src/secrlenv-v24", - "source_sha256": "2" * 64, - "task_pack": "/data/yeto-rl/taskpacks/" + "3" * 64, - "task_pack_sha256": "3" * 64, - "state_root": f"/data/yeto-rl/secrlenv-runs/{run_id}", - "bind": "127.0.0.1", - "port": 28765, - "operator_image": "secrlenv-operator:pinned", - "operator_image_id": "sha256:" + "4" * 64, - "max_active_episodes": 16, - "enable_dind_debug": False, - "dind_image": None, - "dind_image_id": None, - "dind_debug_script_sha256": None, - } - - -def _enable_secrlenv_agent(plan): - plan["remote_env_file"] = ".config/yeto/rl.env" - plan["learner"].update( - custom_agent_function_path="yeto.rl.harness.codex.agent.run", - custom_generate_function_path=ssh_harness.SECRLENV_GENERATE, - use_session_server=True, - tito_model="org/model", - reward_function=ssh_harness.SECRLENV_REWARD, - dynamic_sampling_filter_path=ssh_harness.SECRLENV_GROUP_FILTER, - dynamic_sampling_max_replacements=0, - secrlenv_max_infrastructure_replacements=1, - over_sampling_batch_size=plan["learner"]["groups_per_round"] + 1, - ) +@pytest.mark.parametrize("field, value, message", [ + ("secrlenv_daemon", {"port": 28765}, "secrlenv daemon plans"), + ("codex_harness", {"cli_version": "x"}, "stock Codex contract"), + ("eval_checkpoint", {"sha256": "f" * 64}, "removed secrlenv agent"), +]) +def test_plans_of_the_removed_codex_harness_are_refused(field, value, message): + plan = _plan() + plan[field] = value + with pytest.raises(ssh_harness.HarnessError, match=message): + ssh_harness._validate_plan(plan) def _local_data_plan(path: Path): @@ -198,75 +166,6 @@ def _local_data_plan(path: Path): return plan -def test_secrlenv_plan_round_trip_requires_exact_replacement_contract(): - plan = _plan() - _enable_secrlenv_agent(plan) - plan["secrlenv_daemon"] = _secrlenv_daemon_contract(plan["run_id"]) - - ssh_harness._validate_plan(plan) - parsed = parse_learner_args(_learner_argv(plan, 0)[3:]) - - assert parsed.dynamic_sampling_filter_path == ssh_harness.SECRLENV_GROUP_FILTER - assert parsed.dynamic_sampling_max_replacements == 0 - assert parsed.secrlenv_max_infrastructure_replacements == 1 - assert parsed.over_sampling_batch_size == parsed.groups_per_round + 1 - - mutations = [ - ( - "custom_generate_function_path", - "miles.rollout.generate_hub.agentic_tool_call.generate", - ), - ("reward_function", "other.reward:score"), - ("dynamic_sampling_filter_path", None), - ("dynamic_sampling_filter_path", "other.filter"), - ("dynamic_sampling_max_replacements", None), - ("dynamic_sampling_max_replacements", False), - ("dynamic_sampling_max_replacements", 1), - ("secrlenv_max_infrastructure_replacements", None), - ("secrlenv_max_infrastructure_replacements", True), - ("secrlenv_max_infrastructure_replacements", 2), - ("over_sampling_batch_size", plan["learner"]["groups_per_round"]), - ] - for name, value in mutations: - invalid = copy.deepcopy(plan) - invalid["learner"][name] = value - with pytest.raises(ssh_harness.HarnessError): - ssh_harness._validate_plan(invalid) - - -def _secrlenv_eval_plan(path: Path): - pytest.importorskip("yeto.rl.harness.codex") - plan = _local_data_plan(path) - _enable_secrlenv_agent(plan) - plan["learner"]["reward_function"] = ssh_harness.SECRLENV_REWARD - plan["reward_sha256"] = python_spec_sha256(ssh_harness.SECRLENV_REWARD) - plan["secrlenv_daemon"] = _secrlenv_daemon_contract(plan["run_id"]) - plan["final_ack_timeout_s"] = 21600 - plan["eval_checkpoint"] = ssh_harness._eval_checkpoint_contract( - "root@h200-n6", - "yeto-rl-ssh-data/train-run/state/state.ckpt", - "f" * 64, - 138_632_000_000, - plan["learner"]["global_rounds"], - ) - plan["learner"]["evaluation"] = { - "eval_only": True, - "dataset_name": "flaky100", - "data": plan["learner"]["data"], - "data_sha256": plan["learner"]["data_sha256"], - "interval": 1, - "samples_per_prompt": 2, - "prompt_count": 1, - "skip_before_train": True, - "temperature": 0.7, - "top_p": 0.95, - "max_prompt_len": 256, - "max_response_len": 256, - "max_context_len": 512, - } - return plan - - def _tms_patch_contract(): return { "repository": ssh_harness.TMS_PRELOAD_PATCH_REPOSITORY, @@ -769,84 +668,6 @@ def test_local_prompt_plan_is_content_bound_without_a_hub_revision(tmp_path): ) -def test_secrlenv_eval_reuses_the_attested_training_dataset_and_extends_final_ack( - tmp_path, -): - prompts = tmp_path / "flaky100.jsonl" - prompts.write_text('{"messages": []}\n', encoding="utf-8") - plan = _secrlenv_eval_plan(prompts) - - ssh_harness._validate_plan(plan) - evaluation = plan["learner"]["evaluation"] - assert evaluation["data"] == plan["learner"]["data"] - assert evaluation["data_sha256"] == plan["learner"]["data_sha256"] - assert evaluation["eval_only"] is True - assert evaluation["interval"] == 1 - assert evaluation["prompt_count"] == 1 - assert plan["learner"]["reward_function"] == ssh_harness.SECRLENV_REWARD - assert plan["reward_sha256"] == python_spec_sha256(ssh_harness.SECRLENV_REWARD) - assert plan["eval_checkpoint"]["global_step"] == plan["learner"]["global_rounds"] - syncer = _syncer_argv(plan) - assert syncer[syncer.index("--final-ack-timeout-s") + 1] == "21600" - - argv = _learner_argv(plan, 0) - parsed = parse_learner_args(argv[3:]) - assert parsed.eval_only is True - assert parsed.eval_dataset_name == "flaky100" - assert parsed.eval_data_sha256 == plan["learner"]["data_sha256"] - assert parsed.eval_interval == 1 - assert parsed.eval_samples_per_prompt == 2 - assert parsed.reward_function == ssh_harness.SECRLENV_REWARD - assert parsed.eval_temperature == 0.7 - assert parsed.eval_top_p == 0.95 - assert parsed.eval_max_prompt_len == 256 - assert parsed.eval_max_response_len == 256 - assert parsed.eval_max_context_len == 512 - - -@pytest.mark.parametrize("field", ["data", "data_sha256"]) -def test_secrlenv_eval_rejects_a_distinct_training_dataset_identity(tmp_path, field): - prompts = tmp_path / "flaky100.jsonl" - prompts.write_text('{"messages": []}\n', encoding="utf-8") - plan = _secrlenv_eval_plan(prompts) - plan["learner"]["evaluation"][field] = ( - "/workspace/data/other.jsonl" if field == "data" else "0" * 64 - ) - - with pytest.raises(HarnessError, match="exact training dataset"): - ssh_harness._validate_plan(plan) - - -def test_secrlenv_eval_requires_an_explicitly_extended_final_ack_timeout(tmp_path): - prompts = tmp_path / "flaky100.jsonl" - prompts.write_text('{"messages": []}\n', encoding="utf-8") - plan = _secrlenv_eval_plan(prompts) - plan.pop("final_ack_timeout_s") - - with pytest.raises(HarnessError, match="extended final_ack_timeout"): - ssh_harness._validate_plan(plan) - - -def test_secrlenv_eval_rejects_a_different_reward_contract(tmp_path): - prompts = tmp_path / "flaky100.jsonl" - prompts.write_text('{"messages": []}\n', encoding="utf-8") - plan = _secrlenv_eval_plan(prompts) - plan["learner"]["reward_function"] = "pkg.reward:score" - - with pytest.raises(HarnessError, match="signed secrlenv reward"): - ssh_harness._validate_plan(plan) - - -def test_secrlenv_eval_checkpoint_step_must_match_terminal_training_step(tmp_path): - prompts = tmp_path / "flaky100.jsonl" - prompts.write_text('{"messages": []}\n', encoding="utf-8") - plan = _secrlenv_eval_plan(prompts) - plan["eval_checkpoint"]["global_step"] -= 1 - - with pytest.raises(HarnessError, match="eval checkpoint identity"): - ssh_harness._validate_plan(plan) - - def test_prepare_maps_a_local_prompt_file_into_the_remote_plan(tmp_path, monkeypatch): prompts = tmp_path / "prompts.jsonl" prompts.write_text('{"messages": []}\n', encoding="utf-8") @@ -898,97 +719,6 @@ def test_prepare_maps_a_local_prompt_file_into_the_remote_plan(tmp_path, monkeyp ) -def test_prepare_wires_final_secrlenv_eval_and_same_dataset_sha(tmp_path, monkeypatch): - pytest.importorskip("yeto.rl.harness.codex") - prompts = tmp_path / "flaky100.jsonl" - prompts.write_text('{"messages": []}\n', encoding="utf-8") - learner = _plan()["learner"] - learner.update( - custom_agent_function_path="yeto.rl.harness.codex.agent.run", - custom_generate_function_path=ssh_harness.SECRLENV_GENERATE, - use_session_server=True, - tito_model="org/model", - reward_function=ssh_harness.SECRLENV_REWARD, - dynamic_sampling_filter_path=ssh_harness.SECRLENV_GROUP_FILTER, - dynamic_sampling_max_replacements=0, - secrlenv_max_infrastructure_replacements=1, - over_sampling_batch_size=learner["groups_per_round"] + 1, - ) - args = SimpleNamespace( - **{ - **learner, - "data": str(prompts), - "data_revision": None, - "total_steps": learner["global_rounds"], - "rollout_batch_size": learner["groups_per_round"], - "n_samples_per_prompt": learner["samples_per_group"], - "rl_image": "docker:" + IMAGE, - "source_sha256": "c" * 64, - "reward_sha256": python_spec_sha256(ssh_harness.SECRLENV_REWARD), - "_provenance": {"dataset": {"source": "local"}}, - } - ) - daemon = _secrlenv_daemon_contract("final-eval") - checkpoint_sha256 = "f" * 64 - namespace = SimpleNamespace( - run_id="final-eval", - host=["alice@a0", "alice@b0"], - gpus_per_node=1, - remote_root=".cache/yeto-rl-ssh", - remote_env_file=".config/yeto/rl.env", - model_manifest_sha256=None, - syncer_address="100.64.0.6:29400", - syncer_host="root@syncer0", - output_dir=tmp_path / "run", - ssh_option=[], - final_ack_timeout_s=21600, - eval_checkpoint_host="root@h200-n6", - eval_checkpoint_path="yeto-rl-ssh-data/train/state/state.ckpt", - eval_checkpoint_sha256=checkpoint_sha256, - eval_checkpoint_size_bytes=138_632_000_000, - eval_checkpoint_global_step=learner["global_rounds"], - eval_dataset_name="flaky100", - eval_interval=1, - eval_samples_per_prompt=2, - eval_temperature=0.7, - eval_top_p=0.95, - eval_max_prompt_len=256, - eval_max_response_len=256, - eval_max_context_len=512, - secrlenv_source_root=daemon["source_root"], - secrlenv_source_sha256=daemon["source_sha256"], - secrlenv_task_pack=daemon["task_pack"], - secrlenv_task_pack_sha256=daemon["task_pack_sha256"], - secrlenv_operator_image=daemon["operator_image"], - secrlenv_operator_image_id=daemon["operator_image_id"], - secrlenv_port=daemon["port"], - secrlenv_max_active_episodes=daemon["max_active_episodes"], - secrlenv_enable_dind_debug=False, - secrlenv_dind_image=None, - secrlenv_dind_image_id=None, - secrlenv_dind_debug_script_sha256=None, - ) - monkeypatch.setattr(ssh_harness, "_resolved_launch_args", lambda *unused: args) - monkeypatch.setattr(ssh_harness, "_syncer_source_sha256", lambda: "1" * 64) - - plan_path = ssh_harness.prepare(namespace) - _, plan = load_plan(plan_path) - evaluation = plan["learner"]["evaluation"] - - assert evaluation["data"] == plan["learner"]["data"] - assert evaluation["data_sha256"] == plan["learner"]["data_sha256"] - assert evaluation["eval_only"] is True - assert evaluation["interval"] == 1 - assert evaluation["prompt_count"] == 1 - assert evaluation["skip_before_train"] is True - assert plan["final_ack_timeout_s"] == 21600 - assert plan["eval_checkpoint"]["source_host"] == "root@h200-n6" - assert plan["eval_checkpoint"]["source_path"].endswith("train/state/state.ckpt") - assert plan["eval_checkpoint"]["sha256"] == checkpoint_sha256 - assert plan["eval_checkpoint"]["global_step"] == learner["global_rounds"] - assert plan["secrlenv_daemon"]["placement"] == "island-heads" - - def test_local_prompt_directory_hash_is_stable_and_rejects_symlinks(tmp_path): first = tmp_path / "first" second = tmp_path / "second" @@ -1106,183 +836,6 @@ def test_deploy_copies_local_prompts_to_every_host_and_mounts_read_only( assert all(command[-1].endswith("/data/dataset.jsonl") for command in copies) -def test_eval_deploy_copies_terminal_checkpoint_only_to_fresh_syncer( - tmp_path, monkeypatch -): - prompts = tmp_path / "flaky100.jsonl" - prompts.write_text('{"messages": []}\n', encoding="utf-8") - plan = _secrlenv_eval_plan(prompts) - plan["syncer_host"] = "root@syncer0" - plan_path = tmp_path / "plan.json" - _write_plan(plan_path, plan) - commands = [] - scripts = [] - monkeypatch.setattr(ssh_harness, "_require_program", lambda name: None) - monkeypatch.setattr(ssh_harness, "_attest_local", lambda value: None) - monkeypatch.setattr( - ssh_harness, "_run", lambda command, **kwargs: commands.append(command) - ) - monkeypatch.setattr( - ssh_harness, - "_ssh", - lambda _plan, target, script, **kwargs: scripts.append((target, script)), - ) - - ssh_harness.deploy(plan_path) - - checkpoint_copies = [ - command for command in commands if command[:2] == ["scp", "-3"] - ] - assert len(checkpoint_copies) == 1 - assert checkpoint_copies[0][-2] == ( - "root@h200-n6:yeto-rl-ssh-data/train-run/state/state.ckpt" - ) - assert checkpoint_copies[0][-1].endswith( - "root@syncer0:.cache/yeto-rl-ssh/acceptance/state/state.ckpt" - ) - source_attestation = next( - script for target, script in scripts if target == "root@h200-n6" - ) - assert plan["eval_checkpoint"]["sha256"] in source_attestation - assert str(plan["eval_checkpoint"]["size_bytes"]) in source_attestation - assert "syncer checkpoint global step mismatch" in source_attestation - syncer_script = ssh_harness._syncer_start_script(plan) - assert plan["eval_checkpoint"]["sha256"] in syncer_script - assert "sha256sum --check -" in syncer_script - - -def test_eval_collect_whitelists_small_evidence_and_excludes_checkpoint( - tmp_path, monkeypatch -): - prompts = tmp_path / "flaky100.jsonl" - prompts.write_text('{"messages": []}\n', encoding="utf-8") - plan = _secrlenv_eval_plan(prompts) - plan["islands"] = [plan["islands"][0]] - plan_path = tmp_path / "plan.json" - _write_plan(plan_path, plan) - commands = [] - ssh_scripts = [] - - def ssh(_plan, _target, script, **_kwargs): - ssh_scripts.append(script) - if script.startswith("docker inspect"): - stdout = json.dumps( - { - "Status": "exited", - "ExitCode": 0, - "OOMKilled": False, - "RestartCount": 0, - } - ) - elif "final_ack_count" in script: - stdout = json.dumps({"schema": 1, "resume_count": 1, "final_ack_count": 1}) - else: - stdout = "" - return SimpleNamespace(stdout=stdout, stderr="", returncode=0) - - monkeypatch.setattr(ssh_harness, "_require_program", lambda _name: None) - monkeypatch.setattr(ssh_harness, "_ssh", ssh) - monkeypatch.setattr( - ssh_harness, "_run", lambda command, **_kwargs: commands.append(command) - ) - - ssh_harness.collect(plan_path) - - remote_sources = [command[-2] for command in commands] - assert not any("/island-0/state/" in source for source in remote_sources) - assert not any("/island-0/audit/" in source for source in remote_sources) - assert not any("/island-0/output/" in source for source in remote_sources) - learner_filters = [script for script in ssh_scripts if "eval-learner" in script] - assert len(learner_filters) == 1 - syncer = next(command for command in commands if command[-2].endswith("/state/")) - assert "--include=/syncer.exit" in syncer - assert "--include=/syncer.log" not in syncer - assert "--include=/events.jsonl" not in syncer - assert "--exclude=*" in syncer - assert not any("state.ckpt" in value for value in syncer) - assert not any("docker logs" in script for script in ssh_scripts) - inspect_scripts = [ - script for script in ssh_scripts if script.startswith("docker inspect") - ] - assert inspect_scripts - assert all("--format" in script for script in inspect_scripts) - for inspection in (tmp_path / "artifacts").glob("island-*/node-*.inspect.json"): - assert set(json.loads(inspection.read_text())) == { - "Status", - "ExitCode", - "OOMKilled", - "RestartCount", - } - lifecycle = json.loads( - (tmp_path / "artifacts/syncer/lifecycle-evidence.json").read_text() - ) - assert lifecycle == {"schema": 1, "resume_count": 1, "final_ack_count": 1} - assert not list((tmp_path / "artifacts").glob("island-*/state")) - - -def test_eval_verify_uses_scalar_final_ack_evidence_without_checkpoint( - tmp_path, capsys -): - prompts = tmp_path / "flaky100.jsonl" - prompts.write_text('{"messages": []}\n', encoding="utf-8") - plan = _secrlenv_eval_plan(prompts) - plan["islands"] = [plan["islands"][0]] - plan_path = tmp_path / "plan.json" - _write_plan(plan_path, plan) - artifacts = tmp_path / "artifacts" - island = artifacts / "island-0" - output = island / "output" - output.mkdir(parents=True) - for node_id in range(len(plan["islands"][0]["hosts"])): - (island / f"node-{node_id}.inspect.json").write_text( - json.dumps( - { - "Status": "exited", - "ExitCode": 0, - "OOMKilled": False, - "RestartCount": 0, - } - ) - ) - _write_jsonl( - output / "events.jsonl", - [ - { - "event": "rl_policy_apply", - "island_id": 0, - "policy_version": plan["learner"]["global_rounds"], - "sync/global_policy_hash": "a" * 64, - }, - { - "event": "rl_eval_result", - "island_id": 0, - "rollout_id": 0, - "policy_version": plan["learner"]["global_rounds"], - "dataset_name": "flaky100", - "sample_count": 2, - "rl/eval/result": 0.5, - "rl/eval/pass_at_1": 0.5, - }, - ], - ) - syncer = artifacts / "syncer" - syncer.mkdir() - (syncer / "syncer.exit").write_text( - "result=success\n" - "exit_code=exited\n" - "exit_status=0\n" - "timestamp=2026-08-17T00:00:00Z\n" - ) - (syncer / "lifecycle-evidence.json").write_text( - json.dumps({"schema": 1, "resume_count": 1, "final_ack_count": 1}) - ) - - ssh_harness.verify(plan_path) - - assert "verified terminal eval flaky100" in capsys.readouterr().out - assert not (syncer / "state.ckpt").exists() - - def test_local_prompt_learner_command_omits_hub_revision_and_mounts_read_only( tmp_path, ): @@ -2037,264 +1590,6 @@ def test_harness_forwards_agentic_session_contract_to_the_learner(): assert "--env MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1" in script -def test_harness_freezes_stock_codex_identity_and_binary_mount(): - plan = _plan() - plan["remote_env_file"] = ".config/yeto/rl.env" - plan["secrlenv_daemon"] = _secrlenv_daemon_contract() - plan["learner"].update( - { - "rl_model_recipe": "deepseek-v4-flash", - "custom_generate_function_path": rl_config.SECRLENV_GENERATE, - "custom_agent_function_path": rl_config.CODEX_HARNESS_AGENT, - "codex_reasoning_effort": "xhigh", - "use_session_server": True, - "tito_model": "deepseekv4", - "tito_allowed_append_roles": ["tool", "user"], - "tensor_parallel": 8, - "expert_parallel": 8, - "rollout_num_gpus_per_engine": 8, - "sglang_tp_size": 8, - "sglang_dp_size": 1, - "sglang_ep_size": 8, - "lora_targets": "attention-routed-experts", - "apply_chat_template_kwargs": { - "thinking_mode": "thinking", - "reasoning_effort": "max", - "drop_thinking": False, - }, - "reward_function": ssh_harness.SECRLENV_REWARD, - "dynamic_sampling_filter_path": ssh_harness.SECRLENV_GROUP_FILTER, - "dynamic_sampling_max_replacements": 0, - "secrlenv_max_infrastructure_replacements": 1, - "over_sampling_batch_size": plan["learner"]["groups_per_round"] + 1, - } - ) - plan["codex_harness"] = { - "agent_function_path": rl_config.CODEX_HARNESS_AGENT, - "agent_source_sha256": rl_config.CODEX_HARNESS_AGENT_SHA256, - "controller_binary_path": "/controller/codex", - "controller_package_manifest_path": "/controller/codex-package.json", - "controller_app_server_schema_path": "/controller/app-server-v2.json", - "bundle_binary_path": ssh_harness.CODEX_REMOTE_BUNDLE_PATH, - "bundle_package_manifest_path": ssh_harness.CODEX_REMOTE_MANIFEST_PATH, - "bundle_app_server_schema_path": ssh_harness.CODEX_REMOTE_SCHEMA_PATH, - "container_binary_path": rl_config.CODEX_CONTAINER_BINARY_PATH, - "container_app_server_schema_path": ( - rl_config.CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH - ), - "binary_sha256": rl_config.CODEX_LINUX_BINARY_SHA256, - "binary_size_bytes": rl_config.CODEX_LINUX_BINARY_SIZE_BYTES, - "cli_version": rl_config.CODEX_CLI_VERSION, - "npm_package": rl_config.CODEX_NPM_PACKAGE, - "target": rl_config.CODEX_LINUX_TARGET, - "npm_tarball_sha256": rl_config.CODEX_NPM_TARBALL_SHA256, - "package_manifest_sha256": rl_config.CODEX_PACKAGE_MANIFEST_SHA256, - "app_server_protocol_revision": (rl_config.CODEX_APP_SERVER_PROTOCOL_REVISION), - "app_server_schema_sha256": rl_config.CODEX_APP_SERVER_SCHEMA_SHA256, - "base_instructions_sha256": rl_config.CODEX_BASE_INSTRUCTIONS_SHA256, - "terminal_exec_tool_schema_sha256": ( - rl_config.CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256 - ), - "submit_tool_schema_sha256": (rl_config.CODEX_SUBMIT_TOOL_SCHEMA_SHA256), - "dynamic_tools_schema_sha256": (rl_config.CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256), - "reasoning_effort": "xhigh", - "backend": { - "profile": "deepseekv4", - "tito_model": "deepseekv4", - "model": "deepseekv4", - "max_tokens": plan["learner"]["rollout_max_response_len"], - "reasoning_effort": "max", - "thinking": {"type": "enabled"}, - "chat_template": "deepseekv4", - "chat_template_kwargs": { - "thinking_mode": "thinking", - "reasoning_effort": "max", - "drop_thinking": False, - }, - "tito_allowed_append_roles": ["tool", "user"], - }, - } - - ssh_harness._validate_plan(plan) - argv = _learner_argv(plan, 0) - parsed = parse_learner_args(argv[3:]) - - assert parsed.codex_harness_contract == plan["codex_harness"] - for node_id in (0, 1): - script = _node_start_script(plan, 0, node_id) - assert rl_config.CODEX_CONTAINER_BINARY_PATH in script - assert "YETO_CODEX_REASONING_EFFORT=xhigh" in script - assert "YETO_CODEX_APP_SERVER_SCHEMA_SHA256=" in script - assert "YETO_CODEX_HARNESS_CONTRACT_SHA256=" in script - syntax = subprocess.run( - ["bash", "-n"], - input=script, - text=True, - capture_output=True, - check=False, - ) - assert syntax.returncode == 0, syntax.stderr - - drifted = json.loads(json.dumps(plan)) - drifted["codex_harness"]["backend"]["chat_template_kwargs"]["drop_thinking"] = True - with pytest.raises(HarnessError, match="backend/TITO contract drifted"): - ssh_harness._validate_plan(drifted) - - for field in ( - "agent_source_sha256", - "base_instructions_sha256", - "terminal_exec_tool_schema_sha256", - "submit_tool_schema_sha256", - "dynamic_tools_schema_sha256", - ): - drifted = json.loads(json.dumps(plan)) - drifted["codex_harness"][field] = "0" * 64 - with pytest.raises(HarnessError, match="pinned stock Codex runtime"): - ssh_harness._validate_plan(drifted) - - qwen = copy.deepcopy(plan) - qwen["learner"].update( - { - "model": QWEN38_MODEL, - "model_revision": QWEN38_REVISION, - "rollout_model": None, - "rollout_model_revision": None, - "rl_model_recipe": "generic", - "tito_model": "qwen38small", - "codex_backend_profile": "qwen38", - "apply_chat_template_kwargs": { - "enable_thinking": True, - "preserve_thinking": True, - "reasoning_effort": "xhigh", - }, - "lora_targets": "attention", - "expert_full_count": 0, - "expert_parallel": 1, - "tensor_parallel": 4, - "pipeline_parallel": 2, - "rollout_num_gpus_per_engine": 1, - "sglang_tp_size": 1, - "sglang_dp_size": 1, - "sglang_ep_size": 1, - "sglang_attention_backend": None, - "sglang_page_size": None, - } - ) - qwen["codex_harness"]["backend"] = stock_codex_backend_contract( - "qwen38", - qwen["learner"]["rollout_max_response_len"], - ) - - ssh_harness._validate_plan(qwen) - qwen_args = parse_learner_args(_learner_argv(qwen, 0)[3:]) - assert qwen_args.model == QWEN38_MODEL - assert qwen_args.model_revision == QWEN38_REVISION - assert qwen_args.tito_model == "qwen38small" - assert qwen_args.apply_chat_template_kwargs == { - "enable_thinking": True, - "preserve_thinking": True, - "reasoning_effort": "xhigh", - } - - openenv = copy.deepcopy(qwen) - openenv.pop("secrlenv_daemon") - openenv["learner"].update( - { - "model": QWEN35_08B_MODEL, - "model_revision": QWEN35_08B_REVISION, - "tito_model": "qwen35", - "codex_backend_profile": "qwen35_08b", - "apply_chat_template_kwargs": {"clear_thinking": False}, - "custom_generate_function_path": ( - "miles.rollout.generate_hub.agentic_tool_call.generate" - ), - "custom_agent_function_path": rl_config.CODEX_OPENENV_AGENT, - "reward_function": "openenv_generate.reward_func", - "dynamic_sampling_filter_path": ( - "openenv_generate.check_terminal_bench_episode" - ), - "secrlenv_max_infrastructure_replacements": None, - } - ) - openenv["codex_harness"]["backend"] = stock_codex_backend_contract( - "qwen35_08b", - openenv["learner"]["rollout_max_response_len"], - ) - openenv["codex_harness"]["openenv_identity_env"] = dict( - rl_config.CODEX_OPENENV_IDENTITY_ENV - ) - - ssh_harness._validate_plan(openenv) - openenv_args = parse_learner_args(_learner_argv(openenv, 0)[3:]) - assert openenv_args.tito_model == "qwen35" - assert openenv_args.codex_backend_profile == "qwen35_08b" - openenv_script = _node_start_script(openenv, 0, 0) - assert "YETO_CODEX_OPENENV_BACKEND_PROFILE=qwen35_08b" in openenv_script - assert "YETO_CODEX_OPENENV_MODEL_ID=Qwen/Qwen3.5-0.8B" in openenv_script - - drifted_openenv = copy.deepcopy(openenv) - drifted_openenv["codex_harness"]["openenv_identity_env"][ - "YETO_CODEX_OPENENV_MODEL_REVISION" - ] = "0" * 40 - with pytest.raises(HarnessError, match="pinned Codex OpenEnv surface"): - ssh_harness._validate_plan(drifted_openenv) - - qwen["learner"]["model_revision"] = "0" * 40 - with pytest.raises(HarnessError, match="model profile drifted"): - ssh_harness._validate_plan(qwen) - - -def test_codex_controller_artifacts_are_attested_before_plan_write( - tmp_path, monkeypatch -): - binary = tmp_path / "codex" - manifest = tmp_path / "codex-package.json" - schema = tmp_path / "app-server-v2.json" - binary.write_bytes(b"stock-linux-codex") - binary.chmod(0o555) - manifest.write_bytes(b"signed-package-manifest") - schema.write_bytes(b'{"b": 1, "a": [2]}') - digest = lambda path: hashlib.sha256(path.read_bytes()).hexdigest() - monkeypatch.setattr( - ssh_harness, "CODEX_LINUX_BINARY_SIZE_BYTES", binary.stat().st_size - ) - monkeypatch.setattr(ssh_harness, "CODEX_LINUX_BINARY_SHA256", digest(binary)) - monkeypatch.setattr(ssh_harness, "CODEX_PACKAGE_MANIFEST_SHA256", digest(manifest)) - from yeto.provenance import canonical_json_sha256 - - schema_digest = canonical_json_sha256(schema) - assert schema_digest != digest(schema) # pin is canonical-form, not byte-form - monkeypatch.setattr(ssh_harness, "CODEX_APP_SERVER_SCHEMA_SHA256", schema_digest) - namespace = SimpleNamespace( - codex_harness_binary=str(binary), - codex_package_manifest=str(manifest), - codex_app_server_schema=str(schema), - ) - args = SimpleNamespace( - custom_agent_function_path=rl_config.CODEX_OPENENV_AGENT, - codex_reasoning_effort="xhigh", - codex_backend_profile="qwen35_08b", - tito_model="qwen35", - rollout_max_response_len=4096, - tito_allowed_append_roles=["tool", "user"], - apply_chat_template_kwargs={"clear_thinking": False}, - ) - - contract = ssh_harness._codex_harness_contract(namespace, args) - - assert contract["controller_binary_path"] == str(binary) - assert contract["binary_sha256"] == digest(binary) - assert contract["app_server_schema_sha256"] == schema_digest - assert contract["backend"]["reasoning_effort"] == "xhigh" - assert contract["openenv_identity_env"] == rl_config.CODEX_OPENENV_IDENTITY_ENV - - binary.chmod(0o755) - binary.write_bytes(b"drifted") - binary.chmod(0o555) - with pytest.raises(HarnessError, match="binary size"): - ssh_harness._codex_harness_contract(namespace, args) - - def test_harness_forwards_deepseek_v4_recipe_to_the_learner(): plan = _plan() plan["learner"]["rl_model_recipe"] = "deepseek-v4-flash" @@ -2400,206 +1695,6 @@ def test_decoupled_plan_propagates_fragment_and_variance_filter_settings(): assert learner[learner.index("--rl-distributed-timeout-minutes") + 1] == "7" -def test_secrlenv_daemon_contract_is_required_and_propagated_to_nodes(): - plan = _plan() - _enable_secrlenv_agent(plan) - - with pytest.raises(HarnessError, match="pinned daemon contract"): - ssh_harness._validate_plan(plan) - - plan["secrlenv_daemon"] = _secrlenv_daemon_contract() - script = ssh_harness._secrlenv_daemon_script(plan) - node = _node_start_script(plan, 0, 0) - - assert "python3 -m secrlenv_rl.server" in script - assert "\n --enable-dind-debug" not in script - assert "\n --dind-image" not in script - assert "task_pack_sha256" in script - assert "python3 -m yeto.rl.secrlenv_task_images" in script - assert "secrlenv_task_images=ready" not in script - assert ( - script.index('test "$SOURCE_SHA"') - < script.index("python3 -m yeto.rl.secrlenv_task_images") - < script.index("python3 -m secrlenv_rl.server") - ) - assert "secrlenv_daemon=ready" in script - assert 'ENV_FILE="$HOME/.config/yeto/rl.env"' in script - assert 'TOKEN_FILE="$STATE_ROOT/daemon.token"' in script - assert "secrets.token_hex(32)" in script - assert '--token-file "$TOKEN_FILE"' in script - assert "--env SECRLENV_DAEMON_URL=http://127.0.0.1:28765" in node - assert ("--env SECRLENV_TASK_PACK_SHA256=" + "3" * 64) in node - assert "--env SECRLENV_BEARER_TOKEN_FILE=/run/secrlenv/daemon.token" in node - assert ( - "/data/yeto-rl/secrlenv-runs/acceptance/daemon.token:" - "/run/secrlenv/daemon.token:ro" in node - ) - - -def test_secrlenv_dind_debug_is_explicit_pinned_and_propagated(): - plan = _plan() - _enable_secrlenv_agent(plan) - plan["secrlenv_daemon"] = { - **_secrlenv_daemon_contract(), - "enable_dind_debug": True, - "dind_image": "docker:24-dind", - "dind_image_id": "sha256:" + "5" * 64, - "dind_debug_script_sha256": "6" * 64, - } - - ssh_harness._validate_plan(plan) - script = ssh_harness._secrlenv_daemon_script(plan) - - assert "--enable-dind-debug" in script - assert "--dind-image sha256:" + "5" * 64 in script - assert "DIND_IMAGE_ID=" in script - assert "sha256:" + "5" * 64 in script - assert "DIND_DEBUG_SCRIPT=" in script - assert "6" * 64 in script - assert "existing secrlenv daemon has the wrong DinD debug identity" in script - - -@pytest.mark.parametrize( - ("updates", "message"), - [ - ({"enable_dind_debug": 1}, "explicit boolean"), - ( - {"enable_dind_debug": True}, - "requires a pinned image and image ID", - ), - ( - { - "enable_dind_debug": True, - "dind_image": "docker:24-dind", - "dind_image_id": "sha256:" + "5" * 64, - }, - "requires a pinned debug script", - ), - ( - {"dind_image": "docker:24-dind"}, - "must not carry unused debug identities", - ), - ], -) -def test_secrlenv_dind_debug_contract_fails_closed(updates, message): - plan = _plan() - _enable_secrlenv_agent(plan) - plan["secrlenv_daemon"] = {**_secrlenv_daemon_contract(), **updates} - - with pytest.raises(HarnessError, match=message): - ssh_harness._validate_plan(plan) - - -def test_secrlenv_island_head_placement_is_explicit_and_fail_closed( - tmp_path, monkeypatch -): - plan = _plan() - _enable_secrlenv_agent(plan) - plan["secrlenv_daemon"] = { - **_secrlenv_daemon_contract(), - "placement": "island-heads", - } - ssh_harness._validate_plan(plan) - - assert ssh_harness._secrlenv_daemon_hosts(plan) == [ - "alice@a0", - "alice@b0", - ] - for learner_id in range(2): - head = _node_start_script(plan, learner_id, 0) - worker = _node_start_script(plan, learner_id, 1) - assert "SECRLENV_DAEMON_URL=" in head - assert "/run/secrlenv/daemon.token:ro" in head - assert "SECRLENV_DAEMON_URL=" not in worker - assert "/run/secrlenv/daemon.token:ro" not in worker - - calls = [] - - def fake_ssh(plan, target, script, **kwargs): - calls.append((target, script)) - return SimpleNamespace(stdout="", stderr="") - - monkeypatch.setattr(ssh_harness, "_ssh", fake_ssh) - ssh_harness._start_secrlenv_daemons(plan) - assert [target for target, _script in calls] == ["alice@a0", "alice@b0"] - assert all( - "python3 -m yeto.rl.secrlenv_task_images" in script for _target, script in calls - ) - - calls.clear() - plan_path = _write_plan(tmp_path / "plan.json", plan) - ssh_harness.status(plan_path) - daemon_status_targets = [ - target for target, script in calls if "secrlenv_daemon=running" in script - ] - assert daemon_status_targets == ["alice@a0", "alice@b0"] - - calls.clear() - ssh_harness._stop_secrlenv_daemons(plan) - assert [target for target, _script in calls] == ["alice@a0", "alice@b0"] - - plan["secrlenv_daemon"]["placement"] = "unknown" - with pytest.raises(HarnessError, match="daemon contract"): - ssh_harness._validate_plan(plan) - with pytest.raises(HarnessError, match="invalid placement"): - ssh_harness._secrlenv_daemon_hosts(plan) - - -def test_legacy_secrlenv_plan_without_placement_retains_all_host_semantics( - monkeypatch, -): - plan = _plan() - _enable_secrlenv_agent(plan) - plan["secrlenv_daemon"] = _secrlenv_daemon_contract() - ssh_harness._validate_plan(plan) - - assert ssh_harness._secrlenv_daemon_hosts(plan) == ssh_harness._all_hosts(plan) - assert "SECRLENV_DAEMON_URL=" in _node_start_script(plan, 0, 1) - - calls = [] - monkeypatch.setattr( - ssh_harness, - "_ssh", - lambda plan, target, script, **kwargs: calls.append((target, script)), - ) - ssh_harness._start_secrlenv_daemons(plan) - assert [target for target, _script in calls] == ssh_harness._all_hosts(plan) - - calls.clear() - ssh_harness._stop_secrlenv_daemons(plan) - assert [target for target, _script in calls] == ssh_harness._all_hosts(plan) - - -def test_start_attests_secrlenv_daemons_before_host_gpu_setup(tmp_path, monkeypatch): - plan = _plan() - _enable_secrlenv_agent(plan) - plan["secrlenv_daemon"] = _secrlenv_daemon_contract() - path = _write_plan(tmp_path / "plan.json", plan) - events = [] - - monkeypatch.setattr(ssh_harness, "deploy", lambda _path: events.append("deploy")) - monkeypatch.setattr( - ssh_harness, - "_start_secrlenv_daemons", - lambda _plan: events.append("daemon"), - ) - monkeypatch.setattr(ssh_harness, "_check_island_cards", lambda _plan: None) - monkeypatch.setattr(ssh_harness, "_host_setup_script", lambda *_args: "setup") - monkeypatch.setattr( - ssh_harness, - "_ssh", - lambda *_args, **_kwargs: events.append("host_setup"), - ) - monkeypatch.setattr(ssh_harness, "_build_syncer", lambda _plan: None) - monkeypatch.setattr(ssh_harness, "_start_syncer", lambda _plan: None) - monkeypatch.setattr(ssh_harness, "_wait_for_syncer", lambda _plan: None) - monkeypatch.setattr(ssh_harness, "_start_island", lambda *_args: None) - - ssh_harness.start(path) - - assert events[:3] == ["deploy", "daemon", "host_setup"] - - def test_start_sets_up_a_dedicated_syncer_without_gpu_requirements( tmp_path, monkeypatch ): @@ -2609,7 +1704,6 @@ def test_start_sets_up_a_dedicated_syncer_without_gpu_requirements( calls = [] monkeypatch.setattr(ssh_harness, "deploy", lambda _path: None) - monkeypatch.setattr(ssh_harness, "_start_secrlenv_daemons", lambda _plan: None) monkeypatch.setattr(ssh_harness, "_check_island_cards", lambda _plan: None) monkeypatch.setattr( ssh_harness, "_host_setup_script", lambda *_args: "learner-setup" diff --git a/tests/test_rl_synthetic_segments.py b/tests/test_rl_synthetic_segments.py index 3f2f5e1fa..e08075451 100644 --- a/tests/test_rl_synthetic_segments.py +++ b/tests/test_rl_synthetic_segments.py @@ -27,8 +27,8 @@ def test_labels_alternate_and_keep_gsm8k_reward(): def test_launcher_accepts_synthetic_only_with_per_sample_variant(): - assert L.codex_harness_launch(_args("cross_segment_per_sample"), environ={}) is None + assert L.check_compaction_launch(_args("cross_segment_per_sample"), environ={}) is None with pytest.raises(ValueError, match="only valid"): - L.codex_harness_launch(_args("length_adaptive"), environ={}) + L.check_compaction_launch(_args("length_adaptive"), environ={}) with pytest.raises(ValueError, match="needs compacted rollouts"): - L.codex_harness_launch(_args("cross_segment_per_sample", "yeto.rl.gsm8k_reward:score"), environ={}) + L.check_compaction_launch(_args("cross_segment_per_sample", "yeto.rl.gsm8k_reward:score"), environ={}) diff --git a/tests/test_rl_tb2_baseline.py b/tests/test_rl_tb2_baseline.py deleted file mode 100644 index 1d42c7cf7..000000000 --- a/tests/test_rl_tb2_baseline.py +++ /dev/null @@ -1,127 +0,0 @@ -"""S24 T-TB2: TB2 base-model baseline runner (CPU only; fake serving, provider and agent).""" - -from __future__ import annotations - -import asyncio -import json -from pathlib import Path -from types import SimpleNamespace - -import pytest - -from yeto.rl.eval import tb2_baseline as tb -from yeto.rl.eval.island import EvalTask - -REPO = Path(__file__).resolve().parents[1] -SPLIT = {"benchmark": "tb2", "items": [{"task_id": "a", "difficulty": "easy", "eval_bucket": "tb2-easy"}, - {"task_id": "b", "difficulty": "hard", "eval_bucket": "tb2-hard"}]} -ROWS = [{"metadata": {"task_id": "a"}}, {"metadata": {"task_id": "b"}}] - - -def test_load_tasks_requires_rows_and_known_ids(): - tasks = tb.load_tasks(SPLIT, ROWS) - assert [(t.task_id, t.difficulty) for t in tasks] == [("a", "easy"), ("b", "hard")] - assert tb.load_tasks(SPLIT, ROWS, ["b"])[0].row == ROWS[1] - with pytest.raises(ValueError, match="without a jsonl row"): - tb.load_tasks(SPLIT, ROWS[:1]) - with pytest.raises(ValueError, match="not in the split"): - tb.load_tasks(SPLIT, ROWS, ["a", "zz"]) - - -def test_shipped_train_split_has_46_tasks_with_rows(): - split = json.loads((REPO / "data/eval/tb2-train.json").read_text()) - rows = [json.loads(x) for x in (REPO / "data/tb2/tb2-train.jsonl").read_text().splitlines() if x.strip()] - assert len(tb.load_tasks(split, rows)) == 46 - - -def test_run_attempts_bounds_concurrency_and_keeps_failures(): - live, peak, seen = [0], [0], [] - - async def attempt(task, trial, *, policy_version, policy_token): - assert policy_version == tb.BASE_POLICY_VERSION and policy_token == "base" - live[0] += 1 - peak[0] = max(peak[0], live[0]) - await asyncio.sleep(0.01) - live[0] -= 1 - if task.task_id == "b" and trial == 1: - raise RuntimeError("sandbox create failed") - return {"success": task.task_id == "a", "reward": float(task.task_id == "a"), "end_reason": "completed"} - - tasks = tb.load_tasks(SPLIT, ROWS) - traces = asyncio.run(tb.run_attempts(tasks, 3, attempt, concurrency=2, on_trace=seen.append)) - assert len(traces) == len(seen) == 6 and peak[0] == 2 - infra = [t for t in traces if t["end_reason"] == "infra_error"] - assert len(infra) == 1 and "sandbox create failed" in infra[0]["error"] - - -def test_summarize_excludes_infra_and_counts_mixed_tasks(): - def t(task, ok, end="completed", diff="easy"): - return {"task_id": task, "difficulty": diff, "success": ok, "end_reason": end, "turns": 3, - "timing": {"agent_s": 10.0}} - - traces = [t("a", True), {**t("a", False), "agent_detail": {"end_reason": "CodexModelFailure: x"}}, t("b", False, diff="hard"), t("b", False, "infra_error", "hard"), - t("c", True), t("c", True)] - s = tb.summarize(traces, trials=2) - assert (s["attempts"], s["judged"], s["infra_errors"], s["passed"]) == (6, 5, 1, 3) - assert s["pass_rate"] == 0.6 and s["by_difficulty"]["hard"] == {"passed": 0, "judged": 1, "pass_rate": 0.0} - assert s["mixed_task_ids"] == ["a"] and s["tasks_any_pass"] == 2 and s["tasks_all_pass"] == 1 - assert s["end_reasons"] == {"completed": 5, "infra_error": 1} and s["agent_s"]["n"] == 5 - assert s["driver_end_reasons"] == {"None": 5, "CodexModelFailure: x": 1} - - -def test_run_from_config_writes_traces_and_summary_and_closes(tmp_path): - (tmp_path / "split.json").write_text(json.dumps(SPLIT)) - (tmp_path / "rows.jsonl").write_text("".join(json.dumps(r) + "\n" for r in ROWS)) - closed = [] - loader = SimpleNamespace(start_base=lambda: {"sglang_start_s": 1.0}, close=lambda: closed.append(1)) - - class Lease: - verifier = SimpleNamespace(evaluate=None) - - async def destroy(self): - pass - - async def evaluate(_episode): - return {"passed": True} - - Lease.verifier.evaluate = evaluate - - async def acquire(task_id, tid): - return Lease() - - async def agent(lease, task, trial, *, policy_token): - return {"episode_id": "e", "end_reason": "completed", "turns": 2, "tokens": 100} - - out = tmp_path / "out" - summary = tb.run_from_config({"split": str(tmp_path / "split.json"), "rows": str(tmp_path / "rows.jsonl"), - "out_dir": str(out), "trials": 2, "concurrency": 3, - "infer": {"base_model": "B"}}, - provider=SimpleNamespace(acquire=acquire), loader=loader, agent=agent) - assert closed == [1] and summary["pass_rate"] == 1.0 and summary["attempts"] == 4 - assert len((out / tb.TRACES).read_text().splitlines()) == 4 - assert json.loads((out / tb.SUMMARY).read_text())["serving"] == {"sglang_start_s": 1.0} - - -def test_run_from_config_closes_serving_when_start_fails(tmp_path): - (tmp_path / "split.json").write_text(json.dumps(SPLIT)) - (tmp_path / "rows.jsonl").write_text("".join(json.dumps(r) + "\n" for r in ROWS)) - closed = [] - - def boom(): - raise RuntimeError("sglang exited") - - loader = SimpleNamespace(start_base=boom, close=lambda: closed.append(1)) - with pytest.raises(RuntimeError, match="sglang exited"): - tb.run_from_config({"split": str(tmp_path / "split.json"), "rows": str(tmp_path / "rows.jsonl"), - "out_dir": str(tmp_path / "o"), "infer": {"base_model": "B"}}, provider=object(), loader=loader, agent=object()) - assert closed == [1] - - -def test_function_spec_env_pins_tb2_dir(): - from yeto.cloud.modal_tb2_baseline import BaselineFunctionSpec - - spec = BaselineFunctionSpec(app_name="a", workdir="w", codex_dir="c", tb2_dir="t", image_ref="img", - envs={"SECRLENV_MAX_TURNS": "32"}) - env = spec.container_env() - assert env["YETO_HARNESS_TB2_TASKS_DIR"] == "/root/tb2-data" and env["SECRLENV_MAX_TURNS"] == "32" - assert "PYTHONPATH" not in env and spec.gpu == "H200" diff --git a/tests/test_rl_tbench_direct_preflight.py b/tests/test_rl_tbench_direct_preflight.py deleted file mode 100644 index 638cb3e50..000000000 --- a/tests/test_rl_tbench_direct_preflight.py +++ /dev/null @@ -1,137 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -import pytest - -from yeto.rl import CODEX_OPENENV_AGENT, CODEX_OPENENV_IDENTITY_ENV -from yeto.rl import tbench_direct_preflight as preflight -from yeto.rl.codex_backend import QWEN35_08B_MODEL, QWEN35_08B_REVISION - - -def _argv() -> list[str]: - return [ - "--custom-generate-function-path", - "miles.rollout.generate_hub.agentic_tool_call.generate", - "--custom-agent-function-path", - CODEX_OPENENV_AGENT, - "--custom-rm-path", - "openenv_generate.reward_func", - "--dynamic-sampling-filter-path", - "openenv_generate.check_terminal_bench_episode", - "--input-key", - "messages", - "--tito-model", - "qwen35", - "--sao-online-recipe", - "coding", - "--max-seq-len", - "8192", - "--num-gpus-per-node", - "1", - "--sglang-mem-fraction-static", - "0.15", - "--sglang-max-total-tokens", - "393216", - "--sglang-max-mamba-cache-size", - "256", - "--sao-compaction", - "--sao-one-gpu-island", - "--colocate", - ] - - -@pytest.fixture(autouse=True) -def _environment(monkeypatch): - values = { - **CODEX_OPENENV_IDENTITY_ENV, - "YETO_CODEX_COMPACTION_ENABLED": "1", - "YETO_CODEX_COMPACTION_TRIGGER_TOKENS": "6144", - "YETO_CODEX_COMPACTION_SUMMARY_MAX_TOKENS": "1024", - "YETO_CODEX_MAX_COMPACTIONS": "3", - "OPENENV_MAX_ROLLOUT_TIME_SECONDS": "1800", - "SECRLENV_MAX_ROLLOUT_TIME_SECONDS": "1800", - } - for name, value in values.items(): - monkeypatch.setenv(name, value) - - -def _context() -> dict[str, str]: - return { - "benchmark": "terminal-bench-2.1", - "model": QWEN35_08B_MODEL, - "base_model_revision": QWEN35_08B_REVISION, - "rollout_model_revision": QWEN35_08B_REVISION, - } - - -def _runtime(): - return SimpleNamespace( - trajectory_evidence_kind="terminal-bench-2.1", - trajectory_evidence_schema_version=2, - ) - - -def test_exact_contract_attests_before_launch(tmp_path, monkeypatch): - calls = [] - monkeypatch.setattr(preflight, "_attest_adapter", lambda root: calls.append(root)) - preflight.preflight_tbench_codex_streaming( - _argv(), - sao_context=_context(), - streaming_runtime=_runtime(), - miles_root=tmp_path, - ) - assert calls == [tmp_path.resolve()] - - -@pytest.mark.parametrize( - ("old", "new", "message"), - ( - ("messages", "prompt", "input_key"), - ("qwen35", "qwen35_08b", "tito_model"), - ("8192", "16384", "max_seq_len"), - ("0.15", "0.16", "sglang_mem_fraction_static"), - ("393216", "393217", "sglang_max_total_tokens"), - ("256", "257", "sglang_max_mamba_cache_size"), - ), -) -def test_argument_drift_fails_closed(tmp_path, monkeypatch, old, new, message): - monkeypatch.setattr(preflight, "_attest_adapter", lambda _root: None) - argv = _argv() - argv[argv.index(old)] = new - with pytest.raises(ValueError, match=message): - preflight.preflight_tbench_codex_streaming( - argv, - sao_context=_context(), - streaming_runtime=_runtime(), - miles_root=tmp_path, - ) - - -def test_environment_and_model_drift_fail_before_adapter_attestation( - tmp_path, monkeypatch -): - monkeypatch.setattr( - preflight, - "_attest_adapter", - lambda _root: pytest.fail("adapter attestation ran after an earlier gate failed"), - ) - monkeypatch.setenv("OPENENV_MAX_ROLLOUT_TIME_SECONDS", "1799") - with pytest.raises(ValueError, match="OPENENV_MAX_ROLLOUT_TIME_SECONDS"): - preflight.preflight_tbench_codex_streaming( - _argv(), - sao_context=_context(), - streaming_runtime=_runtime(), - miles_root=tmp_path, - ) - - monkeypatch.setenv("OPENENV_MAX_ROLLOUT_TIME_SECONDS", "1800") - context = _context() - context["rollout_model_revision"] = "0" * 40 - with pytest.raises(ValueError, match="model identity drifted"): - preflight.preflight_tbench_codex_streaming( - _argv(), - sao_context=context, - streaming_runtime=_runtime(), - miles_root=tmp_path, - ) diff --git a/tests/test_rl_trajectory_evidence.py b/tests/test_rl_trajectory_evidence.py index 105c81c51..8253a7cab 100644 --- a/tests/test_rl_trajectory_evidence.py +++ b/tests/test_rl_trajectory_evidence.py @@ -7,7 +7,7 @@ import pytest from yeto.rl import trajectory_evidence -from yeto.rl.harness.codex import reward +from yeto.rl import secrlenv_outcome as reward from yeto.rl.trajectory_evidence import ( build_trajectory_batch_evidence, read_trajectory_batch_evidence, diff --git a/tests/test_rl_trajectory_timing.py b/tests/test_rl_trajectory_timing.py index 22fdf2113..073e90420 100644 --- a/tests/test_rl_trajectory_timing.py +++ b/tests/test_rl_trajectory_timing.py @@ -1,6 +1,5 @@ """agentic-rollout-utilization 1.3: phase clock and tool-wait wiring.""" -from types import SimpleNamespace from yeto.rl.engine.trajectory_timing import PhaseClock, phase_totals, timing_fields @@ -28,18 +27,6 @@ def test_timing_fields_pick_typed_values_from_metadata_and_agent_metrics(): assert timing_fields(meta) == {"trajectory_started_at": 1.0, "turn_tool_seconds": [1.0, 2.5]} -def test_codex_openenv_agent_reads_the_island_tool_wait_board(): - from yeto.rl import CODEX_OPENENV_AGENT - from yeto.rl.adapters.miles.elastic_wiring import LazyBoardActor - from yeto.rl.adapters.miles.entry import load_tool_wait_source - - args = SimpleNamespace(custom_generate_function_path="miles.x.generate", - custom_agent_function_path=CODEX_OPENENV_AGENT, yeto_rl_learner_id=1) - assert isinstance(load_tool_wait_source(args), LazyBoardActor) - args.custom_agent_function_path = "other.agent.run" - assert load_tool_wait_source(args) is None - - def test_engine_kv_fields_and_round_peaks(): from yeto.rl.adapters.miles.rollout import _engine_kv_capacity, _engine_kv_used from yeto.rl.engine.timeline import load_peaks, validate_load_sample diff --git a/tests/test_rlf_baseline.py b/tests/test_rlf_baseline.py deleted file mode 100644 index bfb0b3e23..000000000 --- a/tests/test_rlf_baseline.py +++ /dev/null @@ -1,484 +0,0 @@ -"""RLF Qwen3-30B-A3B baseline runner (W5, rlf-codex-rl 5.6), CPU only. - -Fake provider / fake drive / fake session-server HTTP: serving argv, budget -env, episode orchestration, trace content, flag hygiene (the flag never lands -in any persisted record), classification (success / failed / timeout / -infrastructure / rejected) and the run summary. No Modal, no GPU, no SGLang. -""" - -from __future__ import annotations - -import asyncio -import json -from pathlib import Path -from types import SimpleNamespace -from typing import Any - -import pytest - -from yeto.rl.eval import rlf_baseline as rb -from yeto.rl.codex_backend import stock_codex_backend_profile -from yeto.rl.harness.codex import codex_harness_agent as harness -from yeto.rl.harness.codex.environment import HttpTerminalEnvironment, serve_environment -from yeto.rl.rlf import modal as rlf_modal - -FLAG = "flag{rlf_baseline_cpu_test_9f2e}" -TASK = "fixture-ctf-ret2win" - - -def _run(coro): - return asyncio.run(coro) - - -def _cfg(**over) -> rb.BaselineConfig: - kwargs: dict[str, Any] = {"tasks_dir": "/nonexistent", "flags": {TASK: FLAG}, - "model_path": "/models/qwen3-30b", "start_timeout_s": 5.0} - kwargs.update(over) - return rb.BaselineConfig(**kwargs) - - -# --- serving shape ------------------------------------------------------------- - -def test_sglang_argv_is_lora_free_32k_shape(tmp_path): - local = tmp_path / "model" - local.mkdir() - cfg = _cfg(model_path=str(local)) - argv = rb.sglang_argv(cfg, python="python3", - serving={"reasoning_parser": "qwen3", "tool_call_parser": "qwen3", - "chat_template_path": "/tpl.jinja"}) - assert argv[:4] == ["python3", "-m", "sglang.launch_server", "--model-path"] - assert not any(a.startswith(("--enable-lora", "--max-lora", "--lora-")) for a in argv) - assert argv[argv.index("--context-length") + 1] == "32768" - assert argv[argv.index("--mem-fraction-static") + 1] == "0.9" - assert "--revision" not in argv # local snapshot dir: no hub revision - assert argv[argv.index("--reasoning-parser") + 1] == "qwen3" - missing = _cfg(model_path="Qwen/Qwen3-30B-A3B") # not a local path - argv2 = rb.sglang_argv(missing, serving={}) - assert argv2[argv2.index("--revision") + 1] == rb.QWEN3_30B_A3B_REVISION - - -def test_session_server_config_is_lora_free(): - cfg = _cfg() - conf = rb.session_server_config(cfg, serving={"chat_template_path": "/tpl.jinja", - "chat_template_kwargs": {"enable_thinking": True}}) - # lora_rank is an inert int (pydantic rejects None; no adapter is ever loaded) - assert isinstance(conf["lora_rank"], int) and conf["lora_adapter_path"] is None - assert conf["lora_train_only"] is False - assert conf["backend_url"] == cfg.sglang_url and conf["tito_model"] == "qwen3" - assert conf["hf_checkpoint"] == cfg.model_path and conf["use_session_server"] is True - assert conf["apply_chat_template_kwargs"] == {"enable_thinking": True} - - -def test_config_defaults_come_from_the_profile(): - cfg = rb.BaselineConfig.for_profile(tasks_dir="/t", flags={}) - profile = stock_codex_backend_profile("qwen3_30b_a3b") - assert (cfg.context_length, cfg.max_tokens, cfg.max_turns) == \ - (profile["max_seq_len"], profile["max_tokens_per_turn"], profile["max_turns"]) == \ - (32768, 8192, 48) - assert cfg.mem_fraction_static == 0.9 and cfg.episode_seconds == 1800.0 - assert cfg.start_timeout_s == 900.0 # 15 min serving budget - assert cfg.model_path.endswith("/hf/Qwen3-30B-A3B/ad44e777") - assert cfg.policy_id == f"{profile['model_identifier']}@{profile['model_revision']}" - env = cfg.worker_budget_env() - assert env == {"YETO_CODEX_CHAT_TEMPLATE": "qwen3_30b_a3b", - "SECRLENV_MAX_TURNS": "48", "YETO_CODEX_BACKEND_MAX_TOKENS": "8192"} - - -class _FakeProc: - def __init__(self): - self.killed = 0 - - def poll(self): - return None - - def terminate(self): - self.killed += 1 - - def wait(self, timeout=None): - return 0 - - -def test_serving_stack_ready_and_start_timeout(): - cfg = _cfg() - spawned = [] - stack = rb.ServingStack(cfg, popen=lambda argv, **kw: spawned.append(argv) or _FakeProc(), - http=lambda *a, **kw: (200, {}), sleep=lambda s: None, - clock=iter([0, 1, 2, 3]).__next__, serving=lambda t: {}) - info = stack.start() - assert info["sglang_start_s"] == 1 and "--context-length" in spawned[0] - assert "session.server" in spawned[1][2] - stack.close() - - # never ready within start_timeout_s -> InfrastructureError (the run aborts) - clock_state = {"t": 0.0} - - def clock(): - clock_state["t"] += 400.0 - return clock_state["t"] - - stack2 = rb.ServingStack(_cfg(), popen=lambda argv, **kw: _FakeProc(), - http=lambda *a, **kw: (None, None), sleep=lambda s: None, - clock=clock, serving=lambda t: {}) - with pytest.raises(rb.InfrastructureError, match="infrastructure"): - stack2.start() - - # process exits before ready -> InfrastructureError too - class _Dead(_FakeProc): - returncode = 3 - - def poll(self): - return 3 - - stack3 = rb.ServingStack(_cfg(), popen=lambda argv, **kw: _Dead(), - http=lambda *a, **kw: (None, None), sleep=lambda s: None, - serving=lambda t: {}) - with pytest.raises(rb.InfrastructureError, match="exited"): - stack3.start() - - -# --- fakes: provider / session server / drive ------------------------------------ - -class _FakeExec: - def __init__(self, output=""): - self.output = output - self.commands: list[str] = [] - - def __call__(self, command: str, *, timeout_s: float) -> rlf_modal.VmExecResult: - self.commands.append(command) - return rlf_modal.VmExecResult(0, self.output) - - -class _FakeLease: - def __init__(self, relay, runner, url, token, *, flag=FLAG, trajectory_id=""): - self.env_url = url - self.env_token = token - self.environment = relay - self.verifier = rlf_modal.RlfFlagVerifier(relay, flag=flag, task_id=TASK, - episode_id="ep-lease-1", - trajectory_id=trajectory_id or "traj-lease-1") - self.episode_id = "ep-lease-1" - self.trajectory_id = trajectory_id or "traj-lease-1" - self.deadline_seconds = 1800.0 - self.worker_env = None - self._runner = runner - self.destroyed = 0 - - async def destroy(self): - self.destroyed += 1 - await self._runner.cleanup() - - async def describe(self): - return "gone" - - -class FakeProvider: - """One freshly-served relay environment per acquire (the real wire, no VM).""" - - def __init__(self, *, exec_output="", flag=FLAG): - self.exec_output = exec_output - self.flag = flag - self.leases: list[_FakeLease] = [] - self.acquired: list[tuple[str, str]] = [] - - async def acquire(self, task_id: str, trajectory_id: str = "") -> _FakeLease: - self.acquired.append((task_id, trajectory_id)) - relay = rlf_modal.CtfRelayEnvironment(_FakeExec(self.exec_output)) - runner, url = await serve_environment(relay, "tok-" + trajectory_id[:8], host="127.0.0.1") - lease = _FakeLease(relay, runner, url, "tok-" + trajectory_id[:8], - flag=self.flag, trajectory_id=trajectory_id) - self.leases.append(lease) - return lease - - -class FakeSessions: - """The session server's /sessions routes; the drive never reaches chat completions.""" - - def __init__(self): - self.created = 0 - self.deleted: list[str] = [] - - def __call__(self, method: str, url: str, body: Any = None, *, timeout: float = 60.0): - if method == "POST" and url.endswith("/sessions"): - self.created += 1 - return 200, {"session_id": f"s-{self.created}"} - if method == "DELETE" and "/sessions/" in url: - self.deleted.append(url.rsplit("/", 1)[-1]) - return 200, {} - raise AssertionError(f"unexpected {method} {url}") - - -GOOD_METRICS = {"turns": 3, "tool_calls": 3, "submit_calls": 1, "parse_failures": 0, - "max_model_total_tokens": 12345, "max_seq_len_hit": 0, - "turn_completion_tokens": [100, 200, 300], - "turn_context_tokens": [1000, 2000, 3000], - "turn_tool_output_bytes": [50, 60, 70], - "end_reason": "submitted", "end_kind": "submit"} - - -def make_drive(*, submit=FLAG, commands=("cat /challenge/flag.txt",), metrics=None, - recorded=None): - """The untrusted worker: talks to the lease over the real HTTP wire, then reports.""" - - async def drive(job, trajectory_id, board, worker_env): - if recorded is not None: - recorded.append({"job": dict(job), "worker_env": dict(worker_env or {})}) - async with HttpTerminalEnvironment(job["env_url"], job["env_token"]) as env: - for command in commands: - await env.execute(job["episode_id"], command, timeout_seconds=5, - output_bytes=8192) - if submit is not None: - await env.submit(job["episode_id"], {"flag": submit}) - return {"status": "completed", "episode_id": job["episode_id"], - "metrics": dict(GOOD_METRICS if metrics is None else metrics)} - - return drive - - -ROW = {"metadata": {"task_id": TASK, "rlf_task_id": TASK, "grader_mode": "flag", - "prompt": "Connect to the target and obtain the flag."}, - "prompt": "Connect to the target and obtain the flag."} - - -# --- episode orchestration -------------------------------------------------------- - -def test_episode_success_trace_and_flag_hygiene(): - provider = FakeProvider(exec_output=f"... {FLAG} ...") # the flag appears in tool output - sessions = FakeSessions() - recorded: list[dict] = [] - trace = _run(rb.run_episode(_cfg(), provider, ROW, session_url="http://sess", - drive=make_drive(recorded=recorded), http=sessions,)) - assert trace["classification"] == "success" and trace["success"] is True - assert trace["reward"] == 1.0 and trace["end_reason"] == "completed" - # RLF identity quadruple + baseline policy identity - assert trace["rlf_task_id"] == TASK and trace["rlf_episode_id"] == "ep-lease-1" - assert trace["rlf_trajectory_id"] and trace["policy_id"].startswith("Qwen/Qwen3-30B-A3B@") - assert trace["grader_evidence"]["grader"] == "rlf-flag" - # per-turn token rows and per-tool trace events - assert [t["completion_tokens"] for t in trace["turns"]] == [100, 200, 300] - assert trace["turn_count"] == 3 and trace["tool_calls"] == 3 and trace["submit_calls"] == 1 - assert len(trace["tool_events"]) == 1 - event = trace["tool_events"][0] - assert event["exit_code"] == 0 and event["flag_seen"] is True - assert event["output_preview"] == "... *** ..." # scrubbed, never the raw flag - # the submission turn marker: 1 tool call happened before the submit - assert trace["submissions"] == [{"after_tool_calls": 1, "has_flag": True, "flag_correct": True}] - assert trace["max_model_total_tokens"] == 12345 and trace["max_seq_len_hit"] == 0 - assert trace["compactions"] == 0 - # the flag is nowhere in anything that would be persisted - assert FLAG not in json.dumps(trace) - # the job and the worker budget env reached the drive - job, worker_env = recorded[0]["job"], recorded[0]["worker_env"] - assert job["episode_id"] == "ep-lease-1" and job["max_seq_len"] == 32768 - assert job["request_kwargs"]["max_tokens"] == 8192 and job["driver"] == "stock" - assert job["base_url"] == "http://sess/sessions/s-1" - assert worker_env["SECRLENV_MAX_TURNS"] == "48" - assert worker_env["YETO_CODEX_BACKEND_MAX_TOKENS"] == "8192" - assert worker_env["YETO_CODEX_CHAT_TEMPLATE"] == "qwen3_30b_a3b" - # teardown: session deleted, lease destroyed exactly once - assert sessions.deleted == ["s-1"] and provider.leases[0].destroyed == 1 - - -def test_episode_wrong_flag_is_failed_with_zero_reward(): - provider = FakeProvider() - trace = _run(rb.run_episode(_cfg(), provider, ROW, session_url="http://sess", - drive=make_drive(submit="flag{nope}"), http=FakeSessions(),)) - assert trace["classification"] == "failed" and trace["success"] is False - assert trace["reward"] == 0.0 - assert trace["submissions"][0]["flag_correct"] is False - assert "flag{nope}" not in json.dumps(trace) # submitted values are never recorded - - -def test_episode_hard_deadline_times_out_and_skips_grading(): - provider = FakeProvider() - - async def slow_drive(job, trajectory_id, board, worker_env): - await asyncio.sleep(30) - - cfg = _cfg(episode_seconds=0.2) - trace = _run(rb.run_episode(cfg, provider, ROW, session_url="http://sess", - drive=slow_drive, http=FakeSessions())) - assert trace["status"] == "timeout" and trace["end_reason"] == "timed_out" - assert trace["classification"] == "failed" and trace["reward"] == 0.0 - assert provider.leases[0].verifier.evaluations == 0 # a timeout never reaches the grader - assert provider.leases[0].destroyed == 1 - - -def test_episode_lease_deadline_tightens_the_hard_limit(): - provider = FakeProvider() - seen = {} - - async def slow_drive(job, trajectory_id, board, worker_env): - await asyncio.sleep(30) - - cfg = _cfg(episode_seconds=1800.0) - orig_acquire = provider.acquire - - async def acquire(task_id, trajectory_id=""): - lease = await orig_acquire(task_id, trajectory_id) - lease.deadline_seconds = 0.2 # the lease's own deadline wins over 30 min - return lease - - provider.acquire = acquire - trace = _run(rb.run_episode(cfg, provider, ROW, session_url="http://sess", - drive=slow_drive, http=FakeSessions())) - assert trace["status"] == "timeout" - - -def test_episode_harness_error_is_infrastructure(): - provider = FakeProvider() - - async def broken_drive(job, trajectory_id, board, worker_env): - raise harness.CodexHarnessError("worker exited without a result") - - trace = _run(rb.run_episode(_cfg(), provider, ROW, session_url="http://sess", - drive=broken_drive, http=FakeSessions())) - assert trace["classification"] == "infrastructure" and trace["reward"] is None - assert "worker exited" in trace["error"] - assert provider.leases[0].destroyed == 1 - - -def test_episode_acquire_failure_is_infrastructure_and_never_graded(): - class DenyProvider: - async def acquire(self, task_id, trajectory_id=""): - raise rlf_modal.OrchestrationError(f"no trusted flag configured, FLAG={FLAG}") - - trace = _run(rb.run_episode(_cfg(), DenyProvider(), ROW, session_url="http://sess", - drive=make_drive(), http=FakeSessions())) - assert trace["classification"] == "infrastructure" - assert trace.get("rlf_episode_id") is None # no lease, no identity - assert "acquire" in trace["error"] and FLAG not in json.dumps(trace) # scrubbed - - -def test_episode_malformed_grade_is_rejected_not_failed(): - provider = FakeProvider() - - async def drive(job, trajectory_id, board, worker_env): - return {"status": "completed", "episode_id": job["episode_id"], "metrics": {}} - - orig_acquire = provider.acquire - - async def acquire(task_id, trajectory_id=""): - lease = await orig_acquire(task_id, trajectory_id) - lease.verifier = SimpleNamespace(evaluate=lambda episode_id: _bad_grade()) - return lease - - async def _bad_grade(): - return {"rlf_grade": {"status": "completed"}} # malformed: no reward/identity - - provider.acquire = acquire - trace = _run(rb.run_episode(_cfg(), provider, ROW, session_url="http://sess", - drive=drive, http=FakeSessions())) - assert trace["classification"] == "rejected" and trace["reward"] is None - assert "grade rejected" in trace["error"] - - -# --- run loop + summary ----------------------------------------------------------- - -def _catalog_package(root: Path) -> None: - pkg = root / TASK - (pkg / "public").mkdir(parents=True) - (pkg / "public" / "README.md").write_text("# ret2win\nConnect and obtain the flag.") - (pkg / "Dockerfile").write_text("FROM scratch\n") - (pkg / "task.yaml").write_text(f"task_id: {TASK}\ngrader_mode: flag\ntarget_port: 31337\n") - (pkg / "MANIFEST.json").write_text(json.dumps({"task_id": TASK, "grader_mode": "flag"})) - - -def test_run_from_config_end_to_end(tmp_path): - _catalog_package(tmp_path) - provider = FakeProvider() - config = {"tasks_dir": str(tmp_path), "flags": {TASK: FLAG}, - "evidence": {"runtime_quality_gate": {"healthy": True}}, - "out_dir": str(tmp_path / "out"), "serve": "external", - "session_url": "http://sess", "baseline": {"repeat": 2, "model_path": "/m"}} - result = rb.run_from_config(config, provider=provider, drive=make_drive(), - http=FakeSessions(), progress=None) - assert result["ok"] is True - summary = result["summary"] - assert summary["episodes"] == 2 and summary["evaluated"] == 2 - assert summary["success_count"] == 2 and summary["success_rate"] == 1.0 - assert summary["reward_distribution"] == {"mean": 1.0, "counts": {"1.0": 2}} - assert summary["end_reasons"] == {"completed": 2} - assert summary["max_seq_len_hit"] == {"count": 0, "rate": 0.0} - assert summary["per_task"][TASK]["successes"] == 2 - assert summary["policy_id"].startswith("Qwen/Qwen3-30B-A3B@") - traces = [json.loads(line) for line in - Path(result["paths"]["traces"]).read_text().splitlines()] - assert len(traces) == 2 - assert {t["rlf_trajectory_id"] for t in traces} and \ - len({t["rlf_episode_id"] for t in traces}) == 1 # fake provider reuses its lease id - assert len({t["rlf_trajectory_id"] for t in traces}) == 2 # but trajectories are unique - on_disk = Path(result["paths"]["traces"]).read_text() + \ - Path(result["paths"]["summary"]).read_text() - assert FLAG not in on_disk - - -def test_run_from_config_serving_timeout_aborts_as_infrastructure(tmp_path, monkeypatch): - _catalog_package(tmp_path) - - class NeverReady: - def __init__(self, cfg): - self.cfg = cfg - - def start(self): - raise rb.InfrastructureError("sglang not ready after 900.0 s: infrastructure") - - def close(self): - pass - - monkeypatch.setattr(rb, "ServingStack", NeverReady) - config = {"tasks_dir": str(tmp_path), "flags": {TASK: FLAG}, - "evidence": {"runtime_quality_gate": {"healthy": True}}, - "out_dir": str(tmp_path / "out"), "serve": "start", - "baseline": {"model_path": "/m"}} - result = rb.run_from_config(config, provider=FakeProvider(), drive=make_drive(), - progress=None) - assert result["ok"] is False - aborted = result["summary"]["aborted"] - assert aborted["classification"] == "infrastructure" and aborted["stage"] == "serving" - assert result["summary"]["episodes"] == 0 - assert Path(result["paths"]["summary"]).is_file() - - -def test_summarize_mixed_classifications(): - traces = [ - {"task_id": "a", "classification": "success", "success": True, "reward": 1.0, - "end_reason": "completed", "max_seq_len_hit": 0, "turn_count": 10, "tool_calls": 9, - "max_model_total_tokens": 20000, "compactions": 0}, - {"task_id": "a", "classification": "failed", "success": False, "reward": 0.0, - "end_reason": "max_seq_len", "max_seq_len_hit": 1, "turn_count": 48, "tool_calls": 47, - "max_model_total_tokens": 32768, "compactions": 2}, - {"task_id": "b", "classification": "infrastructure", "success": False, "reward": None}, - {"task_id": "b", "classification": "rejected", "success": False, "reward": None}, - ] - summary = rb.summarize(traces, policy_id="p@1") - assert summary["episodes"] == 4 and summary["evaluated"] == 2 - assert summary["infrastructure"] == 1 and summary["rejected"] == 1 - assert summary["success_rate"] == 0.5 and summary["success_count"] == 1 - assert summary["reward_distribution"]["counts"] == {"0.0": 1, "1.0": 1} - assert summary["end_reasons"] == {"completed": 1, "max_seq_len": 1} - assert summary["max_seq_len_hit"] == {"count": 1, "rate": 0.5} - assert summary["compactions_total"] == 2 - assert summary["turns"] == {"mean": 29.0, "max": 48.0} - assert summary["per_task"]["a"]["success_rate"] == 0.5 - assert summary["per_task"]["b"]["infrastructure"] == 1 - # an all-infrastructure run has a null success rate, never a fake 0 - empty = rb.summarize([traces[2]]) - assert empty["success_rate"] is None and empty["reward_distribution"]["mean"] is None - - -def test_trace_recorder_bounds_previews_and_scrubs(): - relay = rlf_modal.CtfRelayEnvironment(_FakeExec(output=(FLAG + "x" * 100))) - - async def go(): - recorder = rb.TraceRecorder(relay, secrets=(FLAG,), preview_bytes=20).attach() - out = await relay.execute("ep-1", f"echo {FLAG}", timeout_seconds=5, output_bytes=8192) - await relay.submit("ep-1", {"flag": FLAG}) - return recorder, out - - recorder, _ = _run(go()) - event = recorder.tool_events[0] - assert event["flag_seen"] is True and len(event["output_preview"]) <= 20 - assert event["command"] == "echo ***" - assert recorder.submissions[0]["flag_correct"] is True - assert FLAG not in json.dumps({"tool_events": recorder.tool_events, - "submissions": recorder.submissions}) diff --git a/tests/test_rlf_bridge.py b/tests/test_rlf_bridge.py deleted file mode 100644 index e30891255..000000000 --- a/tests/test_rlf_bridge.py +++ /dev/null @@ -1,33 +0,0 @@ -import json -from pathlib import Path -import pytest -from yeto.rl.rlf import * - -def package(tmp_path, mode='flag'): - (tmp_path/'public').mkdir(); (tmp_path/'public'/'README').write_text('x') - (tmp_path/'task.yaml').write_text(f'task_id: demo\ngrader_mode: {mode}\n') - (tmp_path/'MANIFEST.json').write_text(json.dumps({'task_id':'demo','grader_mode':mode})) - -def test_admission_requires_runtime_evidence(tmp_path): - package(tmp_path) - with pytest.raises(AdmissionError): admit_task_package(tmp_path) - a=admit_task_package(tmp_path,evidence={'runtime_quality_gate':{'healthy':True}},image_digest='sha256:x') - assert a.task_id=='demo' and a.grader_mode=='flag' - -def test_agentflow_trace_cannot_admit(tmp_path): - package(tmp_path) - with pytest.raises(AdmissionError, match='AgentFlow'): admit_task_package(tmp_path,evidence={'runtime_quality_gate':1,'agentflow_node_output':'x'}) - -def test_finish_preserves_identity_and_cleanup(): - l=EpisodeLease('demo','traj-1',secret='private') - r=finish_episode(l,TrustedGrade('completed',1,True,{'flag':'matched'})) - assert r.value==1 and r.metadata['success'] and l.closed and l.secret is None - -def test_gate_rejects_constant_reward_and_accepts_update(): - with pytest.raises(AdmissionError): rl_gate([{'reward':0,'advantage':0,'policy_update':False}]) - out=rl_gate([{'reward':0,'advantage':-.1,'metadata':{'success':False},'policy_update':True},{'reward':1,'advantage':.2,'metadata':{'success':True}}]) - assert out['success_count']==1 and out['advantage_abs_sum']>.0 - -def test_gate_requires_success_count(): - with pytest.raises(AdmissionError, match='success_count_positive'): - rl_gate([{'reward':0,'advantage':-.5,'metadata':{'success':False},'policy_update':True},{'reward':1,'advantage':.5,'metadata':{'success':False},'policy_update':True}]) diff --git a/tests/test_rlf_compaction_profile.py b/tests/test_rlf_compaction_profile.py deleted file mode 100644 index 593c2d8af..000000000 --- a/tests/test_rlf_compaction_profile.py +++ /dev/null @@ -1,7 +0,0 @@ -from yeto.rl.harness.codex.compaction_bridge import qwen3_cyber_compaction_config - -def test_qwen3_cyber_profile_reserves_context_and_segments(): - c=qwen3_cyber_compaction_config() - assert c.context_budget==24576 - assert c.t_comp==4096 - assert c.max_compactions==3 diff --git a/tests/test_rlf_env_lease.py b/tests/test_rlf_env_lease.py deleted file mode 100644 index ec54fc104..000000000 --- a/tests/test_rlf_env_lease.py +++ /dev/null @@ -1,446 +0,0 @@ -"""RLF CTF environment lease (W4, breakpoints 3+5), CPU only. - -Fake VM exec / fake verifier: orchestration order, flag hygiene (never in -stdout/logs/exceptions), fail-closed grade parsing, idempotent teardown, and -the finish_trusted dual path (RLF grade vs byte-identical TB2 fallback). -No Modal, no GPU, no nested Docker. -""" - -from __future__ import annotations - -import asyncio -import base64 -import io -import json -import tarfile -from pathlib import Path -from types import SimpleNamespace -from typing import Any - -import pytest - -from yeto.rl.rlf import modal as rlf_modal -from yeto.rl.rlf.contracts import TrustedGrade -from yeto.rl.rlf.preflight import RolloutRejected -from yeto.rl.harness.codex import codex_openenv_agent_function as adapter -from yeto.rl import tbench_outcome - -FLAG = "flag{trusted_vm_acceptance_7b3d}" -TASK = "fixture-ctf-ret2win" - - -# --- fakes ------------------------------------------------------------------- - -class FakeExec: - """Scripted VmExec: records every command, replays canned results.""" - - def __init__(self, responses: dict[str, tuple[int, str]] | None = None): - self.commands: list[str] = [] - self.responses = responses or {} - - def __call__(self, command: str, *, timeout_s: float) -> rlf_modal.VmExecResult: - self.commands.append(command) - for needle, (rc, out) in self.responses.items(): - if needle in command: - return rlf_modal.VmExecResult(rc, out) - return rlf_modal.VmExecResult(0, "") - - -def healthy_exec() -> FakeExec: - return FakeExec({ - "dockerd --host": (0, ""), - " info ": (0, "Server: Docker"), - " build --network host": (0, "Successfully built"), - " run -d ": (0, "c0ffee123\n"), - " inspect -f ": (0, "running healthy\n"), - "rm -rf": (0, ""), - "printf %s": (0, ""), - "base64 -d": (0, ""), - }) - - -def make_package(root: Path) -> Path: - pkg = root / TASK - (pkg / "public").mkdir(parents=True) - (pkg / "public" / "README.md").write_text("# ret2win\nConnect and obtain the flag.") - (pkg / "public" / "challenge.c").write_text("int main(){}") - (pkg / "Dockerfile").write_text("FROM scratch\n") - (pkg / "entrypoint.sh").write_text("#!/bin/sh\n") - (pkg / "task.yaml").write_text(f"task_id: {TASK}\ngrader_mode: flag\ntarget_port: 31337\n") - (pkg / "private").mkdir() - (pkg / "private" / "solution.py").write_text(f"print({FLAG!r})") - (pkg / "MANIFEST.json").write_text(json.dumps({"task_id": TASK})) - return pkg - - -def ctf_task(root: Path) -> rlf_modal.RlfCtfTask: - return rlf_modal.resolve_ctf_task(TASK, root, {TASK: FLAG}) - - -# --- orchestration (breakpoint 3) -------------------------------------------- - -def test_orchestration_order_and_flag_env_injection(tmp_path): - make_package(tmp_path) - fake = healthy_exec() - orch = rlf_modal.CtfOrchestration(fake, ctf_task(tmp_path), sleep=lambda s: None, - auxiliary_containers=("rlf-ctf-solver", "rlf-ctf-grader")) - cid = orch.up() - assert cid == "c0ffee123" - cmds = fake.commands - - def idx(needle): - return next(i for i, c in enumerate(cmds) if needle in c) - - assert idx("dockerd --host=unix:///run/docker/docker.sock --iptables=false --bridge=none") \ - < idx(" info ") < idx("build --network host") < idx("run -d") < idx("inspect -f") - run_cmd = cmds[idx("run -d")] - import shlex as _sh - assert f"-e FLAG={_sh.quote(FLAG)}" in run_cmd and "--network host" in run_cmd - # the flag crosses onto the VM exactly once, on the target run command - assert sum(FLAG in c for c in cmds) == 1 - # dockerd was reset before launch - assert "rm -rf /var/lib/docker /run/docker" in cmds[0] - - -def test_flag_never_in_logs_or_exceptions(tmp_path): - make_package(tmp_path) - fake = healthy_exec() - fake.responses[" build --network host"] = (1, f"build exploded, FLAG={FLAG} leaked by docker") - orch = rlf_modal.CtfOrchestration(fake, ctf_task(tmp_path), sleep=lambda s: None) - with pytest.raises(rlf_modal.OrchestrationError) as excinfo: - orch.up() - assert FLAG not in str(excinfo.value) and FLAG not in repr(excinfo.value) - assert "***" in str(excinfo.value) - - fake2 = healthy_exec() - fake2.responses[" run -d "] = (1, f"docker: {FLAG} boom") - orch2 = rlf_modal.CtfOrchestration(fake2, ctf_task(tmp_path), sleep=lambda s: None) - with pytest.raises(rlf_modal.OrchestrationError) as excinfo2: - orch2.up() - assert FLAG not in str(excinfo2.value) and FLAG not in repr(excinfo2.value) - - -def test_stage_package_excludes_private_and_admission_metadata(tmp_path): - pkg = make_package(tmp_path) - blob = base64.b64decode(rlf_modal.package_tar_b64(pkg)) - with tarfile.open(fileobj=io.BytesIO(blob), mode="r:gz") as tar: - names = tar.getnames() - assert "Dockerfile" in names and "public/challenge.c" in names - assert not any("private" in n or "MANIFEST" in n or "task.yaml" in n for n in names) - assert FLAG.encode() not in blob - - -def test_health_wait_fails_fast_on_exited_container(tmp_path): - make_package(tmp_path) - fake = healthy_exec() - fake.responses[" inspect -f "] = (0, "exited\n") - orch = rlf_modal.CtfOrchestration(fake, ctf_task(tmp_path), sleep=lambda s: None) - with pytest.raises(rlf_modal.OrchestrationError, match="exited"): - orch.up() - - -def test_teardown_is_idempotent_and_ordered(tmp_path): - make_package(tmp_path) - fake = healthy_exec() - orch = rlf_modal.CtfOrchestration(fake, ctf_task(tmp_path), sleep=lambda s: None, - auxiliary_containers=("rlf-ctf-solver", "rlf-ctf-grader")) - orch.up() - orch.down() - orch.down() # second call is a no-op - rm = [c for c in fake.commands if "rm -f" in c and "docker -H" in c] - assert len(rm) == 3 # solver, grader, target — exactly once each - assert "rm -f rlf-ctf-solver " in rm[0] and "rm -f rlf-ctf-grader " in rm[1] \ - and "rm -f rlf-ctf-target " in rm[2] - - -def test_failed_up_cleans_containers(tmp_path): - make_package(tmp_path) - fake = healthy_exec() - fake.responses[" inspect -f "] = (0, "running starting\n") - orch = rlf_modal.CtfOrchestration(fake, ctf_task(tmp_path), sleep=lambda s: None, - health_attempts=2) - with pytest.raises(rlf_modal.OrchestrationError, match="healthy"): - orch.up() - assert any("rm -f" in c and "rlf-ctf-target" in c for c in fake.commands) - - -def test_resolve_ctf_task_fails_closed(tmp_path): - make_package(tmp_path) - with pytest.raises(rlf_modal.OrchestrationError, match="no trusted flag"): - rlf_modal.resolve_ctf_task(TASK, tmp_path, {}) - with pytest.raises(rlf_modal.OrchestrationError, match="no task.yaml"): - rlf_modal.resolve_ctf_task("nope", tmp_path, {"nope": FLAG}) - pkg = make_package(tmp_path / "hacked") - (pkg / "task.yaml").write_text(f"task_id: {TASK}\ngrader_mode: hacked\n") - with pytest.raises(rlf_modal.OrchestrationError, match="flag only"): - rlf_modal.resolve_ctf_task(TASK, tmp_path / "hacked", {TASK: FLAG}) - - -# --- relay + verifier ---------------------------------------------------------- - -def _run(coro): - return asyncio.run(coro) - - -def test_flag_verifier_grades_submission_and_never_leaks(tmp_path): - relay = rlf_modal.CtfRelayEnvironment(FakeExec()) - verifier = rlf_modal.RlfFlagVerifier(relay, flag=FLAG, task_id=TASK, - episode_id="ep-1", trajectory_id="traj-1") - - async def go(): - # no submission -> fail closed, error_class set - empty = await verifier.evaluate("ep-1") - assert empty["passed"] is False - assert empty["rlf_grade"]["error_class"] == "missing_submission" - # wrong flag -> 0 reward - await relay.submit("ep-1", {"flag": "flag{wrong}"}) - wrong = await verifier.evaluate("ep-1") - assert wrong["passed"] is False and wrong["rlf_grade"]["reward"] == 0.0 - # right flag -> 1 reward - relay2 = rlf_modal.CtfRelayEnvironment(FakeExec()) - verifier2 = rlf_modal.RlfFlagVerifier(relay2, flag=FLAG, task_id=TASK, - episode_id="ep-2", trajectory_id="traj-1") - await relay2.submit("ep-2", {"flag": FLAG}) - right = await verifier2.evaluate("ep-2") - assert right["passed"] is True and right["rlf_grade"]["reward"] == 1.0 - ev = right["rlf_grade"]["grader_evidence"] - assert (ev["rlf_task_id"], ev["rlf_episode_id"], ev["rlf_trajectory_id"]) == (TASK, "ep-2", "traj-1") - # the flag never appears in any payload/log/evidence - blob = json.dumps(right) - assert FLAG not in blob - # verifier is bound to its episode - with pytest.raises(ValueError, match="bound"): - await verifier2.evaluate("ep-other") - - _run(go()) - - -def test_relay_binds_one_episode_and_executes_on_vm(): - fake = FakeExec({"uname": (0, "Linux vm\n")}) - relay = rlf_modal.CtfRelayEnvironment(fake) - - async def go(): - out = await relay.execute("ep-1", "uname", timeout_seconds=5, output_bytes=100) - assert out["exit_code"] == 0 and "Linux vm" in out["output"] - assert any("uname" in c and "timeout -k 2 5" in c for c in fake.commands) - with pytest.raises(Exception, match="unknown episode"): - await relay.execute("ep-2", "uname", timeout_seconds=5, output_bytes=100) - await relay.submit("ep-1", {"flag": "x"}) - with pytest.raises(Exception, match="terminal"): - await relay.execute("ep-1", "uname", timeout_seconds=5, output_bytes=100) - - _run(go()) - - -# --- full provider lease contract (fake VM) ------------------------------------ - -class _FakeProc: - def __init__(self, result): - self._result = result - self.stdout = io.BytesIO(result.output.encode()) - self.stderr = io.BytesIO(b"") - - def wait(self): - return self._result.exit_code - - -class _FakeSandbox: - object_id = "sb-fake" - - def __init__(self, exec_fn): - self._exec = exec_fn - self.terminated = 0 - - def exec(self, *args, **kwargs): - return _FakeProc(self._exec(args[2], timeout_s=30)) - - def poll(self): - return 1 if self.terminated else None - - def terminate(self): - self.terminated += 1 - - -def _fake_vm_provider(monkeypatch, tmp_path): - make_package(tmp_path) - fake = healthy_exec() - sandbox = _FakeSandbox(fake) - monkeypatch.setattr(rlf_modal, "acquire_vm", - lambda task_id, episode_id, *, config=None, image=None, trajectory_id=None: - rlf_modal.RlfVmLease(sandbox, task_id=task_id, episode_id=episode_id, - trajectory_id=trajectory_id)) - provider = rlf_modal.RlfModalProvider( - ctf=rlf_modal.RlfCtfConfig(tasks_dir=str(tmp_path), flags={TASK: FLAG})) - return provider, fake, sandbox - - -def test_provider_full_environment_lease_contract(monkeypatch, tmp_path): - provider, fake, sandbox = _fake_vm_provider(monkeypatch, tmp_path) - - async def go(): - lease = await provider.acquire(TASK, "traj-9") - # full EnvironmentLease surface - assert lease.env_url.startswith("http://127.0.0.1:") and lease.env_token - assert lease.deadline_seconds > 0 and callable(lease.destroy) and callable(lease.describe) - assert lease.task_id == TASK and lease.episode_id and lease.trajectory_id == "traj-9" - # target orchestrated before the lease is handed out - assert any("run -d" in c and FLAG in c for c in fake.commands) - # the worker wire reaches the VM (host namespace) and grading is trusted - from yeto.rl.harness.codex.environment import HttpTerminalEnvironment - async with HttpTerminalEnvironment(lease.env_url, lease.env_token) as env: - out = await env.execute("ep-wire-1", "id", timeout_seconds=5, output_bytes=100) - assert out["exit_code"] == 0 - await env.submit("ep-wire-1", {"flag": FLAG}) - evaluation = await lease.verifier.evaluate("ep-wire-1") - assert evaluation["passed"] is True - # teardown: containers -> runner -> VM, idempotent - await lease.destroy() - await lease.destroy() - assert await lease.describe() == "gone" - - _run(go()) - assert sandbox.terminated == 1 - assert [c for c in fake.commands if "rm -f" in c and "docker -H" in c] - - -def test_provider_acquire_failure_cleans_vm(monkeypatch, tmp_path): - provider, fake, sandbox = _fake_vm_provider(monkeypatch, tmp_path) - fake.responses[" build --network host"] = (1, "nope") - with pytest.raises(rlf_modal.OrchestrationError): - _run(provider.acquire(TASK, "traj-10")) - assert sandbox.terminated == 1 # failed acquire never leaks a VM - assert any("rm -f" in c and "rlf-ctf-target" in c for c in fake.commands) - - -def test_provider_without_ctf_keeps_vm_only_lease(monkeypatch): - sandbox = _FakeSandbox(FakeExec()) - monkeypatch.setattr(rlf_modal, "acquire_vm", - lambda task_id, episode_id, *, config=None, image=None, trajectory_id=None: - rlf_modal.RlfVmLease(sandbox, task_id=task_id, episode_id=episode_id, - trajectory_id=trajectory_id)) - lease = _run(rlf_modal.RlfModalProvider().acquire(TASK, "traj-11")) - assert isinstance(lease, rlf_modal.RlfVmLease) and not hasattr(lease, "env_url") - lease.teardown() - lease.teardown() - assert sandbox.terminated == 1 - - -def test_cloud_factory_env_switch(monkeypatch, tmp_path): - from yeto.cloud import modal_reward_env as mre - - make_package(tmp_path) - monkeypatch.delenv(mre.RLF_CTF_TASKS_DIR_ENV, raising=False) - monkeypatch.delenv(mre.RLF_CTF_FLAGS_ENV, raising=False) - assert mre.rlf_modal_provider(SimpleNamespace()).ctf is None - monkeypatch.setenv(mre.RLF_CTF_TASKS_DIR_ENV, str(tmp_path)) - with pytest.raises(RuntimeError, match="FLAGS"): - mre.rlf_modal_provider(SimpleNamespace()) # fail closed without flags - monkeypatch.setenv(mre.RLF_CTF_FLAGS_ENV, json.dumps({TASK: FLAG})) - provider = mre.rlf_modal_provider(SimpleNamespace()) - assert provider.ctf is not None and provider.ctf.flags[TASK] == FLAG - - -# --- finish_trusted dual path (breakpoint 5) ----------------------------------- - -class _GradeVerifier: - def __init__(self, payload): - self.payload = payload - - async def evaluate(self, episode_id): - return self.payload - - -def _rlf_payload(passed=True): - return {"passed": passed, "log": "flag submission matched", - "rlf_grade": {"status": "completed", "reward": 1.0 if passed else 0.0, - "passed": passed, - "grader_evidence": {"grader": "rlf-flag", "rlf_task_id": TASK, - "rlf_episode_id": "ep-rlf-1", - "rlf_trajectory_id": "traj-rlf-1"}}} - - -def test_finish_trusted_rlf_grade_writes_back_episode(): - untrusted = {"status": "completed", "metrics": {"turns": 3}, "episode_id": "ep-wire"} - md = _run(adapter.finish_trusted(untrusted, _GradeVerifier(_rlf_payload()), - task_id=TASK, sample_id="s-1")) - # RLF write-back: reward/success/grader evidence + identity quadruple - assert md["reward"] == 1.0 and md["success"] is True - assert md["grader_evidence"]["grader"] == "rlf-flag" - assert (md["rlf_task_id"], md["rlf_episode_id"], md["rlf_trajectory_id"]) == \ - (TASK, "ep-rlf-1", "traj-rlf-1") - assert md["rlf_status"] == "completed" and md["exit_status"] == "completed" - # the outcome envelope still validates (tape consumers keep working) - outcome = md[tbench_outcome.OUTCOME_KEY] - assert outcome["reward"] == 1.0 and outcome["episode_id"] == "ep-wire" - assert tbench_outcome.verified_outcome(md) == (outcome, 1.0) - assert "tbench_trusted_outcome_hmac" not in md - assert md["agent_metrics"]["turns"] == 3 and "evaluate_time" in md["agent_metrics"] - - -def test_finish_trusted_malformed_grade_fails_closed(): - untrusted = {"status": "completed", "metrics": {}, "episode_id": "ep-wire"} - bad = {"rlf_grade": {"status": "completed"}} # present but malformed - with pytest.raises(RolloutRejected, match="malformed"): - _run(adapter.finish_trusted(untrusted, _GradeVerifier(bad), - task_id=TASK, sample_id="s-1")) - missing_identity = _rlf_payload() - del missing_identity["rlf_grade"]["grader_evidence"]["rlf_episode_id"] - with pytest.raises(RolloutRejected, match="identity"): - _run(adapter.finish_trusted(untrusted, _GradeVerifier(missing_identity), - task_id=TASK, sample_id="s-1")) - wrong_task = _rlf_payload() - wrong_task["rlf_grade"]["grader_evidence"]["rlf_task_id"] = "other-task" - with pytest.raises(RolloutRejected, match="does not match"): - _run(adapter.finish_trusted(untrusted, _GradeVerifier(wrong_task), - task_id=TASK, sample_id="s-1")) - - -def test_finish_trusted_tb2_fallback_byte_unchanged(): - untrusted = {"status": "completed", "metrics": {"turns": 2}, "episode_id": "ep-wire"} - tb2_payload = {"passed": True, "testsh_rc": 0, "log": "ok"} - md = _run(adapter.finish_trusted(untrusted, _GradeVerifier(tb2_payload), - task_id="tb2-task", sample_id="s-1")) - # no RLF keys at all on the fallback path - assert not any(k.startswith("rlf_") or k == "grader_evidence" for k in md) - assert "success" not in md - outcome = md[tbench_outcome.OUTCOME_KEY] - assert outcome["reward"] == 1.0 and outcome["verifier"] == tbench_outcome.TEST_SH_VERIFIER - assert tbench_outcome.verified_outcome(md) == (outcome, 1.0) - assert "tbench_trusted_outcome_hmac" not in md - # timeout path never calls the verifier and never sees RLF keys - md_timeout = _run(adapter.finish_trusted({"status": "timeout", "metrics": {}, "episode_id": "ep"}, - _GradeVerifier(tb2_payload), - task_id="tb2-task", sample_id="s-1")) - assert md_timeout[tbench_outcome.OUTCOME_KEY]["verifier"] == tbench_outcome.TIMEOUT_VERIFIER - assert not any(k.startswith("rlf_") for k in md_timeout) - - -def test_acquire_vm_idle_timeout_follows_ttl(monkeypatch): - """Sandbox idle kill must not reclaim a VM mid-generation (S21 gate-5).""" - import sys - created: dict[str, Any] = {} - - class _Sb: - object_id = "sb-x" - - def poll(self): - return None - - def terminate(self): - pass - - fake_modal = SimpleNamespace( - App=SimpleNamespace(lookup=lambda *a, **k: object()), - Image=SimpleNamespace(from_registry=lambda *a, **k: object()), - Sandbox=SimpleNamespace(create=lambda *a, **k: (created.update(k) or _Sb())), - ) - monkeypatch.setitem(sys.modules, "modal", fake_modal) - - rlf_modal.acquire_vm("t", "ep") - assert created["timeout"] == 1800 - assert created["idle_timeout"] == 1800 # not the old min(ttl, 300) - - created.clear() - cfg = rlf_modal.RlfModalVmConfig(ttl_s=600, idle_timeout_s=900) - rlf_modal.acquire_vm("t", "ep", config=cfg) - assert created["timeout"] == 600 and created["idle_timeout"] == 900 diff --git a/tests/test_rlf_gate_evidence.py b/tests/test_rlf_gate_evidence.py deleted file mode 100644 index 7270e9d91..000000000 --- a/tests/test_rlf_gate_evidence.py +++ /dev/null @@ -1,151 +0,0 @@ -"""RL gate criteria, structured gate-evidence collection, and the RLF Miles -reward wrapping contract (task 5.8 readiness fixes).""" -import asyncio -import importlib -from types import SimpleNamespace - -import pytest - -from yeto.rl.rlf import AdmissionError, rl_gate -from yeto.rl.rlf.contracts import group_advantages, policy_update_evidence -from yeto.rl.adapters.miles import rewards as miles_rewards -from yeto.rl.adapters.verl import reward_fn as verl_reward_fn -from yeto.rl.rewards.types import Trajectory - - -def _row(reward, advantage, success, policy_update=True): - return {'reward': reward, 'advantage': advantage, - 'metadata': {'success': success}, 'policy_update': policy_update} - - -# -- rl_gate criteria (spec: 非空 reward、std>0、success_count>0、adv_abs_sum>0、update) -- - -def test_gate_passes_with_success_and_variance(): - out = rl_gate([_row(0.0, -0.5, False), _row(1.0, 0.5, True)]) - assert out['passed'] and out['success_count'] == 1 and out['advantage_abs_sum'] == 1.0 - - -def test_gate_rejects_empty_rows(): - with pytest.raises(AdmissionError, match='non_empty'): - rl_gate([]) - - -def test_gate_rejects_all_zero_rewards(): - with pytest.raises(AdmissionError): - rl_gate([_row(0.0, 0.0, False), _row(0.0, 0.0, False)]) - - -def test_gate_rejects_constant_nonzero_rewards(): - with pytest.raises(AdmissionError, match='reward_std_positive'): - rl_gate([_row(1.0, -0.5, True), _row(1.0, 0.5, True)]) - - -def test_gate_rejects_zero_success_count(): - # std>0, advantage>0, policy update=true, but nobody succeeded. - with pytest.raises(AdmissionError, match='success_count_positive'): - rl_gate([_row(0.0, -0.5, False), _row(1.0, 0.5, False)]) - - -def test_gate_rejects_zero_advantage_abs_sum(): - with pytest.raises(AdmissionError, match='advantage_abs_sum_positive'): - rl_gate([_row(0.0, 0.0, False), _row(1.0, 0.0, True)]) - - -def test_gate_rejects_missing_policy_update(): - with pytest.raises(AdmissionError, match='policy_update'): - rl_gate([_row(0.0, -0.5, False, policy_update=False), - _row(1.0, 0.5, True, policy_update=False)]) - - -# -- structured evidence collection -- - -def test_group_advantages_are_reward_minus_group_mean(): - assert group_advantages([0.0, 1.0]) == [-0.5, 0.5] - assert group_advantages([]) == [] - assert sum(group_advantages([0.2, 0.7, 1.0, 0.0])) == pytest.approx(0.0) - - -def test_policy_update_evidence_detects_real_change(): - ev = policy_update_evidence({'bias': 0.0}, {'bias': 0.01}, step=3) - assert ev.updated and ev.step == 3 and ev.changed == ('bias',) - assert ev.before_digest != ev.after_digest - - -def test_policy_update_evidence_rejects_unchanged_params(): - ev = policy_update_evidence({'bias': 0.5}, {'bias': 0.5}, step=0) - assert not ev.updated and ev.before_digest == ev.after_digest and ev.changed == () - - -# -- Miles wrapping of the RLF success reward -- - -def _rlf_metadata(success=True, reward=1.0): - return {'rlf_episode_id': 'ep-0', 'rlf_task_id': 'fixture-ctf-ret2win', - 'rlf_trajectory_id': 'traj-0', 'grader_evidence': {'trusted': True}, - 'success': success, 'reward': reward} - - -def _fake_sample(metadata): - return SimpleNamespace(metadata=dict(metadata), group_index=0, index=0, - prompt='p', response='r', label=None, tokens=[1], - rollout_log_probs=[-0.1], loss_mask=[1], - status=None, reward=None) - - -@pytest.mark.parametrize('success,reward', [(True, 1.0), (False, 0.0)]) -def test_miles_wrapper_matches_neutral_rlf_reward(success, reward): - md = _rlf_metadata(success, reward) - neutral = miles_rewards.rlf_success_reward(Trajectory(metadata=dict(md))) - wrapped = getattr(miles_rewards, 'rlf_success') # module __getattr__ wraps it - sample = _fake_sample(md) - value = asyncio.run(wrapped(None, sample)) - assert value == neutral.value == reward - # Same input, same output: the adapter writes exactly the neutral metadata. - assert sample.metadata['success'] == neutral.metadata['success'] == success - - -def test_miles_wrapper_handles_sample_list(): - wrapped = getattr(miles_rewards, 'rlf_success') - samples = [_fake_sample(_rlf_metadata(True, 1.0)), _fake_sample(_rlf_metadata(False, 0.0))] - assert asyncio.run(wrapped(None, samples)) == [1.0, 0.0] - assert [s.metadata['success'] for s in samples] == [True, False] - - -def test_custom_rm_path_resolves_to_wrapped_callable(): - path = miles_rewards.miles_custom_rm_path('rlf_success') - assert path == 'yeto.rl.adapters.miles.rewards.rlf_success' - # What Miles' load_function does: split at the last dot, getattr the module. - module, _, attr = path.rpartition('.') - fn = getattr(importlib.import_module(module), attr) - assert getattr(fn, '__wrapped_neutral__') is miles_rewards.rlf_success_reward - sample = _fake_sample(_rlf_metadata(True, 1.0)) - assert asyncio.run(fn(None, sample)) == 1.0 and sample.metadata['success'] is True - - -def test_raw_neutral_name_is_not_a_valid_miles_entry(): - # The unwrapped neutral attribute takes one Trajectory; Miles calls fn(args, samples). - with pytest.raises(TypeError): - miles_rewards.rlf_success_reward(None, _fake_sample(_rlf_metadata())) - - -# -- verl path behavior unchanged: bare float rejected, RLF row validation kept -- - -def test_verl_rlf_row_requires_trusted_metadata(): - with pytest.raises(ValueError, match='missing trusted metadata'): - verl_reward_fn.compute_score_rlf('rlf', 'answer', extra_info={'metadata': {}}) - - -def test_verl_rlf_row_accepts_full_metadata(): - score = verl_reward_fn.compute_score_rlf('rlf', 'answer', - extra_info={'metadata': _rlf_metadata(True, 1.0)}) - assert score == 1.0 - - -def test_verl_rlf_row_rejects_bool_reward(): - md = _rlf_metadata(True, True) - with pytest.raises(ValueError, match='must be numeric'): - verl_reward_fn.compute_score_rlf('rlf', 'answer', extra_info={'metadata': md}) - - -def test_verl_unmapped_reward_still_refused(): - with pytest.raises(verl_reward_fn.UnmappedReward): - verl_reward_fn.function_for('not-a-reward') diff --git a/tests/test_rlf_rl_gate.py b/tests/test_rlf_rl_gate.py deleted file mode 100644 index 7c064d216..000000000 --- a/tests/test_rlf_rl_gate.py +++ /dev/null @@ -1,350 +0,0 @@ -"""RLF real-GPU RL gate runner orchestration (task 5.8): fake learner / fake -provider only — no GPU, no Miles, no torch, no Modal. - -These tests cover the orchestration (collection -> gate rows -> rl_gate -> -evidence). The injected collection can never announce a passing verdict; -only the real ``miles_island`` collector on the GPU run may do that. -""" - -from __future__ import annotations - -import json -from pathlib import Path - -import pytest - -from yeto.rl.eval.rlf_rl_gate import ( - EVIDENCE_SCHEMA, - GateConfig, - GateEvidenceError, - build_gate_rows, - collect_policy_update, - collect_policy_update_from_tape, - collect_rollout_samples, - prepare_prompt_data, - run_gate, - write_evidence, -) - -FLAGS = {"fixture-ctf-ezshell": "flag{ezshell-secret}", "fixture-ctf-ret2win": "flag{ret2win-secret}"} - - -def _sample(group, index, reward, success, task_id="fixture-ctf-ezshell"): - return {"reward": reward, - "metadata": {"success": success, "rlf_task_id": task_id, - "rlf_episode_id": f"ep-{group}-{index}", - "rlf_trajectory_id": f"traj-{group}-{index}", - "grader_evidence": {"trusted": True}}} - - -def _payload(groups): - return {"schema_version": 1, "completed_groups": groups} - - -def _passing_groups(): - # ezshell group: one success (variance + success_count); ret2win group: all fail. - return [[_sample(0, 0, 1.0, True), _sample(0, 1, 0.0, False), - _sample(0, 2, 0.0, False), _sample(0, 3, 0.0, False)], - [_sample(1, 0, 0.0, False, "fixture-ctf-ret2win"), - _sample(1, 1, 0.0, False, "fixture-ctf-ret2win")]] - - -def _update_loader(before, after): - tables = {"before": before, "after": after} - return lambda path: tables[str(path)] - - -def _update(updated=True): - before = {"lora_a": 0.5, "lora_b": -0.25} - after = {"lora_a": 0.5 + (0.01 if updated else 0.0), "lora_b": -0.25} - return collect_policy_update("before", "after", step=2, - loader=_update_loader(before, after)) - - -def _config(tmp_path): - return GateConfig(tasks_dir=str(tmp_path), flags=FLAGS) - - -# -- collection from the completed-groups checkpoint ------------------------- - -def test_collect_rollout_samples_reads_real_checkpoint_shape(): - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload(_passing_groups())) - assert len(samples) == 6 - assert samples[0] == {"group_index": 0, "index": 0, "task_id": "fixture-ctf-ezshell", - "reward": 1.0, "success": True} - assert sum(s["success"] for s in samples) == 1 - - -def test_collect_flattens_compaction_segments_sharing_one_outcome(): - group = [[_sample(0, 0, 1.0, True), _sample(0, 0, 1.0, True)], _sample(0, 1, 0.0, False)] - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload([group])) - assert [s["reward"] for s in samples] == [1.0, 0.0] - - -def test_collect_fails_closed_on_missing_success_or_reward(): - bad = _sample(0, 0, 1.0, True) - del bad["metadata"]["success"] - with pytest.raises(GateEvidenceError, match="metadata.success"): - collect_rollout_samples("ckpt", loader=lambda _p: _payload([[bad]])) - bad = _sample(0, 0, 1.0, True) - bad["reward"] = None - with pytest.raises(GateEvidenceError, match="numeric reward"): - collect_rollout_samples("ckpt", loader=lambda _p: _payload([[bad]])) - with pytest.raises(GateEvidenceError, match="no completed groups"): - collect_rollout_samples("ckpt", loader=lambda _p: _payload([])) - - -# -- advantages and policy update are computed, never hand-filled ------------ - -def test_gate_rows_compute_grpo_group_advantages(): - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload(_passing_groups())) - rows = build_gate_rows(samples, _update()) - first = [r["advantage"] for r in rows if r["metadata"]["group_index"] == 0] - assert first == [0.75, -0.25, -0.25, -0.25] # reward minus group mean - assert all(r["policy_update"] for r in rows) - - -def test_policy_update_unchanged_params_marks_rows_not_updated(): - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload(_passing_groups())) - rows = build_gate_rows(samples, _update(updated=False)) - assert not any(r["policy_update"] for r in rows) - - -# -- the five gate criteria --------------------------------------------------- - -def test_all_five_criteria_pass_but_injected_collection_never_announces(tmp_path): - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload(_passing_groups())) - update = _update() - rows = build_gate_rows(samples, update) - evidence = run_gate(rows, collector="injected", config=_config(tmp_path), update=update) - gate = evidence["gate"] - assert gate["passed"] and gate["result"]["reward_std"] > 0 - assert gate["result"]["success_count"] == 1 - assert gate["result"]["advantage_abs_sum"] > 0 - # CPU orchestration only: the 5.8 gate is never announced off the real run. - assert evidence["verdict"] is False - assert "never announced" in evidence["verdict_note"] - - -def test_all_zero_rewards_fail(): - groups = [[_sample(0, i, 0.0, False) for i in range(4)]] - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload(groups)) - evidence = run_gate(build_gate_rows(samples, _update()), collector="miles_island") - assert evidence["verdict"] is False and evidence["gate"]["passed"] is False - assert "reward_std_positive" in evidence["gate"]["error"] - assert "success_count_positive" in evidence["gate"]["error"] - assert "advantage_abs_sum_positive" in evidence["gate"]["error"] - - -def test_constant_nonzero_rewards_fail(): - groups = [[_sample(0, i, 1.0, True) for i in range(4)]] - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload(groups)) - evidence = run_gate(build_gate_rows(samples, _update()), collector="miles_island") - assert evidence["verdict"] is False - assert "reward_std_positive" in evidence["gate"]["error"] - - -def test_no_policy_update_fails_even_with_real_collector(): - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload(_passing_groups())) - update = _update(updated=False) - rows = build_gate_rows(samples, update) - evidence = run_gate(rows, collector="miles_island", update=update) - assert evidence["verdict"] is False - assert "policy_update" in evidence["gate"]["error"] - pu = evidence["policy_update_evidence"] - assert pu["updated"] is False and pu["before_digest"] == pu["after_digest"] - - -def test_real_collector_passing_run_announces_verdict(tmp_path): - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload(_passing_groups())) - update = _update() - rows = build_gate_rows(samples, update) - evidence = run_gate(rows, collector="miles_island", config=_config(tmp_path), update=update) - assert evidence["verdict"] is True and evidence["schema"] == EVIDENCE_SCHEMA - assert evidence["policy_update_evidence"]["before_digest"] != \ - evidence["policy_update_evidence"]["after_digest"] - - -# -- policy-update proof from the island event tape (代拍板 v2) ----------------- - -def _tape(tmp_path, events): - path = tmp_path / "rl-island-0.jsonl" - path.write_text("".join(json.dumps(e) + "\n" for e in events), encoding="utf-8") - return path - - -def _publication(version, digest): - return {"event": "rl_publication", "policy_version": version, - "sync/publication_payload_hash": digest} - - -def test_tape_policy_update_detects_changed_payload(tmp_path): - tape = _tape(tmp_path, [ - {"event": "rl_driver_start"}, - _publication(0, "aaa"), - "not json at all", - _publication(1, "bbb"), - _publication(2, "ccc"), - ]) - update = collect_policy_update_from_tape(tape, step=2) - assert update.updated is True - assert update.before_digest == "aaa" and update.after_digest == "ccc" - assert update.changed == ("policy_payload",) - - -def test_tape_policy_update_unchanged_payload_marks_not_updated(tmp_path): - tape = _tape(tmp_path, [_publication(0, "aaa"), _publication(1, "aaa")]) - update = collect_policy_update_from_tape(tape, step=1) - assert update.updated is False and update.changed == () - - -def test_tape_policy_update_fails_closed_on_single_version(tmp_path): - tape = _tape(tmp_path, [_publication(0, "aaa")]) - with pytest.raises(GateEvidenceError, match="published policy version"): - collect_policy_update_from_tape(tape, step=1) - - -def test_tape_policy_update_fails_closed_on_missing_hash(tmp_path): - tape = _tape(tmp_path, [_publication(0, "aaa"), {"event": "rl_publication", - "policy_version": 1}]) - with pytest.raises(GateEvidenceError, match="payload hash"): - collect_policy_update_from_tape(tape, step=1) - - -def test_tape_policy_update_passes_full_gate(tmp_path): - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload(_passing_groups())) - tape = _tape(tmp_path, [_publication(0, "aaa"), _publication(1, "bbb")]) - update = collect_policy_update_from_tape(tape, step=1) - rows = build_gate_rows(samples, update) - evidence = run_gate(rows, collector="miles_island", update=update) - assert evidence["verdict"] is True - - -# -- evidence on disk: content and flag hygiene -------------------------------- - -def test_write_evidence_scrubs_flags_and_keeps_detail(tmp_path): - config = _config(tmp_path) - samples = collect_rollout_samples("ckpt", loader=lambda _p: _payload(_passing_groups())) - rows = build_gate_rows(samples, _update()) - rows[0]["metadata"]["note"] = f"grader saw {FLAGS['fixture-ctf-ezshell']} then submitted" - evidence = run_gate(rows, collector="injected", config=config, update=_update()) - out = write_evidence(evidence, tmp_path / "evidence.json", secrets=config.flags.values()) - text = out.read_text() - for flag in FLAGS.values(): - assert flag not in text - assert "***" in text - loaded = json.loads(text) - assert loaded["schema"] == EVIDENCE_SCHEMA - assert {r["reward"] for r in loaded["rows"]} == {0.0, 1.0} - assert "flag{" not in json.dumps(loaded["config_digest"]) # digest only - - -def test_config_requires_flags_for_gate_tasks(tmp_path): - with pytest.raises(GateEvidenceError, match="no trusted flag"): - GateConfig(tasks_dir=str(tmp_path), flags={"fixture-ctf-ezshell": "flag{x}"}) - - -# -- prompt data preparation ---------------------------------------------------- - -def _package(root, task_id): - (root / "public").mkdir(parents=True) - (root / "public" / "README.md").write_text(f"# {task_id}\nObtain the flag.") - (root / "task.yaml").write_text(f"task_id: {task_id}\ngrader_mode: flag\n") - (root / "MANIFEST.json").write_text(json.dumps({"task_id": task_id, "grader_mode": "flag"})) - - -def test_prepare_prompt_data_admits_gate_tasks(tmp_path): - for task_id in FLAGS: - _package(tmp_path / task_id, task_id) - out = prepare_prompt_data(_config(tmp_path), tmp_path / "prompt.jsonl") - rows = [json.loads(line) for line in Path(out).read_text().splitlines() if line.strip()] - assert len(rows) == 2 * len(FLAGS) # repeat=2 per task - assert {r["metadata"]["rlf_task_id"] for r in rows} == set(FLAGS) - assert all(r["metadata"]["benchmark"] == "rlf" for r in rows) - - -def test_prepare_prompt_data_rejects_unadmitted_package(tmp_path): - _package(tmp_path / "fixture-ctf-ezshell", "fixture-ctf-ezshell") - from yeto.rl.rlf.contracts import AdmissionError - with pytest.raises(AdmissionError): - prepare_prompt_data(_config(tmp_path), tmp_path / "prompt.jsonl") - - -# --- rl_sample_outcome tape collection (ports engine, S21 gate-7) ----------- - -def _outcome(group, index, reward, success, task="fixture-ctf-ezshell"): - return {"event": "rl_sample_outcome", "policy_version": 0, - "group_index": group, "index": index, - "task_id": task, "reward": reward, "success": success} - - -def test_tape_sample_outcomes_collect(tmp_path): - from yeto.rl.eval.rlf_rl_gate import collect_rollout_samples_from_tape - tape = _tape(tmp_path, [ - {"event": "rl_driver_start"}, - _outcome(0, 0, 1.0, True), - _outcome(0, 1, 0.0, False), - "garbage", - _outcome(1, 0, 0.0, False, task="fixture-ctf-ret2win"), - ]) - records = collect_rollout_samples_from_tape(tape) - assert len(records) == 3 - assert records[0]["reward"] == 1.0 and records[0]["success"] is True - assert records[2]["task_id"] == "fixture-ctf-ret2win" - - -def test_tape_sample_outcomes_fail_closed(tmp_path): - from yeto.rl.eval.rlf_rl_gate import collect_rollout_samples_from_tape - from yeto.rl.eval.rlf_rl_gate import GateEvidenceError - # no outcome records at all - tape = _tape(tmp_path, [{"event": "rl_driver_start"}]) - with pytest.raises(GateEvidenceError, match="no rl_sample_outcome"): - collect_rollout_samples_from_tape(tape) - # missing success must abort, never default to False - tape = _tape(tmp_path, [_outcome(0, 0, 1.0, None)]) - with pytest.raises(GateEvidenceError, match="lacks success"): - collect_rollout_samples_from_tape(tape) - # non-numeric reward aborts - bad = _outcome(0, 0, "1.0", True) - tape = _tape(tmp_path, [bad]) - with pytest.raises(GateEvidenceError, match="numeric reward"): - collect_rollout_samples_from_tape(tape) - - -def _trajectory_reward(group, sample, reward, success, task="fixture-ctf-ezshell"): - return {"event": "rl_trajectory_reward", "rollout_id": 0, "policy_version": 0, - "group_index": group, "sample_index": sample, "trajectory_id": f"t-{sample}", - "task_id": task, "reward": reward, "success": success, "aborted": False} - - -def test_tape_trajectory_rewards_collect(tmp_path): - """rlf-codex-rl 5.8 attempt 8 lesson: on the ports engine the all-samples - hook is rollout_meta_hook (queue_completed_groups never runs); the driver - tapes its per-sample records as rl_trajectory_reward.""" - from yeto.rl.eval.rlf_rl_gate import collect_rollout_samples_from_tape - tape = _tape(tmp_path, [ - {"event": "rl_driver_start"}, - _trajectory_reward(0, 0, 1.0, True), - _trajectory_reward(0, 1, 0.0, False), - _trajectory_reward(1, 0, 0.0, False, task="fixture-ctf-ret2win"), - ]) - records = collect_rollout_samples_from_tape(tape) - assert len(records) == 3 - assert records[0]["reward"] == 1.0 and records[0]["success"] is True - assert records[0]["index"] == 0 and records[1]["index"] == 1 - assert records[2]["task_id"] == "fixture-ctf-ret2win" - - -def test_tape_trajectory_rewards_fail_closed(tmp_path): - from yeto.rl.eval.rlf_rl_gate import collect_rollout_samples_from_tape - from yeto.rl.eval.rlf_rl_gate import GateEvidenceError - # success None (no signed verdict) must abort, never default to False - tape = _tape(tmp_path, [_trajectory_reward(0, 0, 1.0, None)]) - with pytest.raises(GateEvidenceError, match="lacks success"): - collect_rollout_samples_from_tape(tape) - # reward None (non-finite) must abort - tape = _tape(tmp_path, [_trajectory_reward(0, 0, None, True)]) - with pytest.raises(GateEvidenceError, match="numeric reward"): - collect_rollout_samples_from_tape(tape) - - diff --git a/tests/test_rlf_rollout_entry.py b/tests/test_rlf_rollout_entry.py deleted file mode 100644 index 55616e59a..000000000 --- a/tests/test_rlf_rollout_entry.py +++ /dev/null @@ -1,226 +0,0 @@ -"""RLF rollout entry points (W2): catalog -> prompt data, identity producer, -Modal provider factory contract, compaction profile registration, RLF grade -interface, and RLF metadata preservation across compaction segments. - -CPU only: fake samples/fake VM, no Ray, no Miles, no Modal, no torch. -""" - -from __future__ import annotations - -import asyncio -import dataclasses -import json -from types import SimpleNamespace -from typing import Any - -import pytest - -from yeto.rl.compaction import T_COMP -from yeto.rl.harness.codex import codex_harness_agent as harness -from yeto.rl.harness.codex import codex_openenv_generate as generate_wrapper -from yeto.rl.harness.codex import compaction_bridge as cb -from yeto.rl.rlf.catalog import ( - TaskCatalog, import_catalog, prompt_rows, rlf_identity, task_statement, write_prompt_data, -) -from yeto.rl.rlf.contracts import AdmissionError, EpisodeLease, TrustedGrade -from yeto.rl.rlf.preflight import RolloutRejected, trusted_grade_from_evaluation - - -# --- helpers ----------------------------------------------------------------- - -def _package(root, task_id="fixture-ctf-ret2win"): - (root / "public").mkdir(parents=True) - (root / "public" / "README.md").write_text("# ret2win\nConnect to the target and obtain the flag.") - (root / "task.yaml").write_text(f"task_id: {task_id}\ngrader_mode: flag\n") - (root / "MANIFEST.json").write_text(json.dumps({"task_id": task_id, "grader_mode": "flag"})) - - -def _catalog(tmp_path): - _package(tmp_path / "pkg") - return import_catalog([tmp_path / "pkg"], evidence={"runtime_quality_gate": {"healthy": True}}) - - -def _sample(n_prompt, n_gen, version="pv", **meta): - tokens = list(range(n_prompt + n_gen)) - span = SimpleNamespace(version=version, abs_start=n_prompt, abs_end=n_prompt + n_gen) - return SimpleNamespace(tokens=tokens, loss_mask=[1] * n_gen, rollout_log_probs=[-0.1] * n_gen, - response_length=n_gen, weight_versions=[SimpleNamespace(spans=[span])], - metadata=dict(meta), status=None) - - -# --- breakpoint 2: catalog -> rollout prompt data ---------------------------- - -def test_catalog_prompt_rows_feed_miles_prompt_data(tmp_path): - catalog = _catalog(tmp_path) - rows = prompt_rows(catalog) - assert len(rows) == 1 - row = rows[0] - md = row["metadata"] - # harness task_identity reads metadata.task_id; the RLF producer reads rlf_task_id. - assert md["task_id"] == "fixture-ctf-ret2win" == md["rlf_task_id"] - assert md["grader_mode"] == "flag" and md["benchmark"] == "rlf" - assert md["package_digest"] == catalog.get("fixture-ctf-ret2win").package_digest - assert "obtain the flag" in md["prompt"] and row["prompt"] == md["prompt"] - out = write_prompt_data(catalog, tmp_path / "prompt.jsonl", repeat=2) - loaded = [json.loads(line) for line in out.read_text().splitlines() if line.strip()] - assert len(loaded) == 2 and loaded[0] == row - - -def test_catalog_refuses_fabricated_statement_and_unknown_task(tmp_path): - root = tmp_path / "pkg" - _package(root) - (root / "public" / "README.md").unlink() - catalog = import_catalog([root], evidence={"runtime_quality_gate": {"healthy": True}}) - with pytest.raises(AdmissionError, match="task statement"): - prompt_rows(catalog) - with pytest.raises(KeyError): - prompt_rows(catalog, ["no-such-task"]) - - -def test_catalog_episode_produces_rlf_identity(tmp_path): - catalog = _catalog(tmp_path) - lease = catalog.episode("fixture-ctf-ret2win") - assert isinstance(lease, EpisodeLease) and lease.task_id == "fixture-ctf-ret2win" - assert lease.episode_id and lease.trajectory_id - other = catalog.episode("fixture-ctf-ret2win") - assert lease.episode_id != other.episode_id # episodes are never shared - identity = rlf_identity(lease, policy_id="pv") - assert identity == {"rlf_task_id": "fixture-ctf-ret2win", "rlf_episode_id": lease.episode_id, - "rlf_trajectory_id": lease.trajectory_id, "policy_id": "pv"} - - -# --- breakpoint 1: identity producer on the real generate path --------------- - -def test_bookkeeping_produces_rlf_identity_for_catalog_rows(): - inp = SimpleNamespace(group_index=3, index=5, - metadata={"rlf_task_id": "fixture-ctf-ret2win"}) - chains = [_sample(4, 6), _sample(4, 6)] - generate_wrapper.apply_trajectory_bookkeeping(inp, chains, expected_version="pv") - md = [c.metadata for c in chains] - assert all(m["rlf_task_id"] == "fixture-ctf-ret2win" for m in md) - assert len({m["rlf_episode_id"] for m in md}) == 1 # one episode per rollout - assert len({m["rlf_trajectory_id"] for m in md}) == 1 - assert all(m["policy_id"] == "pv" for m in md) - assert all(c.status is None for c in chains) # versions match: no abort - - -def test_bookkeeping_keeps_existing_identity_and_ignores_non_rlf_rows(): - inp = SimpleNamespace(group_index=3, index=5, metadata={ - "rlf_task_id": "fixture-ctf-ret2win", "rlf_episode_id": "ep-fixed", - "rlf_trajectory_id": "traj-fixed", "policy_id": "pv-old"}) - chain = _sample(4, 6) - generate_wrapper.apply_trajectory_bookkeeping(inp, [chain], expected_version="pv") - assert chain.metadata["rlf_episode_id"] == "ep-fixed" - assert chain.metadata["rlf_trajectory_id"] == "traj-fixed" - assert chain.metadata["policy_id"] == "pv-old" # never overwritten - plain_inp = SimpleNamespace(group_index=3, index=5, metadata={"task_id": "tb2-task"}) - plain = _sample(4, 6) - generate_wrapper.apply_trajectory_bookkeeping(plain_inp, [plain], expected_version="pv") - assert not any(k in plain.metadata for k in generate_wrapper.RLF_IDENTITY_KEYS) - - -# --- breakpoint: rlf_modal_provider factory contract -------------------------- - -def test_rlf_modal_provider_matches_preflight_factory_contract(monkeypatch): - from yeto.cloud import modal_reward_env as mre - from yeto.rl.rlf import modal as rlf_modal - - class FakeSandbox: - object_id = "sb-fake" - - def terminate(self): - pass - - calls = {} - - def fake_acquire_vm(task_id, episode_id, *, config=None, image=None, trajectory_id=None): - calls.update(task_id=task_id, episode_id=episode_id, trajectory_id=trajectory_id) - return rlf_modal.RlfVmLease(FakeSandbox(), task_id=task_id, episode_id=episode_id, - trajectory_id=trajectory_id) - - monkeypatch.setattr(rlf_modal, "acquire_vm", fake_acquire_vm) - factory = mre.rlf_modal_provider - provider = factory(SimpleNamespace()) # preflight calls factory(miles_args) - assert hasattr(provider, "acquire") # resolve_environment_provider's only check - lease = asyncio.run(provider.acquire("fixture-ctf-ret2win", "traj-1")) - assert lease.task_id == "fixture-ctf-ret2win" - assert lease.episode_id and lease.trajectory_id == "traj-1" - assert calls["trajectory_id"] == "traj-1" and calls["episode_id"] == lease.episode_id - lease.teardown() # idempotent cleanup path stays intact - assert lease.closed - - -# --- compaction profile registration (5.7) ------------------------------------ - -def test_qwen3_cyber_profile_is_registered_and_selectable(): - assert "qwen3_cyber" in cb.compaction_profiles() - cfg = cb.compaction_config(24_576, env={cb.PROFILE_ENV: "qwen3_cyber"}) - assert (cfg.context_budget, cfg.t_comp, cfg.max_compactions) == (24_576, 4_096, 3) - with pytest.raises(harness.CodexHarnessError, match="context_budget"): - cb.compaction_config(32_768, env={cb.PROFILE_ENV: "qwen3_cyber"}) - with pytest.raises(harness.CodexHarnessError, match="not a registered"): - cb.compaction_config(24_576, env={cb.PROFILE_ENV: "no-such-profile"}) - - -def test_compaction_config_default_derivation_unchanged(): - cfg = cb.compaction_config(32_768, env={}) - assert cfg.context_budget == 32_768 and cfg.t_comp == T_COMP - cfg = cb.compaction_config(32_768, env={cb.T_COMP_ENV: "4096"}) - assert cfg.t_comp == 4_096 - with pytest.raises(harness.CodexHarnessError): - cb.compaction_config(32_768, env={cb.T_COMP_ENV: "x"}) - - -# --- trusted finish: RLF grade interface (breakpoint 5, interface only) ------- - -def test_trusted_grade_from_evaluation_prefers_rlf_and_falls_back(): - assert trusted_grade_from_evaluation({"passed": True, "testsh_rc": 0}) is None # TB2 HMAC fallback - payload = {"status": "completed", "reward": 1.0, "passed": True, - "grader_evidence": {"flag": "matched"}} - grade = trusted_grade_from_evaluation({"rlf_grade": payload, "log": "..."}) - assert isinstance(grade, TrustedGrade) and grade.passed and grade.reward == 1.0 - assert grade.grader_evidence == {"flag": "matched"} - direct = trusted_grade_from_evaluation(payload) # a bare full-grade mapping also parses - assert direct is not None and direct.status == "completed" - with pytest.raises(RolloutRejected, match="malformed"): - trusted_grade_from_evaluation({"rlf_grade": {"status": "completed"}}) - - -# --- RLF metadata survives compaction segment splitting ------------------------ - -@dataclasses.dataclass(frozen=True) -class _Out: - samples: Any - - -def test_compaction_segments_preserve_rlf_identity_and_reward_metadata(): - seg0_meta = { - "expected_policy_version": "pv", "session_server_id": "10.0.0.1:9", - "exit_status": "completed", - "rlf_task_id": "fixture-ctf-ret2win", "rlf_episode_id": "ep-1", - "rlf_trajectory_id": "traj-1", "policy_id": "pv", - "success": True, "reward": 1.0, "grader_evidence": {"flag": "matched"}, - generate_wrapper.SESSIONS_KEY: ["a", "b"], - "agent_metrics": {cb.METRICS_KEY: {"compactions": 1}}, - } - seg0 = _sample(10, 30, **seg0_meta) - s1 = _sample(8, 20) - inp = SimpleNamespace(sample=SimpleNamespace(group_index=7, index=11, - metadata={"expected_policy_version": "pv"})) - - async def upstream(_input): - return _Out(samples=seg0) - - async def collect(_input, router, sid): - return (([s1], {"seg": 1}) if sid == "a" else ([], {})) - - out = asyncio.run(generate_wrapper.generate(inp, upstream=upstream, collect=collect)) - assert out.samples == [seg0, s1] - for sample in out.samples: - md = sample.metadata - assert md["rlf_task_id"] == "fixture-ctf-ret2win" - assert md["rlf_episode_id"] == "ep-1" and md["rlf_trajectory_id"] == "traj-1" - assert md["policy_id"] == "pv" - assert md["success"] is True and md["grader_evidence"] == {"flag": "matched"} - assert [s.metadata["segment_index"] for s in out.samples] == [0, 1] - assert s1.metadata["seg"] == 1 diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index 9ef5c8eab..90995176e 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -251,10 +251,10 @@ def test_forward_tables_name_flags_of_both_layers(): _NON_RL_LITERALS = frozenset({ "--abbrev-ref", "--adapter-dir", "--add", "--address", "--allow-unattested-legacy-adapter", "--allow-unauthenticated-islands", "--bind", "--cap-add", "--carry-gamma", - "--checkpoint-every", "--checkpoint-path", "--codex-harness-binary", "--config-json", + "--checkpoint-every", "--checkpoint-path", "--config-json", "--create", "--dashboard-host", "--dearmor", "--depth", "--detach", "--dtype", "--entrypoint", "--expected-island-contract", "--final-ack-timeout-s", "--final-grace-s", "--follow", "--force", - "--format", "--get", "--global", "--gpus", "--head", "--ignore-installed", + "--format", "--get", "--global", "--gpus", "--head", "--include-dashboard", "--index-url", "--ipc", "--island-lease-s", "--island-scheduling-mode", "--json", "--learner-weight", "--learners", "--level", "--manifest-path", "--mark-final-checkpoint", "--master", "--matrix-merge", diff --git a/tests/test_sandbox_isolation.py b/tests/test_sandbox_isolation.py deleted file mode 100644 index b7b04f1b5..000000000 --- a/tests/test_sandbox_isolation.py +++ /dev/null @@ -1,142 +0,0 @@ -"""rl-agentic-reward-env 5.1/5.2: sandbox egress closed by default, minimal local env.""" - -from __future__ import annotations - -import json -import sys -import types -from pathlib import Path - -import pytest - -from yeto.rl.harness.codex import tb2_provider as tb2 -from yeto.rl.harness.codex.tb2_network_grants import TB2_OPEN_EGRESS_TASKS - - -def _task(tmp_path: Path, task_id: str = "t1") -> tb2.Tb2Task: - return tb2.Tb2Task(task_id=task_id, task_dir=tmp_path, docker_image="img:1", cpus=1, memory_mb=2048, - workdir="/app", agent_timeout_s=60, verifier_timeout_s=60) - - -def test_default_policy_is_closed(): - grant = tb2.NetworkPolicy().grant("anything") - assert grant.closed - assert grant.modal_kwargs() == {"block_network": True} - - -def test_grant_forms(): - policy = tb2.NetworkPolicy.from_mapping({"tasks": { - "pip": {"domains": ["pypi.org", "*.pythonhosted.org"]}, - "net": {"cidrs": ["10.0.0.0/8"]}, - "web": "open", - "off": "closed", - }}) - assert policy.grant("pip").modal_kwargs() == {"outbound_domain_allowlist": ["pypi.org", "*.pythonhosted.org"]} - assert policy.grant("net").modal_kwargs() == {"outbound_cidr_allowlist": ["10.0.0.0/8"]} - assert policy.grant("web").modal_kwargs() == {} - assert policy.grant("off").modal_kwargs() == {"block_network": True} - assert policy.grant("unlisted").modal_kwargs() == {"block_network": True} - - -@pytest.mark.parametrize("bad", [ - {"tasks": {"x": {"domains": ["*"]}}}, - {"tasks": {"x": {"cidrs": ["not-a-cidr"]}}}, - {"tasks": {"x": {"hosts": ["a"]}}}, - {"tasks": {"x": "yes"}}, - {"tasks": [], "default": "closed"}, - {"extra": 1}, -]) -def test_bad_policy_fails_closed(bad): - with pytest.raises(ValueError): - tb2.NetworkPolicy.from_mapping(bad) - - -def test_builtin_grants_every_tb2_task_open_explicitly(): - assert len(TB2_OPEN_EGRESS_TASKS) == 89 == len(set(TB2_OPEN_EGRESS_TASKS)) - policy = tb2.NetworkPolicy.builtin() - assert policy.default.closed - for task_id in ("adaptive-rejection-sampler", "build-pmars", "qemu-startup"): - assert policy.grant(task_id).open - assert policy.grant("not-a-tb2-task").closed - - -def test_policy_from_env(tmp_path, monkeypatch): - monkeypatch.delenv(tb2.NETWORK_POLICY_ENV, raising=False) - assert tb2.NetworkPolicy.from_env() == tb2.NetworkPolicy.builtin() - path = tmp_path / "p.json" - path.write_text(json.dumps({"tasks": {"pip": {"domains": ["pypi.org"]}}})) - monkeypatch.setenv(tb2.NETWORK_POLICY_ENV, str(path)) - policy = tb2.NetworkPolicy.from_env() - assert policy.grant("pip").domains == ("pypi.org",) - assert policy.grant("adaptive-rejection-sampler").closed - - -def _fake_modal(monkeypatch, calls): - mod = types.ModuleType("modal") - - class Sandbox: - @staticmethod - def create(*args, **kwargs): - calls.append(kwargs) - return types.SimpleNamespace(object_id="sb-1") - - class Image: - @staticmethod - def from_registry(name): - return types.SimpleNamespace(name=name, run_commands=lambda *a: "img") - - class App: - @staticmethod - def lookup(name, create_if_missing=False): - return "app" - - mod.Sandbox, mod.Image, mod.App = Sandbox, Image, App - monkeypatch.setitem(sys.modules, "modal", mod) - - -def test_modal_backend_passes_network_kwargs(tmp_path, monkeypatch): - calls: list[dict] = [] - _fake_modal(monkeypatch, calls) - tb2.ModalSandboxBackend().create(_task(tmp_path), "traj") - assert calls[-1]["block_network"] is True - policy = tb2.NetworkPolicy.from_mapping({"tasks": {"t1": {"domains": ["pypi.org"]}}}) - tb2.ModalSandboxBackend(network_policy=policy).create(_task(tmp_path), "traj") - assert calls[-1]["outbound_domain_allowlist"] == ["pypi.org"] - assert "block_network" not in calls[-1] - - -def test_prebaked_modal_backend_passes_network_kwargs(tmp_path, monkeypatch): - from yeto.cloud import modal_reward_env - - calls: list[dict] = [] - _fake_modal(monkeypatch, calls) - (tmp_path / "tests").mkdir() - (tmp_path / "tests" / "test.sh").write_text("#!/bin/bash\necho 1\n") - modal_reward_env.PrebakedModalSandboxBackend(prebake=False).create(_task(tmp_path), "traj") - assert calls[-1]["block_network"] is True - backend = modal_reward_env.PrebakedModalSandboxBackend(prebake=False, network_policy=tb2.NetworkPolicy.builtin()) - backend.create(_task(tmp_path, "adaptive-rejection-sampler"), "traj") - assert not {"block_network", "outbound_domain_allowlist", "outbound_cidr_allowlist"} & set(calls[-1]) - - -def test_minimal_env_drops_tokens(): - parent = {"PATH": "/x/bin", "MODAL_TOKEN_SECRET": "s", "HF_TOKEN": "h", "TBENCH_REWARD_HMAC_KEY": "k", - "WANDB_API_KEY": "w", "LANG": "en_US.UTF-8", "HOME": "/home/u"} - assert tb2.minimal_env(parent) == {"PATH": "/x/bin", "LANG": "en_US.UTF-8"} - assert tb2.minimal_env({})["PATH"] == tb2.DEFAULT_PATH - - -def test_local_sandbox_does_not_inherit_parent_env(tmp_path, monkeypatch): - monkeypatch.setenv("MODAL_TOKEN_SECRET", "secret-value") - monkeypatch.setenv("YETO_FAKE_TOKEN", "x") - backend = tb2.LocalProcessBackend(tmp_path) - sandbox = backend.create(_task(tmp_path), "traj") - try: - result = sandbox.exec("env", timeout_s=10) - finally: - sandbox.terminate() - assert result.exit_code == 0 - names = {line.split("=", 1)[0] for line in result.output.splitlines() if "=" in line} - assert "MODAL_TOKEN_SECRET" not in names and "YETO_FAKE_TOKEN" not in names - assert {"PATH", "HOME", "TB2_TESTS_DIR", "TB2_VERIFIER_LOGS_DIR"} <= names - assert sandbox.network.closed diff --git a/tests/test_secret_handling_hardening.py b/tests/test_secret_handling_hardening.py index a9def9dcb..294fde129 100644 --- a/tests/test_secret_handling_hardening.py +++ b/tests/test_secret_handling_hardening.py @@ -199,37 +199,6 @@ def test_secret_envs_are_split_out_of_plain_envs(): assert secrets == {"HF_TOKEN": "h", "WANDB_API_KEY": "w", "YETO_SANDBOX_MODAL_TOKEN_SECRET": "s"} -def test_main_modal_token_is_not_a_harness_passthrough(): - from yeto import launcher - - assert not set(launcher.HARNESS_PASSTHROUGH_ENV) & launcher.CLOUD_CREDENTIAL_ENV_NAMES - with pytest.raises(ValueError, match="YETO_SANDBOX_MODAL_TOKEN_ID"): - launcher.require_sandbox_modal_token({"MODAL_TOKEN_ID": "a", "MODAL_TOKEN_SECRET": "b"}) - launcher.require_sandbox_modal_token( - {"YETO_SANDBOX_MODAL_TOKEN_ID": "a", "YETO_SANDBOX_MODAL_TOKEN_SECRET": "b"}) - - -def test_sandbox_client_uses_only_the_sandbox_token(monkeypatch): - from yeto.rl.harness.codex import tb2_provider - - assert tb2_provider.sandbox_modal_client({"MODAL_TOKEN_ID": "a", "MODAL_TOKEN_SECRET": "b"}) is None - seen = {} - - class FakeClient: - @staticmethod - def from_credentials(token_id, token_secret): - seen["args"] = (token_id, token_secret) - return "client" - - import sys - import types - - monkeypatch.setitem(sys.modules, "modal", types.SimpleNamespace(Client=FakeClient)) - env = {"YETO_SANDBOX_MODAL_TOKEN_ID": "sid", "YETO_SANDBOX_MODAL_TOKEN_SECRET": "ssec"} - assert tb2_provider.sandbox_modal_client(env) == "client" - assert seen["args"] == ("sid", "ssec") - - # --------------------------------------------------------------------------- D4 island start-up check def test_guard_names_match_the_launcher_list(): diff --git a/tests/test_secrlenv_adapter.py b/tests/test_secrlenv_adapter.py deleted file mode 100644 index 02d81d455..000000000 --- a/tests/test_secrlenv_adapter.py +++ /dev/null @@ -1,1937 +0,0 @@ -from __future__ import annotations - -import asyncio -import copy -import json -import sys -import types -from io import BytesIO -from types import SimpleNamespace - -import pytest - -pytest.importorskip( - "yeto_miles_secrlenv", - reason="the Miles SecRLEnv adapter is an optional external integration", -) - -from yeto_miles_secrlenv import agent, reward -from yeto_miles_secrlenv import generate as secrlenv_generate -from yeto_miles_secrlenv.client import ( - EpisodeAPIError, - EpisodeClient, - EpisodeClientError, - EpisodeTransportError, - require_daemon_ready, -) - -from yeto.rl.harness.codex import reward as secrlenv_reward - - -class FakeMessage: - def __init__(self, value): - self.value = value - - def model_dump(self, *, exclude_none=True): - return copy.deepcopy(self.value) - - -class FakeCompletions: - def __init__(self, messages): - self.responses = list(messages) - self.requests = [] - - async def create(self, **kwargs): - self.requests.append(copy.deepcopy(kwargs)) - return SimpleNamespace( - choices=[SimpleNamespace(message=FakeMessage(self.responses.pop(0)))] - ) - - -class FakePolicy: - def __init__(self, messages): - self.chat = SimpleNamespace(completions=FakeCompletions(messages)) - - -class FakeEpisodeClient: - def __init__(self): - self.exec_calls = [] - self.submissions = [] - - async def execute( - self, episode_id, command, *, timeout_seconds, output_bytes - ): - self.exec_calls.append( - (episode_id, command, timeout_seconds, output_bytes) - ) - return { - "exit_code": 0, - "output": "command output", - "timed_out": False, - "truncated": False, - } - - async def submit(self, episode_id, submission): - self.submissions.append((episode_id, submission)) - return {"accepted": True} - - -def _install_run_openai(monkeypatch): - class RunPolicy: - async def close(self): - return None - - openai = types.ModuleType("openai") - openai.AsyncOpenAI = lambda **_kwargs: RunPolicy() - monkeypatch.setitem(sys.modules, "openai", openai) - - -def _run_metadata(): - return {"task_id": "CVE-2024-1234", "prompt_tier": "l2"} - - -def test_exhausted_precreate_infrastructure_503_returns_signed_admission( - monkeypatch, -): - attempts = [] - - class AdmissionClient: - def __init__(self, **_kwargs): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - async def create(self, task_id, tier): - attempts.append((task_id, tier)) - raise EpisodeAPIError(503, "infrastructure_error", "unavailable") - - _install_run_openai(monkeypatch) - monkeypatch.setattr(agent, "EpisodeClient", AdmissionClient) - monkeypatch.setattr(agent.random, "uniform", lambda *_args: 0.0) - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "a" * 48) - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "30") - - result = asyncio.run( - agent.run("http://127.0.0.1:8000/v1", None, metadata=_run_metadata()) - ) - - assert result is not None - outcome, value = reward._verified_outcome(result) - assert len(attempts) == 2 - assert outcome["schema"] == reward.ADMISSION_SCHEMA - assert outcome["status"] == reward.INFRASTRUCTURE_STATUS - assert outcome["episode_id"] is None - assert outcome[reward.INFRASTRUCTURE_503_RETRIES_KEY] == 1 - assert value == 0.0 - agent.require_no_episode_residue() - - -def test_precreate_transport_failure_remains_unsigned_and_fatal(monkeypatch): - class TransportFailingClient: - def __init__(self, **_kwargs): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - async def create(self, _task_id, _tier): - raise EpisodeTransportError("unavailable") - - _install_run_openai(monkeypatch) - monkeypatch.setattr(agent, "EpisodeClient", TransportFailingClient) - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "a" * 48) - - assert ( - asyncio.run( - agent.run( - "http://127.0.0.1:8000/v1", None, metadata=_run_metadata() - ) - ) - is None - ) - agent.require_no_episode_residue() - - -def test_precreate_infrastructure_retry_timeout_remains_unsigned(monkeypatch): - attempts = 0 - - class AdmissionTimeoutClient: - def __init__(self, **_kwargs): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - async def create(self, _task_id, _tier): - nonlocal attempts - attempts += 1 - if attempts == 1: - raise EpisodeAPIError( - 503, "infrastructure_error", "unavailable" - ) - await asyncio.Event().wait() - - _install_run_openai(monkeypatch) - monkeypatch.setattr(agent, "EpisodeClient", AdmissionTimeoutClient) - monkeypatch.setattr(agent.random, "uniform", lambda *_args: 0.0) - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "a" * 48) - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "0.01") - - result = asyncio.run( - agent.run("http://127.0.0.1:8000/v1", None, metadata=_run_metadata()) - ) - - assert result is None - assert attempts == 2 - agent.require_no_episode_residue() - - -def test_episode_finalization_cannot_be_claimed_by_abort(): - episode_id = "f" * 24 - agent._register_episode(episode_id) - try: - assert agent._claim_episode_finalization(episode_id) - assert agent._claim_driving_episodes_for_abort() == [] - finally: - agent._release_episode(episode_id) - - -def test_abort_claim_preempts_finalization_once(): - episode_id = "a" * 24 - agent._register_episode(episode_id) - try: - assert agent._claim_driving_episodes_for_abort() == [episode_id] - assert not agent._claim_episode_finalization(episode_id) - assert not agent._claim_episode_cleanup(episode_id) - finally: - agent._release_episode(episode_id) - - -def test_finalizing_episode_cannot_be_claimed_for_cancelled_cleanup(): - episode_id = "e" * 24 - agent._register_episode(episode_id) - try: - assert agent._claim_episode_finalization(episode_id) - assert not agent._claim_episode_cleanup(episode_id) - assert agent._claim_episode_finalization(episode_id) - finally: - agent._release_episode(episode_id) - - -def test_policy_completion_claims_finalization_before_abort_callback(): - episode_id = "b" * 24 - abort_claims = [] - - async def complete_policy(): - asyncio.get_running_loop().call_soon( - lambda: abort_claims.extend( - agent._claim_driving_episodes_for_abort() - ) - ) - return "completed" - - async def scenario(): - result = await agent._await_policy_and_claim_finalization( - episode_id, complete_policy() - ) - await asyncio.sleep(0) - return result - - agent._register_episode(episode_id) - try: - assert asyncio.run(scenario()) == "completed" - assert abort_claims == [] - assert agent._claim_episode_finalization(episode_id) - finally: - agent._release_episode(episode_id) - - -def test_registered_policy_task_includes_finalization_boundary(): - episode_id = "6" * 24 - abort_claims = [] - - async def complete_policy(): - asyncio.get_running_loop().call_soon( - lambda: abort_claims.extend( - agent._claim_driving_episodes_for_abort() - ) - ) - return "completed" - - async def scenario(): - policy_task = asyncio.create_task( - agent._await_policy_and_claim_finalization( - episode_id, complete_policy() - ) - ) - agent._register_episode(episode_id, policy_task) - result = await asyncio.wait_for(policy_task, timeout=1.0) - await asyncio.sleep(0) - return result - - try: - assert asyncio.run(scenario()) == "completed" - assert abort_claims == [] - assert agent._claim_episode_finalization(episode_id) - finally: - agent._release_episode(episode_id) - - -def test_cancelled_run_can_transfer_its_drained_boundary_to_cleanup(): - episode_id = "5" * 24 - - async def scenario(): - policy_task = asyncio.create_task( - agent._await_policy_and_claim_finalization( - episode_id, asyncio.sleep(0, result="completed") - ) - ) - agent._register_episode(episode_id, policy_task) - await policy_task - assert not agent._claim_episode_cleanup(episode_id) - assert agent._claim_cancelled_policy_cleanup(episode_id, policy_task) - - try: - asyncio.run(scenario()) - assert agent._EPISODE_PHASES[episode_id] == "aborting" - finally: - agent._release_episode(episode_id) - - -def test_run_cancellation_drains_evaluation_before_cleanup_close(monkeypatch): - episode_id = "3" * 24 - evaluation_started = None - evaluation_live = False - close_observations = [] - - class RunClient: - def __init__(self, **_kwargs): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, _exc_type, _exc, _traceback): - return False - - async def create(self, _task_id, _tier): - return {"episode_id": episode_id, "prompt": "task"} - - async def evaluate(self, value): - assert value == episode_id - nonlocal evaluation_live - evaluation_live = True - evaluation_started.set() - try: - await asyncio.Event().wait() - finally: - evaluation_live = False - - async def close(self, value): - assert value == episode_id - close_observations.append(evaluation_live) - return {} - - async def complete_policy(*_args, **_kwargs): - return "completed" - - async def scenario(): - nonlocal evaluation_started - evaluation_started = asyncio.Event() - run_task = asyncio.create_task( - agent.run("http://127.0.0.1:8000/v1", None, metadata=_run_metadata()) - ) - await evaluation_started.wait() - run_task.cancel() - with pytest.raises(asyncio.CancelledError): - await run_task - - _install_run_openai(monkeypatch) - monkeypatch.setattr(agent, "EpisodeClient", RunClient) - monkeypatch.setattr(agent, "_drive_policy", complete_policy) - asyncio.run(scenario()) - assert close_observations == [False] - assert episode_id not in agent._EPISODE_PHASES - - -def test_run_cancellation_drains_normal_close_before_cleanup_retry(monkeypatch): - episode_id = "2" * 24 - first_close_started = None - first_close_live = False - close_observations = [] - - class RunClient: - def __init__(self, **_kwargs): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, _exc_type, _exc, _traceback): - return False - - async def create(self, _task_id, _tier): - return {"episode_id": episode_id, "prompt": "task"} - - async def evaluate(self, value): - assert value == episode_id - return {} - - async def close(self, value): - assert value == episode_id - nonlocal first_close_live - close_observations.append(first_close_live) - if len(close_observations) == 1: - first_close_live = True - first_close_started.set() - try: - await asyncio.Event().wait() - finally: - first_close_live = False - assert not first_close_live - return {} - - async def complete_policy(*_args, **_kwargs): - return "completed" - - async def scenario(): - nonlocal first_close_started - first_close_started = asyncio.Event() - run_task = asyncio.create_task( - agent.run("http://127.0.0.1:8000/v1", None, metadata=_run_metadata()) - ) - await first_close_started.wait() - run_task.cancel() - with pytest.raises(asyncio.CancelledError): - await run_task - - _install_run_openai(monkeypatch) - monkeypatch.setattr(agent, "EpisodeClient", RunClient) - monkeypatch.setattr(agent, "_drive_policy", complete_policy) - asyncio.run(scenario()) - assert close_observations == [False, False] - assert episode_id not in agent._EPISODE_PHASES - - -def test_run_close_failure_retains_cleanup_pending_ownership(monkeypatch): - episode_id = "1" * 24 - - class RunClient: - def __init__(self, **_kwargs): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, _exc_type, _exc, _traceback): - return False - - async def create(self, _task_id, _tier): - return {"episode_id": episode_id, "prompt": "task"} - - async def evaluate(self, value): - assert value == episode_id - return { - "task_id": "CVE-2024-1234", - "episode_id": episode_id, - "reward": 0.0, - "passed": False, - } - - async def close(self, value): - assert value == episode_id - raise EpisodeTransportError("close failed") - - async def complete_policy(*_args, **_kwargs): - return "completed" - - _install_run_openai(monkeypatch) - monkeypatch.setattr(agent, "EpisodeClient", RunClient) - monkeypatch.setattr(agent, "_drive_policy", complete_policy) - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "k" * 48) - try: - result = asyncio.run( - agent.run("http://127.0.0.1:8000/v1", None, metadata=_run_metadata()) - ) - assert result is not None - assert result["exit_status"] == reward.CLEANUP_ERROR_STATUS - assert ( - result[reward.OUTCOME_KEY]["status"] - == reward.CLEANUP_ERROR_STATUS - ) - assert agent._EPISODE_PHASES[episode_id] == "cleanup_pending" - assert agent._claim_episodes_and_tasks_for_abort() == [ - (episode_id, agent._EPISODE_POLICY_TASKS[episode_id]) - ] - finally: - agent._release_episode(episode_id) - - -def test_run_recovers_close_before_returning_retryable_infrastructure(monkeypatch): - episode_id = "2" * 24 - close_calls = [] - - class RunClient: - def __init__(self, **_kwargs): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, _exc_type, _exc, _traceback): - return False - - async def create(self, _task_id, _tier): - return {"episode_id": episode_id, "prompt": "task"} - - async def evaluate(self, value): - assert value == episode_id - return { - "task_id": "CVE-2024-1234", - "episode_id": episode_id, - "reward": 0.0, - "passed": False, - } - - async def close(self, value): - assert value == episode_id - close_calls.append(value) - if len(close_calls) == 1: - raise EpisodeTransportError("close failed once") - return {"closed": True} - - async def complete_policy(*_args, **_kwargs): - return "completed" - - _install_run_openai(monkeypatch) - monkeypatch.setattr(agent, "EpisodeClient", RunClient) - monkeypatch.setattr(agent, "_drive_policy", complete_policy) - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "k" * 48) - - result = asyncio.run( - agent.run("http://127.0.0.1:8000/v1", None, metadata=_run_metadata()) - ) - - assert result is not None - assert result["exit_status"] == reward.INFRASTRUCTURE_STATUS - assert len(close_calls) == 2 - agent.require_no_episode_residue() - - -def test_second_run_cancellation_drains_cleanup_close_before_handoff(monkeypatch): - episode_id = "0" * 24 - policy_started = None - cleanup_close_started = None - close_live = False - close_observations = [] - - class RunClient: - def __init__(self, **_kwargs): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, _exc_type, _exc, _traceback): - return False - - async def create(self, _task_id, _tier): - return {"episode_id": episode_id, "prompt": "task"} - - async def close(self, value): - assert value == episode_id - nonlocal close_live - close_observations.append(close_live) - if len(close_observations) == 1: - close_live = True - cleanup_close_started.set() - try: - await asyncio.Event().wait() - finally: - close_live = False - assert not close_live - return {} - - async def blocked_policy(*_args, **_kwargs): - policy_started.set() - await asyncio.Event().wait() - - async def scenario(): - nonlocal policy_started, cleanup_close_started - policy_started = asyncio.Event() - cleanup_close_started = asyncio.Event() - run_task = asyncio.create_task( - agent.run("http://127.0.0.1:8000/v1", None, metadata=_run_metadata()) - ) - await policy_started.wait() - run_task.cancel() - await cleanup_close_started.wait() - run_task.cancel() - with pytest.raises(asyncio.CancelledError): - await run_task - assert not close_live - assert agent._EPISODE_PHASES[episode_id] == "cleanup_pending" - await agent.abort() - - _install_run_openai(monkeypatch) - monkeypatch.setattr(agent, "EpisodeClient", RunClient) - monkeypatch.setattr(agent, "_drive_policy", blocked_policy) - asyncio.run(scenario()) - assert close_observations == [False, False] - assert episode_id not in agent._EPISODE_PHASES - - -def test_abort_drains_pending_policy_before_close(monkeypatch): - episode_id = "d" * 24 - events = [] - - class AbortClient: - def __init__(self, *, total_timeout_seconds): - assert total_timeout_seconds == 180.0 - - async def __aenter__(self): - return self - - async def __aexit__(self, _exc_type, _exc, _traceback): - return False - - async def close(self, value): - assert value == episode_id - assert policy_task.done() - events.append("close") - return {} - - async def scenario(): - started = asyncio.Event() - - async def pending_policy(): - events.append("policy_started") - started.set() - try: - await asyncio.Event().wait() - finally: - events.append("policy_drained") - - nonlocal policy_task - policy_task = asyncio.create_task(pending_policy()) - agent._register_episode(episode_id, policy_task) - await started.wait() - await agent.abort() - assert policy_task.cancelled() - - policy_task = None - monkeypatch.setattr(agent, "EpisodeClient", AbortClient) - try: - asyncio.run(scenario()) - assert events == ["policy_started", "policy_drained", "close"] - assert episode_id not in agent._EPISODE_PHASES - finally: - agent._release_episode(episode_id) - - -def test_abort_close_retries_episode_conflict_without_releasing_ownership( - monkeypatch, -): - episode_id = "9" * 24 - phases = [] - - class ConflictClient: - async def close(self, value): - assert value == episode_id - phases.append(agent._EPISODE_PHASES.get(episode_id)) - if len(phases) < 3: - raise EpisodeAPIError(409, "episode_conflict", "still active") - return {} - - async def no_wait(_delay): - return None - - monkeypatch.setattr(agent.asyncio, "sleep", no_wait) - agent._register_episode(episode_id) - try: - assert agent._claim_driving_episodes_for_abort() == [episode_id] - assert asyncio.run( - agent._close_aborted_episode( - ConflictClient(), - episode_id, - retry_timeout_seconds=1.0, - ) - ) - assert phases == ["aborting", "aborting", "aborting"] - assert episode_id not in agent._EPISODE_PHASES - finally: - agent._release_episode(episode_id) - - -@pytest.mark.parametrize( - ("status", "code"), - [(409, "episode_conflict"), (500, "internal_error")], -) -def test_failed_abort_close_retains_retryable_cleanup_ownership(status, code): - episode_id = "8" * 24 - - class FailingClient: - async def close(self, value): - assert value == episode_id - raise EpisodeAPIError(status, code, "not closed") - - agent._register_episode(episode_id) - try: - assert agent._claim_driving_episodes_for_abort() == [episode_id] - assert not asyncio.run( - agent._close_aborted_episode( - FailingClient(), - episode_id, - retry_timeout_seconds=0.0, - ) - ) - assert agent._EPISODE_PHASES[episode_id] == "aborting" - assert agent._claim_episodes_and_tasks_for_abort() == [] - agent._mark_episode_cleanup_pending(episode_id) - assert agent._EPISODE_PHASES[episode_id] == "cleanup_pending" - assert agent._claim_episodes_and_tasks_for_abort() == [(episode_id, None)] - assert agent._EPISODE_PHASES[episode_id] == "aborting" - finally: - agent._release_episode(episode_id) - - -def test_failed_close_cannot_be_reclaimed_before_owner_finishes(monkeypatch): - episode_id = "4" * 24 - - class FailingClient: - async def close(self, value): - assert value == episode_id - raise EpisodeAPIError(500, "internal_error", "not closed") - - class UnexpectedAbortClient: - def __init__(self, **_kwargs): - raise AssertionError("overlapping abort reclaimed live cleanup ownership") - - async def scenario(): - assert not await agent._close_aborted_episode( - FailingClient(), - episode_id, - retry_timeout_seconds=0.0, - ) - assert agent._EPISODE_PHASES[episode_id] == "aborting" - await agent.abort() - assert agent._EPISODE_PHASES[episode_id] == "aborting" - agent._mark_episode_cleanup_pending(episode_id) - - monkeypatch.setattr(agent, "EpisodeClient", UnexpectedAbortClient) - agent._register_episode(episode_id) - try: - assert agent._claim_driving_episodes_for_abort() == [episode_id] - asyncio.run(scenario()) - assert agent._EPISODE_PHASES[episode_id] == "cleanup_pending" - finally: - agent._release_episode(episode_id) - - -def test_timed_out_policy_drain_retains_live_owner_until_abort_returns(): - episode_id = "a" * 23 + "0" - - async def scenario(): - release = asyncio.Event() - - async def slow_cancellation(): - try: - await asyncio.Event().wait() - except asyncio.CancelledError: - await release.wait() - - task = asyncio.create_task(slow_cancellation()) - agent._register_episode(episode_id, task) - await asyncio.sleep(0) - assert agent._claim_episodes_and_tasks_for_abort() == [(episode_id, task)] - assert not await agent._drain_aborted_policy_task(task, timeout_seconds=0.0) - assert agent._EPISODE_PHASES[episode_id] == "aborting" - assert agent._claim_episodes_and_tasks_for_abort() == [] - release.set() - await task - - try: - asyncio.run(scenario()) - finally: - agent._release_episode(episode_id) - - -def test_abort_without_claimable_episodes_does_not_initialize_client(monkeypatch): - class UnexpectedClient: - def __init__(self, **_kwargs): - raise AssertionError("abort client initialized without cleanup work") - - monkeypatch.setattr(agent, "EpisodeClient", UnexpectedClient) - asyncio.run(agent.abort()) - - -def test_cancelled_abort_retains_cleanup_pending_ownership(monkeypatch): - episode_id = "7" * 24 - - class BlockingClient: - def __init__(self, *, total_timeout_seconds): - assert total_timeout_seconds == 180.0 - - async def __aenter__(self): - return self - - async def __aexit__(self, _exc_type, _exc, _traceback): - return False - - async def close(self, value): - assert value == episode_id - close_started.set() - await asyncio.Event().wait() - - async def scenario(): - nonlocal close_started - close_started = asyncio.Event() - abort_task = asyncio.create_task(agent.abort()) - await close_started.wait() - abort_task.cancel() - with pytest.raises(asyncio.CancelledError): - await abort_task - - close_started = None - monkeypatch.setattr(agent, "EpisodeClient", BlockingClient) - agent._register_episode(episode_id) - try: - asyncio.run(scenario()) - assert agent._EPISODE_PHASES[episode_id] == "cleanup_pending" - assert agent._claim_episodes_and_tasks_for_abort() == [(episode_id, None)] - finally: - agent._release_episode(episode_id) - - -def test_abort_client_initialization_failure_retains_cleanup_ownership(monkeypatch): - episode_id = "c" * 24 - - class FailingEpisodeClient: - def __init__(self, *, total_timeout_seconds): - assert total_timeout_seconds == 180.0 - - async def __aenter__(self): - raise RuntimeError("client unavailable") - - async def __aexit__(self, _exc_type, _exc, _traceback): - return False - - monkeypatch.setattr(agent, "EpisodeClient", FailingEpisodeClient) - agent._register_episode(episode_id) - try: - asyncio.run(agent.abort()) - assert not agent._claim_episode_finalization(episode_id) - assert agent._EPISODE_PHASES[episode_id] == "cleanup_pending" - assert agent._claim_episodes_and_tasks_for_abort() == [(episode_id, None)] - finally: - agent._release_episode(episode_id) - - -def _tool_call(identifier, name, arguments): - return { - "id": identifier, - "type": "function", - "function": {"name": name, "arguments": json.dumps(arguments)}, - } - - -def test_adapter_is_network_only_and_enforces_l2(): - prompt_surface = agent.SYSTEM_PROMPT + json.dumps(agent.TOOLS) - assert "DOCKER_HOST" not in prompt_surface - assert "DEBUG_URL" not in prompt_surface - assert "debug daemon" not in prompt_surface - assert agent._task_identity({"task_id": "CVE-2024-1234"}) == ( - "CVE-2024-1234", - "l2", - ) - assert agent._task_identity( - {"task_id": "CVE-2024-1234", "prompt_tier": "l2"} - ) == ("CVE-2024-1234", "l2") - for tier in ("l0", "l1"): - with pytest.raises(ValueError, match="must use prompt_tier l2"): - agent._task_identity( - {"task_id": "CVE-2024-1234", "prompt_tier": tier} - ) - - -def test_agent_preserves_full_assistant_message_and_uses_tool_roles(monkeypatch): - monkeypatch.setenv("SECRLENV_MAX_TURNS", "4") - policy = FakePolicy( - [ - { - "role": "assistant", - "content": "", - "reasoning_content": "inspect first", - "tool_calls": [ - _tool_call("call-1", "terminal.exec", {"command": "id"}) - ], - }, - { - "role": "assistant", - "content": "", - "reasoning_content": "submit evidence", - "tool_calls": [ - _tool_call( - "call-2", "submit", {"evidence": "target returned nonce"} - ) - ], - }, - ] - ) - client = FakeEpisodeClient() - metrics = agent.AgentMetrics() - status = asyncio.run( - agent._drive_policy( - policy, - client, - {"episode_id": "a" * 24, "prompt": "solve this"}, - {"temperature": 0.7, "stream": True}, - metrics, - ) - ) - assert status == "completed" - assert client.exec_calls[0][1] == "id" - assert client.submissions == [ - ("a" * 24, {"evidence": "target returned nonce"}) - ] - second_messages = policy.chat.completions.requests[1]["messages"] - first_assistant = second_messages[2] - assert first_assistant["reasoning_content"] == "inspect first" - assert first_assistant["tool_calls"][0]["id"] == "call-1" - assert second_messages[3] == { - "role": "tool", - "tool_call_id": "call-1", - "content": json.dumps( - { - "exit_code": 0, - "output": "command output", - "timed_out": False, - "truncated": False, - }, - separators=(",", ":"), - ), - } - assert "stream" not in policy.chat.completions.requests[0]["extra_body"] - assert metrics.terminal_calls == 1 - assert metrics.submit_calls == 1 - - -def test_agent_retries_non_tool_prose_as_user_turn(monkeypatch): - monkeypatch.setenv("SECRLENV_MAX_TURNS", "3") - monkeypatch.setenv("SECRLENV_MAX_NO_TOOL_RETRIES", "2") - policy = FakePolicy( - [ - {"role": "assistant", "content": "I would run id"}, - {"role": "assistant", "content": "Still no tool"}, - ] - ) - metrics = agent.AgentMetrics() - status = asyncio.run( - agent._drive_policy( - policy, - FakeEpisodeClient(), - {"episode_id": "b" * 24, "prompt": "solve"}, - {}, - metrics, - ) - ) - assert status == "max_turns" - assert policy.chat.completions.requests[1]["messages"][-1]["role"] == "user" - assert metrics.parse_failures == 2 - - -def test_signed_zero_reward_is_valid_and_tampering_is_rejected(monkeypatch): - key = "k" * 48 - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", key) - outcome = { - "schema": 1, - "status": "completed", - "episode_id": "c" * 24, - "task_id": "CVE-2024-1234", - "reward": 0.0, - "passed": False, - "class": "rce", - reward.INFRASTRUCTURE_503_RETRIES_KEY: 0, - } - metadata = { - reward.OUTCOME_KEY: outcome, - reward.MAC_KEY: reward.sign_outcome(outcome), - } - assert reward.verify_outcome(metadata) == 0.0 - sample = SimpleNamespace(metadata=metadata) - assert asyncio.run(reward.reward_func(None, sample)) == 0.0 - - forged = copy.deepcopy(metadata) - forged[reward.OUTCOME_KEY]["reward"] = 1.0 - with pytest.raises(reward.UntrustedOutcome, match="signature mismatch"): - reward.verify_outcome(forged) - - -def test_sequence_cap_outcome_contract_and_unknown_status(monkeypatch): - key = "s" * 48 - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", key) - task_id = "CVE-2024-1234" - episode_id = "m" * 24 - outcome = agent._validated_outcome( - { - "task_id": task_id, - "episode_id": episode_id, - "reward": 0.25, - "passed": False, - "grader": {"class": "rce"}, - }, - task_id, - episode_id, - "max_seq_len", - ) - metadata = { - reward.OUTCOME_KEY: outcome, - reward.MAC_KEY: reward.sign_outcome(outcome), - } - - assert reward.verify_outcome(metadata) == 0.25 - - unknown = copy.deepcopy(outcome) - unknown["status"] = "unknown_termination" - unknown_metadata = { - reward.OUTCOME_KEY: unknown, - reward.MAC_KEY: reward.sign_outcome(unknown), - } - with pytest.raises(reward.UntrustedOutcome, match="invalid schema or status"): - reward.verify_outcome(unknown_metadata) - - -def test_missing_reward_never_defaults_to_zero(monkeypatch): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "z" * 48) - sample = SimpleNamespace(metadata={"reward": 0.0}) - with pytest.raises(reward.UntrustedOutcome, match="no signed"): - asyncio.run(reward.reward_func(None, sample)) - - -def _install_filter_types(monkeypatch): - class DynamicFilterOutput: - def __init__(self, *, keep, reason): - self.keep = keep - self.reason = reason - - class Sample: - class Status: - ABORTED = "aborted" - PENDING = "pending" - COMPLETED = "completed" - TRUNCATED = "truncated" - FAILED = "failed" - - miles = types.ModuleType("miles") - rollout = types.ModuleType("miles.rollout") - filter_hub = types.ModuleType("miles.rollout.filter_hub") - base_types = types.ModuleType("miles.rollout.filter_hub.base_types") - utils = types.ModuleType("miles.utils") - sample_types = types.ModuleType("miles.utils.types") - base_types.DynamicFilterOutput = DynamicFilterOutput - sample_types.Sample = Sample - monkeypatch.setitem(sys.modules, "miles", miles) - monkeypatch.setitem(sys.modules, "miles.rollout", rollout) - monkeypatch.setitem(sys.modules, "miles.rollout.filter_hub", filter_hub) - monkeypatch.setitem( - sys.modules, "miles.rollout.filter_hub.base_types", base_types - ) - monkeypatch.setitem(sys.modules, "miles.utils", utils) - monkeypatch.setitem(sys.modules, "miles.utils.types", sample_types) - - -def _signed_sample( - index, value, *, status="completed", outcome_status="completed" -): - outcome = { - "schema": 1, - "status": outcome_status, - "episode_id": f"episode-{index:024d}", - "task_id": "CVE-2019-7859", - "reward": value, - reward.INFRASTRUCTURE_503_RETRIES_KEY: 0, - } - return SimpleNamespace( - index=index, - status=status, - metadata={ - reward.OUTCOME_KEY: outcome, - reward.MAC_KEY: reward.sign_outcome(outcome), - }, - ) - - -def test_signed_infrastructure_outcome_is_only_an_abort_signal(monkeypatch): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "i" * 48) - _install_filter_types(monkeypatch) - outcome = agent._infrastructure_outcome( - "CVE-2024-1234", "f" * 24 - ) - metadata = { - reward.OUTCOME_KEY: outcome, - reward.MAC_KEY: reward.sign_outcome(outcome), - } - sample = SimpleNamespace(index=30, status="completed", metadata=metadata) - - assert asyncio.run(reward.reward_func(None, sample)) == 0.0 - assert sample.status == "aborted" - with pytest.raises( - reward.UntrustedOutcome, match="not a grader verdict" - ): - reward.verify_outcome(metadata) - - result = reward.check_group( - SimpleNamespace( - yeto_rl_policy_version=10, - yeto_rl_dynamic_sampling_max_replacements=0, - _yeto_secrlenv_retry_callback=lambda *_args, **_kwargs: None, - ), - [sample], - ) - assert (result.keep, result.reason) == ( - False, - "secrlenv_infrastructure_failure", - ) - - direct_sample = SimpleNamespace( - index=31, status="completed", metadata=copy.deepcopy(metadata) - ) - direct = reward.check_group( - SimpleNamespace( - yeto_rl_policy_version=10, - yeto_rl_dynamic_sampling_max_replacements=0, - _yeto_secrlenv_retry_callback=lambda *_args, **_kwargs: None, - ), - [direct_sample], - ) - assert (direct.keep, direct.reason) == ( - False, - "secrlenv_infrastructure_failure", - ) - - forged = copy.deepcopy(metadata) - forged[reward.OUTCOME_KEY]["task_id"] = "CVE-2024-9999" - forged_sample = SimpleNamespace( - index=32, status="completed", metadata=forged - ) - with pytest.raises(reward.UntrustedOutcome, match="signature mismatch"): - asyncio.run(reward.reward_func(None, forged_sample)) - assert forged_sample.status == "completed" - - -def test_admission_outcome_schema_is_strict_and_nonce_distinguishes_attempts( - monkeypatch, -): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "n" * 48) - first = agent._admission_infrastructure_outcome( - "CVE-2024-1234", 1, nonce="1" * 32 - ) - second = agent._admission_infrastructure_outcome( - "CVE-2024-1234", 1, nonce="2" * 32 - ) - first_metadata = { - reward.OUTCOME_KEY: first, - reward.MAC_KEY: reward.sign_outcome(first), - } - second_metadata = { - reward.OUTCOME_KEY: second, - reward.MAC_KEY: reward.sign_outcome(second), - } - - assert reward._verified_outcome(first_metadata)[1] == 0.0 - assert reward._group_key( - [SimpleNamespace(index=9, metadata=first_metadata)] - ) != reward._group_key([SimpleNamespace(index=9, metadata=second_metadata)]) - - for field, invalid in ( - (reward.ADMISSION_PHASE_KEY, "post_create"), - (reward.ADMISSION_ERROR_CODE_KEY, "capacity_reached"), - (reward.ADMISSION_NONCE_KEY, "bad"), - (reward.INFRASTRUCTURE_503_RETRIES_KEY, 0), - ("episode_id", "not-null"), - ): - forged = copy.deepcopy(first) - forged[field] = invalid - metadata = { - reward.OUTCOME_KEY: forged, - reward.MAC_KEY: reward.sign_outcome(forged), - } - with pytest.raises(reward.UntrustedOutcome, match="admission outcome"): - reward._verified_outcome(metadata) - - extra = copy.deepcopy(first) - extra["unexpected"] = True - with pytest.raises(reward.UntrustedOutcome, match="admission outcome"): - reward._verified_outcome( - { - reward.OUTCOME_KEY: extra, - reward.MAC_KEY: reward.sign_outcome(extra), - } - ) - - -def _install_generate_upstream(monkeypatch, upstream): - package = types.ModuleType("miles") - rollout = types.ModuleType("miles.rollout") - hub = types.ModuleType("miles.rollout.generate_hub") - upstream_module = types.ModuleType( - "miles.rollout.generate_hub.agentic_tool_call" - ) - upstream_module.generate = upstream - monkeypatch.setitem(sys.modules, "miles", package) - monkeypatch.setitem(sys.modules, "miles.rollout", rollout) - monkeypatch.setitem(sys.modules, "miles.rollout.generate_hub", hub) - monkeypatch.setitem( - sys.modules, "miles.rollout.generate_hub.agentic_tool_call", upstream_module - ) - - -def test_generate_wrapper_preserves_only_verified_metadata_on_aborted_fallback( - monkeypatch, -): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "g" * 48) - sample = SimpleNamespace(status=SimpleNamespace(value="aborted"), metadata={}) - outcome = agent._admission_infrastructure_outcome( - "CVE-2024-1234", 1, nonce="3" * 32 - ) - signed = { - reward.OUTCOME_KEY: outcome, - reward.MAC_KEY: reward.sign_outcome(outcome), - "exit_status": reward.INFRASTRUCTURE_STATUS, - "agent_metrics": {}, - } - - async def upstream(_input): - secrlenv_generate.capture_agent_metadata(signed) - return SimpleNamespace(samples=sample) - - _install_generate_upstream(monkeypatch, upstream) - output = asyncio.run(secrlenv_generate.generate(SimpleNamespace())) - - assert output.samples is sample - assert sample.status.value == "aborted" - assert reward._verified_outcome(sample.metadata)[0] == outcome - assert set(sample.metadata) == {reward.OUTCOME_KEY, reward.MAC_KEY} - - -def test_generate_wrapper_rejects_forged_or_missing_fallback_metadata(monkeypatch): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "g" * 48) - captured = {"value": None} - sample_metadata = {"value": {}} - - async def upstream(_input): - if captured["value"] is not None: - secrlenv_generate.capture_agent_metadata(captured["value"]) - return SimpleNamespace( - samples=SimpleNamespace( - status=SimpleNamespace(value="aborted"), - metadata=copy.deepcopy(sample_metadata["value"]), - ) - ) - - _install_generate_upstream(monkeypatch, upstream) - with pytest.raises(reward.UntrustedOutcome): - asyncio.run(secrlenv_generate.generate(SimpleNamespace())) - - outcome = agent._admission_infrastructure_outcome( - "CVE-2024-1234", 1, nonce="4" * 32 - ) - captured["value"] = { - reward.OUTCOME_KEY: outcome, - reward.MAC_KEY: "0" * 64, - } - with pytest.raises(reward.UntrustedOutcome, match="signature mismatch"): - asyncio.run(secrlenv_generate.generate(SimpleNamespace())) - - sample_metadata["value"] = {reward.OUTCOME_KEY: outcome} - with pytest.raises(reward.UntrustedOutcome, match="lost signed outcome metadata"): - asyncio.run(secrlenv_generate.generate(SimpleNamespace())) - - -def test_admission_replacement_nonce_preserves_shared_retry_exhaustion(monkeypatch): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "r" * 48) - _install_filter_types(monkeypatch) - calls = [] - - def retry(group, **kwargs): - assert kwargs["infrastructure_503_retries"] == {201: 1} - if calls: - raise secrlenv_reward.SecRLEnvReplacementExhausted( - "SecRLEnv rollout replacement budget was exhausted" - ) - calls.append((group, kwargs)) - - def admission_sample(nonce): - outcome = agent._admission_infrastructure_outcome( - "CVE-2024-1234", 1, nonce=nonce - ) - return SimpleNamespace( - index=201, - status="aborted", - metadata={ - reward.OUTCOME_KEY: outcome, - reward.MAC_KEY: reward.sign_outcome(outcome), - }, - ) - - args = SimpleNamespace( - yeto_rl_policy_version=92, - yeto_rl_dynamic_sampling_max_replacements=0, - _yeto_secrlenv_retry_callback=retry, - ) - first = [admission_sample("5" * 32)] - exhausted = [admission_sample("6" * 32)] - - initial = reward.check_group(args, first) - with pytest.raises( - secrlenv_reward.SecRLEnvReplacementExhausted, - match="budget was exhausted", - ): - reward.check_group(args, exhausted) - - assert (initial.keep, initial.reason) == ( - False, - "secrlenv_infrastructure_failure", - ) - assert len(calls) == 1 - - -def test_signed_cleanup_failure_is_terminal_without_replacement(monkeypatch): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "c" * 48) - _install_filter_types(monkeypatch) - outcome = agent._cleanup_error_outcome("CVE-2024-1234", "e" * 24) - metadata = { - reward.OUTCOME_KEY: outcome, - reward.MAC_KEY: reward.sign_outcome(outcome), - } - sample = SimpleNamespace(index=33, status="completed", metadata=metadata) - replacements = [] - - assert asyncio.run(reward.reward_func(None, sample)) == 0.0 - with pytest.raises( - secrlenv_reward.SecRLEnvRolloutCleanupError, - match="cleanup could not be verified", - ): - reward.check_group( - SimpleNamespace( - yeto_rl_policy_version=10, - yeto_rl_dynamic_sampling_max_replacements=0, - _yeto_secrlenv_retry_callback=lambda *_args, **_kwargs: ( - replacements.append(True) - ), - ), - [sample], - ) - - assert replacements == [] - - -def test_signed_group_filter_accepts_valid_reward_variance(monkeypatch): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "v" * 48) - _install_filter_types(monkeypatch) - args = SimpleNamespace( - yeto_rl_policy_version=7, - yeto_rl_dynamic_sampling_max_replacements=2, - ) - - result = reward.check_group( - args, [_signed_sample(1, 0.0), _signed_sample(2, 1.0)] - ) - - assert result.keep is True - assert result.reason is None - assert args._yeto_secrlenv_filter_state["rejections"] == 0 - - -def test_signed_group_filter_accepts_completed_and_truncated_statuses(monkeypatch): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "t" * 48) - _install_filter_types(monkeypatch) - args = SimpleNamespace( - yeto_rl_policy_version=71, - yeto_rl_dynamic_sampling_max_replacements=0, - ) - - result = reward.check_group( - args, - [ - _signed_sample(101, 0.0, status="completed"), - _signed_sample(102, 1.0, status="truncated"), - ], - ) - - assert (result.keep, result.reason) == (True, None) - - -def test_signed_zero_variance_groups_have_a_memoized_replacement_bound(monkeypatch): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "b" * 48) - _install_filter_types(monkeypatch) - args = SimpleNamespace( - yeto_rl_policy_version=8, - yeto_rl_dynamic_sampling_max_replacements=2, - ) - first = [_signed_sample(10, 0.0), _signed_sample(11, 0.0)] - second = [_signed_sample(12, 1.0), _signed_sample(13, 1.0)] - fallback = [_signed_sample(14, 0.0), _signed_sample(15, 0.0)] - - decisions = [ - reward.check_group(args, first), - reward.check_group(args, second), - reward.check_group(args, fallback), - reward.check_group(args, fallback), - ] - - assert [decision.keep for decision in decisions] == [False, False, True, True] - assert decisions[2].reason == "secrlenv_bounded_fallback_after_2_replacements" - assert args._yeto_secrlenv_filter_state["rejections"] == 2 - assert args._yeto_secrlenv_filter_state["forced"] == 1 - - -def test_aborted_and_untrusted_groups_never_use_bounded_fallback(monkeypatch): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "u" * 48) - _install_filter_types(monkeypatch) - args = SimpleNamespace( - yeto_rl_policy_version=9, - yeto_rl_dynamic_sampling_max_replacements=0, - _yeto_secrlenv_retry_callback=lambda *_args, **_kwargs: None, - ) - aborted = [_signed_sample(20, 0.0, status="aborted")] - forged = _signed_sample(21, 0.0) - forged.metadata[reward.OUTCOME_KEY]["reward"] = 1.0 - - aborted_result = reward.check_group(args, aborted) - with pytest.raises( - secrlenv_reward.SecRLEnvUntrustedEvidence, - match="authentication failed", - ): - reward.check_group(args, [forged]) - - assert (aborted_result.keep, aborted_result.reason) == ( - False, - "secrlenv_aborted", - ) - assert aborted[0].status == "aborted" - assert args._yeto_secrlenv_filter_state["nontrainable_rejections"] == 1 - - -def test_signed_infrastructure_replacement_is_memoized_and_exhausts(monkeypatch): - monkeypatch.setenv("SECRLENV_REWARD_HMAC_KEY", "r" * 48) - _install_filter_types(monkeypatch) - calls = [] - - def retry(group, **kwargs): - if calls: - raise secrlenv_reward.SecRLEnvReplacementExhausted( - "SecRLEnv rollout replacement budget was exhausted" - ) - calls.append((group, kwargs)) - - args = SimpleNamespace( - yeto_rl_policy_version=91, - yeto_rl_dynamic_sampling_max_replacements=0, - _yeto_secrlenv_retry_callback=retry, - ) - first = [ - _signed_sample( - 201, - 0.0, - status="aborted", - outcome_status=reward.INFRASTRUCTURE_STATUS, - ) - ] - second = [ - _signed_sample( - 202, - 0.0, - status="aborted", - outcome_status=reward.INFRASTRUCTURE_STATUS, - ) - ] - - initial = reward.check_group(args, first) - duplicate = reward.check_group(args, first) - with pytest.raises( - secrlenv_reward.SecRLEnvReplacementExhausted, - match="budget was exhausted", - ): - reward.check_group(args, second) - - assert (initial.keep, duplicate.keep) == (False, False) - assert len(calls) == 1 - assert first[0].status == "aborted" - - -def test_evaluation_identity_and_range_are_checked(): - with pytest.raises(EpisodeClientError, match="identity"): - agent._validated_outcome( - { - "task_id": "CVE-2024-9999", - "episode_id": "d" * 24, - "reward": 0.5, - "passed": False, - }, - "CVE-2024-1234", - "d" * 24, - "completed", - ) - with pytest.raises(EpisodeClientError, match="invalid verdict"): - agent._validated_outcome( - { - "task_id": "CVE-2024-1234", - "episode_id": "d" * 24, - "reward": 2.0, - "passed": False, - }, - "CVE-2024-1234", - "d" * 24, - "completed", - ) - - -def test_client_refuses_non_loopback_or_credentialed_origins(): - token = "t" * 48 - for value in ( - "http://example.com:8765", - "https://127.0.0.1:8765", - "http://user:pass@127.0.0.1:8765", - "http://127.0.0.1:8765/path", - ): - with pytest.raises(EpisodeClientError): - EpisodeClient(value, token) - - -def test_daemon_readiness_fails_fast_without_a_listener(monkeypatch): - monkeypatch.setenv("SECRLENV_DAEMON_URL", "http://127.0.0.1:8765") - - def unavailable(*_args, **_kwargs): - raise OSError - - monkeypatch.setattr("yeto_miles_secrlenv.client.urlopen", unavailable) - with pytest.raises(EpisodeTransportError, match="not healthy"): - require_daemon_ready(timeout_seconds=0.1) - - -def test_daemon_readiness_attests_health_and_task_pack(monkeypatch): - expected = "a" * 64 - monkeypatch.setenv("SECRLENV_DAEMON_URL", "http://127.0.0.1:8765") - monkeypatch.setenv("SECRLENV_TASK_PACK_SHA256", expected) - - class Response: - status = 200 - - def __enter__(self): - return self - - def __exit__(self, *_args): - return None - - def read(self): - return BytesIO( - json.dumps({"ok": True, "task_pack_sha256": expected}).encode() - ).read() - - monkeypatch.setattr( - "yeto_miles_secrlenv.client.urlopen", lambda *_a, **_k: Response() - ) - require_daemon_ready(timeout_seconds=0.1) - - monkeypatch.setenv("SECRLENV_TASK_PACK_SHA256", "b" * 64) - with pytest.raises(EpisodeTransportError, match="task-pack identity"): - require_daemon_ready(timeout_seconds=0.1) - - -def test_create_retry_default_covers_sixty_rollouts_at_capacity_sixteen( - monkeypatch, -): - monkeypatch.delenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", raising=False) - observed = [] - - def positive_env(name, default): - observed.append((name, default)) - return 0.001 - - class UnavailableClient: - calls = 0 - - async def create(self, *_args): - self.calls += 1 - raise EpisodeTransportError("offline") - - client = UnavailableClient() - monkeypatch.setattr(agent, "_positive_env", positive_env) - with pytest.raises(EpisodeTransportError, match="offline"): - asyncio.run( - agent._create_with_capacity_retry( - client, "CVE-2024-1234", "l2" - ) - ) - assert client.calls == 1 - assert observed == [("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", 14400.0)] - rollout_count = 60 - max_active = 16 - rollout_timeout_seconds = 3600.0 - admission_waves = (rollout_count + max_active - 1) // max_active - assert observed[0][1] >= admission_waves * rollout_timeout_seconds - - -def test_create_retries_capacity_then_one_infrastructure_error_same_identity( - monkeypatch, -): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "30") - expected = {"episode_id": "a" * 24} - responses = [ - EpisodeAPIError(503, "capacity_reached", "busy"), - EpisodeAPIError(503, "infrastructure_error", "provisioning failed"), - expected, - ] - calls = [] - delays = [] - - class RecoveringClient: - async def create(self, task_id, tier): - calls.append((task_id, tier)) - response = responses.pop(0) - if isinstance(response, Exception): - raise response - return response - - async def no_wait(delay): - delays.append(delay) - - monkeypatch.setattr(agent.asyncio, "sleep", no_wait) - monkeypatch.setattr(agent.random, "uniform", lambda *_args: 0.5) - result = asyncio.run( - agent._create_with_capacity_retry( - RecoveringClient(), "CVE-2024-1234", "l2" - ) - ) - assert result == expected - assert calls == [("CVE-2024-1234", "l2")] * 3 - assert delays == [0.5, 0.5] - - -def test_capacity_retry_stops_at_the_configured_deadline(monkeypatch): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "0.02") - calls = 0 - - class BusyClient: - async def create(self, *_args): - nonlocal calls - calls += 1 - raise EpisodeAPIError(503, "capacity_reached", "busy") - - monkeypatch.setattr(agent.random, "uniform", lambda *_args: 0.005) - with pytest.raises(EpisodeAPIError) as caught: - asyncio.run( - agent._create_with_capacity_retry( - BusyClient(), "CVE-2024-1234", "l2" - ) - ) - assert caught.value.code == "capacity_reached" - assert 2 <= calls < 20 - - -def test_capacity_retry_accepts_success_exactly_at_deadline(monkeypatch): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "1") - clock = [0.0] - failure = EpisodeAPIError(503, "capacity_reached", "busy") - calls = [] - - class DeadlineClient: - async def create(self, *_args): - calls.append(clock[0]) - if len(calls) == 1: - raise failure - clock[0] = 1.0 - return {"episode_id": "a" * 24} - - async def advance(delay): - clock[0] += delay - - monkeypatch.setattr( - agent, "time", SimpleNamespace(monotonic=lambda: clock[0]) - ) - monkeypatch.setattr(agent.asyncio, "sleep", advance) - monkeypatch.setattr(agent.random, "uniform", lambda *_args: 0.1) - result = asyncio.run( - agent._create_with_capacity_retry( - DeadlineClient(), "CVE-2024-1234", "l2" - ) - ) - assert result == {"episode_id": "a" * 24} - assert calls == [0.0, 0.1] - - -def test_infrastructure_retry_accepts_late_create_instead_of_stale_503( - monkeypatch, -): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "1") - clock = [0.0] - calls = [] - ledger = {reward.INFRASTRUCTURE_503_RETRIES_KEY: 0} - - class DeadlineClient: - async def create(self, *_args): - calls.append(clock[0]) - if len(calls) == 1: - raise EpisodeAPIError( - 503, "infrastructure_error", "unavailable" - ) - clock[0] = 1.0 - return {"episode_id": "b" * 24} - - async def advance(delay): - clock[0] += delay - - monkeypatch.setattr( - agent, "time", SimpleNamespace(monotonic=lambda: clock[0]) - ) - monkeypatch.setattr(agent.asyncio, "sleep", advance) - monkeypatch.setattr(agent.random, "uniform", lambda *_args: 0.1) - - result = asyncio.run( - agent._create_with_capacity_retry( - DeadlineClient(), - "CVE-2024-1234", - "l2", - retry_ledger=ledger, - ) - ) - - assert result == {"episode_id": "b" * 24} - assert calls == [0.0, 0.1] - assert ledger == {reward.INFRASTRUCTURE_503_RETRIES_KEY: 1} - - -def test_capacity_retry_does_not_create_at_the_deadline(monkeypatch): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "1") - clock = [0.0] - failure = EpisodeAPIError(503, "capacity_reached", "busy") - calls = 0 - - class DeadlineClient: - async def create(self, *_args): - nonlocal calls - calls += 1 - if calls > 1: - raise AssertionError("a retry began at or after its deadline") - raise failure - - async def advance(delay): - clock[0] += delay - - monkeypatch.setattr( - agent, "time", SimpleNamespace(monotonic=lambda: clock[0]) - ) - monkeypatch.setattr(agent.asyncio, "sleep", advance) - monkeypatch.setattr(agent.random, "uniform", lambda *_args: 3.0) - with pytest.raises(EpisodeAPIError) as caught: - asyncio.run( - agent._create_with_capacity_retry( - DeadlineClient(), "CVE-2024-1234", "l2" - ) - ) - assert caught.value is failure - assert calls == 1 - - -def test_capacity_retry_bounds_a_blocked_retry_by_remaining_time(monkeypatch): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "0.05") - failure = EpisodeAPIError(503, "capacity_reached", "busy") - calls = 0 - retry_cancelled = False - - class BlockingClient: - async def create(self, *_args): - nonlocal calls, retry_cancelled - calls += 1 - if calls == 1: - raise failure - try: - await asyncio.Event().wait() - except asyncio.CancelledError: - retry_cancelled = True - raise - - async def no_wait(_delay): - return None - - monkeypatch.setattr(agent.asyncio, "sleep", no_wait) - with pytest.raises(EpisodeAPIError) as caught: - asyncio.run( - agent._create_with_capacity_retry( - BlockingClient(), "CVE-2024-1234", "l2" - ) - ) - assert caught.value is failure - assert calls == 2 - assert retry_cancelled is True - - -def test_capacity_retry_preserves_external_cancellation(monkeypatch): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "30") - retry_started = asyncio.Event() - retry_cancelled = False - calls = 0 - - class BlockingClient: - async def create(self, *_args): - nonlocal calls, retry_cancelled - calls += 1 - if calls == 1: - raise EpisodeAPIError(503, "capacity_reached", "busy") - retry_started.set() - try: - await asyncio.Event().wait() - except asyncio.CancelledError: - retry_cancelled = True - raise - - async def no_wait(_delay): - return None - - async def drive_and_cancel(): - task = asyncio.create_task( - agent._create_with_capacity_retry( - BlockingClient(), "CVE-2024-1234", "l2" - ) - ) - await retry_started.wait() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - monkeypatch.setattr(agent.asyncio, "sleep", no_wait) - asyncio.run(drive_and_cancel()) - assert calls == 2 - assert retry_cancelled is True - - -def test_infrastructure_error_is_retried_at_most_once(monkeypatch): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "30") - calls = 0 - delays = [] - - class BrokenClient: - async def create(self, *_args): - nonlocal calls - calls += 1 - raise EpisodeAPIError( - 503, "infrastructure_error", "provisioning failed" - ) - - async def no_wait(delay): - delays.append(delay) - - monkeypatch.setattr(agent.asyncio, "sleep", no_wait) - monkeypatch.setattr(agent.random, "uniform", lambda *_args: 0.5) - ledger = {reward.INFRASTRUCTURE_503_RETRIES_KEY: 0} - with pytest.raises(EpisodeAPIError) as caught: - asyncio.run( - agent._create_with_capacity_retry( - BrokenClient(), - "CVE-2024-1234", - "l2", - retry_ledger=ledger, - ) - ) - assert caught.value.code == "infrastructure_error" - assert calls == 2 - assert delays == [0.5] - assert ledger == {reward.INFRASTRUCTURE_503_RETRIES_KEY: 1} - - -def test_same_task_retry_shares_infrastructure_503_attempt_ledger(monkeypatch): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "30") - calls = 0 - - class BrokenClient: - async def create(self, *_args): - nonlocal calls - calls += 1 - raise EpisodeAPIError( - 503, "infrastructure_error", "provisioning failed" - ) - - async def unexpected_wait(_delay): - raise AssertionError("a consumed infrastructure retry must not back off") - - monkeypatch.setattr(agent.asyncio, "sleep", unexpected_wait) - ledger = {reward.INFRASTRUCTURE_503_RETRIES_KEY: 1} - - with pytest.raises(EpisodeAPIError) as caught: - asyncio.run( - agent._create_with_capacity_retry( - BrokenClient(), - "CVE-2024-1234", - "l2", - retry_ledger=ledger, - ) - ) - - assert caught.value.code == "infrastructure_error" - assert calls == 1 - assert ledger == {reward.INFRASTRUCTURE_503_RETRIES_KEY: 1} - - -@pytest.mark.parametrize( - "failure", - [ - EpisodeAPIError(400, "invalid_request", "bad task"), - EpisodeAPIError(400, "capacity_reached", "wrong status"), - EpisodeAPIError(500, "internal_error", "unexpected failure"), - EpisodeTransportError("ambiguous transport failure"), - ], - ids=("invalid-request", "wrong-capacity-status", "internal-error", "transport"), -) -def test_create_does_not_retry_nonretryable_or_transport_errors( - monkeypatch, failure -): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "30") - calls = 0 - - class FailingClient: - async def create(self, *_args): - nonlocal calls - calls += 1 - raise failure - - async def unexpected_wait(_delay): - raise AssertionError("single-shot failures must not back off") - - monkeypatch.setattr(agent.asyncio, "sleep", unexpected_wait) - with pytest.raises(type(failure)) as caught: - asyncio.run( - agent._create_with_capacity_retry( - FailingClient(), "CVE-2024-1234", "l2" - ) - ) - assert caught.value is failure - assert calls == 1 - - -@pytest.mark.parametrize("value", ["0", "-1", "nan", "inf", "invalid"]) -def test_create_retry_rejects_invalid_capacity_deadline_before_create( - monkeypatch, value -): - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", value) - - class UnexpectedClient: - async def create(self, *_args): - raise AssertionError("invalid retry configuration must fail first") - - with pytest.raises(ValueError, match="SECRLENV_CAPACITY_MAX_WAIT_SECONDS"): - asyncio.run( - agent._create_with_capacity_retry( - UnexpectedClient(), "CVE-2024-1234", "l2" - ) - ) diff --git a/tests/test_secrlenv_codex_harness.py b/tests/test_secrlenv_codex_harness.py deleted file mode 100644 index c7b1136ba..000000000 --- a/tests/test_secrlenv_codex_harness.py +++ /dev/null @@ -1,1790 +0,0 @@ -from __future__ import annotations - -import asyncio -import copy -import errno -import hashlib -import json -import os -import shutil -import subprocess -import sys -from collections.abc import Iterable -from pathlib import Path -from types import SimpleNamespace -from typing import Any - -import aiohttp -import pytest -from aiohttp import web - - -from yeto.rl.harness.codex import codex_harness_agent as harness -from yeto.rl.harness.codex import reward as secrlenv_reward - - -def test_stock_codex_qwen38_adapter_process_binds_exact_xhigh_profile(): - environment = dict(os.environ) - environment["YETO_CODEX_CHAT_TEMPLATE"] = "qwen38" - result = subprocess.run( - [ - sys.executable, - "-c", - """ -from yeto.rl.harness.codex import codex_harness_agent as adapter -assert adapter.BACKEND_MODEL == "qwen38" -assert adapter.BACKEND_REASONING_EFFORT == "xhigh" -assert adapter.BACKEND_CHAT_TEMPLATE_KWARGS == { - "enable_thinking": True, - "preserve_thinking": True, - "reasoning_effort": "xhigh", -} -""", - ], - check=False, - capture_output=True, - text=True, - timeout=15, - env=environment, - ) - assert result.returncode == 0, result.stderr - - -def test_stock_codex_qwen35_adapter_process_binds_exact_fixed_profile(): - environment = dict(os.environ) - environment["YETO_CODEX_CHAT_TEMPLATE"] = "qwen35" - result = subprocess.run( - [ - sys.executable, - "-c", - """ -from yeto.rl.harness.codex import codex_harness_agent as adapter -assert adapter.BACKEND_MODEL == "qwen35" -assert adapter.BACKEND_REASONING_EFFORT == "xhigh" -assert adapter.BACKEND_CHAT_TEMPLATE_KWARGS == {"clear_thinking": False} -""", - ], - check=False, - capture_output=True, - text=True, - timeout=15, - env=environment, - ) - assert result.returncode == 0, result.stderr - - -def _completion( - name: str, - raw_arguments: str, - call_id: str, - reasoning: str, - total_tokens: int, -) -> dict[str, Any]: - return { - "id": f"chatcmpl-{call_id}", - "object": "chat.completion", - "created": 1, - "model": harness.BACKEND_MODEL, - "choices": [ - { - "index": 0, - "finish_reason": "tool_calls", - "message": { - "role": "assistant", - "content": "", - "reasoning_content": reasoning, - "tool_calls": [ - { - "id": call_id, - "type": "function", - "function": { - "name": name, - "arguments": raw_arguments, - }, - } - ], - }, - } - ], - "usage": { - "prompt_tokens": total_tokens - 5, - "completion_tokens": 5, - "total_tokens": total_tokens, - }, - } - - -def _fake_miles_sse(completion: dict[str, Any]) -> bytes: - choice = completion["choices"][0] - message = choice["message"] - delta = { - "role": message.get("role", "assistant"), - "content": message.get("content"), - } - if message.get("reasoning_content") is not None: - delta["reasoning_content"] = message["reasoning_content"] - if message.get("tool_calls"): - delta["tool_calls"] = [ - {**tool_call, "index": index} - for index, tool_call in enumerate(message["tool_calls"]) - ] - chunk = { - "id": completion.get("id"), - "object": "chat.completion.chunk", - "created": completion.get("created", 1), - "model": completion.get("model", harness.BACKEND_MODEL), - "choices": [ - { - "index": 0, - "delta": delta, - "finish_reason": choice.get("finish_reason"), - } - ], - "usage": completion.get("usage"), - } - return ( - b"data: " - + json.dumps( - chunk, - ensure_ascii=False, - allow_nan=False, - separators=(",", ":"), - ).encode() - + b"\n\ndata: [DONE]\n\n" - ) - - -async def _fake_miles( - completions: Iterable[dict[str, Any]], -) -> tuple[web.AppRunner, str, list[dict[str, Any]]]: - queue = list(completions) - requests: list[dict[str, Any]] = [] - - async def chat(request: web.Request) -> web.Response: - requests.append(await request.json()) - if not queue: - return web.json_response({"error": "unexpected sample"}, status=409) - return web.Response( - body=_fake_miles_sse(queue.pop(0)), - content_type="text/event-stream", - ) - - app = web.Application() - app.router.add_post("/v1/chat/completions", chat) - runner = web.AppRunner(app, access_log=None) - await runner.setup() - site = web.TCPSite(runner, "127.0.0.1", 0) - await site.start() - sockets = site._server.sockets if site._server is not None else [] - assert len(sockets) == 1 - return runner, f"http://127.0.0.1:{sockets[0].getsockname()[1]}", requests - - -async def _fake_miles_wire( - body: bytes, - *, - content_type: str = "text/event-stream", -) -> tuple[web.AppRunner, str]: - async def chat(_request: web.Request) -> web.Response: - return web.Response(body=body, headers={"Content-Type": content_type}) - - app = web.Application() - app.router.add_post("/v1/chat/completions", chat) - runner = web.AppRunner(app, access_log=None) - await runner.setup() - site = web.TCPSite(runner, "127.0.0.1", 0) - await site.start() - sockets = site._server.sockets if site._server is not None else [] - assert len(sockets) == 1 - return runner, f"http://127.0.0.1:{sockets[0].getsockname()[1]}" - - -def _codex_body(input_items: list[dict[str, Any]]) -> dict[str, Any]: - return { - "model": harness.BACKEND_MODEL, - "instructions": harness.BASE_INSTRUCTIONS, - "input": input_items, - "parallel_tool_calls": False, - "reasoning": {"effort": "xhigh", "summary": "none"}, - "store": False, - "stream": True, - "tool_choice": "auto", - "tools": [ - { - "type": "function", - "name": "update_plan", - "description": "Update the task plan.", - "strict": False, - "parameters": {"type": "object"}, - }, - *copy.deepcopy(harness._EXPECTED_RESPONSE_TOOLS), - ], - } - - -def _sse_events(raw: str) -> list[dict[str, Any]]: - events: list[dict[str, Any]] = [] - for block in raw.split("\n\n"): - data = [line[6:] for line in block.splitlines() if line.startswith("data: ")] - if data and data[0] != "[DONE]": - events.append(json.loads("\n".join(data))) - return events - - -def _set_bridge_env(monkeypatch: pytest.MonkeyPatch, max_tokens: int = 256) -> None: - monkeypatch.setenv("YETO_CODEX_BACKEND_MAX_TOKENS", str(max_tokens)) - - -def test_identity_is_live_exact_and_fail_closed(monkeypatch): - identity = harness.codex_harness_identity() - assert set(identity) == { - "base_instructions_sha256", - "terminal_exec_tool_schema_sha256", - "submit_tool_schema_sha256", - "dynamic_tools_schema_sha256", - } - assert all(len(value) == 64 and int(value, 16) >= 0 for value in identity.values()) - assert harness._MILES_TOOLS == harness.legacy.TOOLS - assert "model-facing tools, `terminal.exec` and `submit`" in ( - harness.BASE_INSTRUCTIONS - ) - assert "Codex's internal\n`terminal_exec` dynamic-tool alias" in ( - harness.BASE_INSTRUCTIONS - ) - monkeypatch.setattr(harness, "BASE_INSTRUCTIONS", harness.BASE_INSTRUCTIONS + " drift") - with pytest.raises(harness.CodexHarnessError, match="identity constants drifted"): - harness.codex_harness_identity() - - -def test_runtime_attests_live_binary_version_and_all_env(tmp_path, monkeypatch): - harness._RUNTIME_ATTESTATION_CACHE.clear() - binary = tmp_path / "codex-fake" - binary.write_text( - "#!/bin/sh\n" - "if [ \"$1\" = \"--version\" ]; then\n" - " printf 'codex-cli 0.145.0\\n'\n" - " exit 0\n" - "fi\n" - "exit 1\n" - ) - binary.chmod(0o700) - for name, value in harness._IDENTITY_ENV.items(): - monkeypatch.setenv(name, value) - monkeypatch.setenv("YETO_CODEX_BINARY_PATH", str(binary)) - monkeypatch.setenv( - "YETO_CODEX_BINARY_SHA256", hashlib.sha256(binary.read_bytes()).hexdigest() - ) - monkeypatch.setenv("YETO_CODEX_BINARY_SIZE_BYTES", str(binary.stat().st_size)) - monkeypatch.setenv("YETO_CODEX_VERSION", harness.CODEX_CLI_VERSION) - monkeypatch.setenv("YETO_CODEX_BACKEND_MAX_TOKENS", "256") - real_run = harness.subprocess.run - version_calls = 0 - - def counted_run(*args, **kwargs): - nonlocal version_calls - version_calls += 1 - return real_run(*args, **kwargs) - - monkeypatch.setattr(harness.subprocess, "run", counted_run) - assert harness._attest_runtime() == binary - assert harness._attest_runtime() == binary - assert version_calls == 1 - - binary.write_text(binary.read_text().replace("0.145.0", "0.145.1")) - monkeypatch.setenv( - "YETO_CODEX_BINARY_SHA256", hashlib.sha256(binary.read_bytes()).hexdigest() - ) - monkeypatch.setenv("YETO_CODEX_BINARY_SIZE_BYTES", str(binary.stat().st_size)) - with pytest.raises(harness.CodexHarnessError, match="version drifted"): - harness._attest_runtime() - assert version_calls == 2 - harness._RUNTIME_ATTESTATION_CACHE.clear() - - -def test_history_and_responses_contract_reject_compaction_and_extra_tools(): - with pytest.raises(harness.CodexHarnessError, match="compaction"): - harness._canonical_history_item({"type": "compaction"}) - body = _codex_body( - [ - { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "task"}], - } - ] - ) - body["tools"].append( - { - "type": "function", - "name": "shell", - "strict": False, - "parameters": {"type": "object"}, - } - ) - with pytest.raises(harness.CodexHarnessError, match="tool surface drifted"): - harness._validate_codex_request(body, harness.BACKEND_MODEL) - body = _codex_body(body["input"]) - body["reasoning"]["effort"] = "max" - with pytest.raises(harness.CodexHarnessError, match="xhigh"): - harness._validate_codex_request(body, harness.BACKEND_MODEL) - - -def test_sampling_cannot_override_signed_miles_fields(monkeypatch): - _set_bridge_env(monkeypatch) - metrics = harness.legacy.AgentMetrics() - with pytest.raises(harness.CodexHarnessError, match="override"): - harness._ResponsesBridge( - "http://127.0.0.1:1", - "task", - {"max_tokens": 64, "parallel_tool_calls": True}, - metrics, - max_seq_len=None, - ) - with pytest.raises(harness.CodexHarnessError, match="unknown sampling"): - harness._ResponsesBridge( - "http://127.0.0.1:1", - "task", - {"max_tokens": 64, "untrusted_body_field": "value"}, - metrics, - max_seq_len=None, - ) - - -def test_miles_fake_stream_round_trips_usage_reasoning_and_tool_call(): - raw_arguments = '{ "command": "printf café" }' - completion = _completion( - "terminal.exec", - raw_arguments, - "call-stream", - "inspect café", - 321, - ) - - normalized = harness._parse_miles_sse(_fake_miles_sse(completion)) - - assert normalized["object"] == "chat.completion" - assert normalized["usage"] == completion["usage"] - choice = normalized["choices"][0] - assert choice["finish_reason"] == "tool_calls" - assert choice["message"] == { - "role": "assistant", - "content": "", - "reasoning_content": "inspect café", - "tool_calls": [ - { - "id": "call-stream", - "type": "function", - "function": { - "name": "terminal.exec", - "arguments": raw_arguments, - }, - } - ], - } - - -def test_completion_accepts_parser_whitespace_but_rejects_real_prose(monkeypatch): - _set_bridge_env(monkeypatch) - whitespace_completion = _completion( - "terminal.exec", - '{"command":"id"}', - "call-whitespace", - "inspect", - 32, - ) - whitespace_completion["choices"][0]["message"]["content"] = "\n\n" - bridge = harness._ResponsesBridge( - "http://127.0.0.1:1", - "task", - {"max_tokens": 64}, - harness.legacy.AgentMetrics(), - max_seq_len=None, - ) - - output = bridge._translate_completion(whitespace_completion) - - assert output[-1]["type"] == "function_call" - assert output[-1]["name"] == "terminal_exec" - assert bridge._messages[-1]["content"] == "\n\n" - - prose_completion = copy.deepcopy(whitespace_completion) - prose_completion["choices"][0]["message"]["content"] = "I will describe it." - prose_bridge = harness._ResponsesBridge( - "http://127.0.0.1:1", - "task", - {"max_tokens": 64}, - harness.legacy.AgentMetrics(), - max_seq_len=None, - ) - with pytest.raises(harness.CodexModelFailure, match="mixed prose"): - prose_bridge._translate_completion(prose_completion) - - -def test_only_opted_in_profiles_keep_prose_before_the_tool_call(monkeypatch): - from yeto.rl.codex_backend import _PROFILES - - assert [name for name, p in _PROFILES.items() if p.get("allow_prose_with_tool_call")] == ["qwen38"] - assert harness.ALLOW_PROSE_WITH_TOOL_CALL is False # module default profile stays strict - _set_bridge_env(monkeypatch) - completion = _completion("terminal.exec", '{"command":"id"}', "call-prose", "inspect", 12) - completion["choices"][0]["message"]["content"] = "I found the cause; checking it now." - monkeypatch.setattr(harness, "ALLOW_PROSE_WITH_TOOL_CALL", True) - bridge = harness._ResponsesBridge("http://127.0.0.1:1", "task", {"max_tokens": 64}, - harness.legacy.AgentMetrics(), max_seq_len=None) - output = bridge._translate_completion(copy.deepcopy(completion)) - assert output[-1]["type"] == "function_call" and output[-1]["name"] == "terminal_exec" - assert bridge._messages[-1]["content"] == "I found the cause; checking it now." - - -def test_miles_fake_stream_preserves_length_as_sequence_policy_boundary(monkeypatch): - _set_bridge_env(monkeypatch) - completion = _completion( - "terminal.exec", - '{"command":"id"}', - "call-length", - "inspect", - 321, - ) - completion["choices"][0]["finish_reason"] = "length" - - normalized = harness._parse_miles_sse(_fake_miles_sse(completion)) - - assert normalized["choices"][0]["finish_reason"] == "length" - bridge = harness._ResponsesBridge( - "http://127.0.0.1:1", - "task", - {"max_tokens": 64}, - harness.legacy.AgentMetrics(), - max_seq_len=512, - ) - with pytest.raises(harness.CodexSequenceLimit, match="truncated"): - bridge._translate_completion(normalized) - - -def test_miles_fake_stream_rejects_malformed_framing_and_shapes(): - valid = _fake_miles_sse( - _completion( - "terminal.exec", - '{"command":"id"}', - "call-stream", - "inspect", - 32, - ) - ) - json_event, done_event, _empty = valid.split(b"\n\n") - chunk = json.loads(json_event[len(b"data: ") :]) - invalid_chunks = [] - for mutate in ( - lambda value: value.update({"extra": True}), - lambda value: value.update({"choices": []}), - lambda value: value["choices"][0].update({"index": False}), - lambda value: value["choices"][0]["delta"]["tool_calls"][0].update( - {"index": False} - ), - lambda value: value["choices"][0].update({"delta": {"role": "assistant"}}), - ): - value = copy.deepcopy(chunk) - mutate(value) - invalid_chunks.append( - b"data: " - + json.dumps(value, separators=(",", ":")).encode() - + b"\n\ndata: [DONE]\n\n" - ) - malformed = [ - b"", - b"data: [DONE]\n\n", - json_event + b"\n\n", - done_event + b"\n\n" + json_event + b"\n\n", - json_event + b"\n\n" + done_event + b"\n\n" + done_event + b"\n\n", - json_event + b"\n\n" + json_event + b"\n\n" + done_event + b"\n\n", - b"data: \xff\n\ndata: [DONE]\n\n", - b"data: {not-json}\n\ndata: [DONE]\n\n", - b'data: {"id":"first","id":"second"}\n\ndata: [DONE]\n\n', - b'data: {"created":NaN}\n\ndata: [DONE]\n\n', - *invalid_chunks, - ] - - for raw in malformed: - with pytest.raises(harness.CodexHarnessError, match="Miles session returned"): - harness._parse_miles_sse(raw) - - -def test_miles_fake_stream_omits_large_tito_metadata_from_client(monkeypatch): - _set_bridge_env(monkeypatch, max_tokens=32_768) - - async def scenario() -> None: - completion = _completion( - "terminal.exec", - '{"command":"id"}', - "call-large-record", - "inspect", - 32_768, - ) - completion["choices"][0]["meta_info"] = { - "output_token_logprobs": "x" * (harness.MAX_MILES_RESPONSE_BYTES + 1) - } - assert len(json.dumps(completion).encode()) > (harness.MAX_MILES_RESPONSE_BYTES) - runner, miles_url, miles_requests = await _fake_miles([completion]) - try: - async with harness._ResponsesBridge( - miles_url, - "task", - {"max_tokens": 32_768}, - harness.legacy.AgentMetrics(), - max_seq_len=65_536, - ) as bridge: - normalized = await bridge._sample_miles() - finally: - await runner.cleanup() - - assert normalized["choices"][0]["message"]["reasoning_content"] == ("inspect") - assert "meta_info" not in normalized["choices"][0] - assert miles_requests[0]["stream"] is True - - asyncio.run(scenario()) - - -def test_miles_response_reader_rejects_exact_cap_plus_one_without_unbounded_read(): - class NeverReadContent: - calls = 0 - - def iter_chunked(self, _size): - self.calls += 1 - raise AssertionError("oversized Content-Length must not be read") - - async def read(self): - raise AssertionError("unbounded read is forbidden") - - declared_content = NeverReadContent() - declared = SimpleNamespace( - content_length=harness.MAX_MILES_RESPONSE_BYTES + 1, - content=declared_content, - ) - with pytest.raises(harness.CodexHarnessError, match="oversized"): - asyncio.run(harness._read_bounded_miles_response(declared)) - assert declared_content.calls == 0 - - class ChunkedContent: - def __init__(self): - self.yields = 0 - self.read_past_limit = False - - def iter_chunked(self, size): - assert size == harness._MILES_RESPONSE_READ_CHUNK_BYTES - - async def chunks(): - block = b"x" * size - for _ in range(harness.MAX_MILES_RESPONSE_BYTES // size): - self.yields += 1 - yield block - self.yields += 1 - yield b"x" - self.read_past_limit = True - yield b"must-not-be-read" - - return chunks() - - async def read(self): - raise AssertionError("unbounded read is forbidden") - - chunked_content = ChunkedContent() - chunked = SimpleNamespace(content_length=None, content=chunked_content) - with pytest.raises(harness.CodexHarnessError, match="oversized"): - asyncio.run(harness._read_bounded_miles_response(chunked)) - assert chunked_content.yields == 65 - assert chunked_content.read_past_limit is False - - -def test_miles_non_200_response_is_closed_without_reading_body(monkeypatch): - _set_bridge_env(monkeypatch) - - class NeverReadContent: - calls = 0 - - def iter_chunked(self, _size): - self.calls += 1 - raise AssertionError("non-200 response body must not be read") - - async def read(self): - raise AssertionError("unbounded read is forbidden") - - class Response: - status = 503 - content_type = "application/json" - content_length = harness.MAX_MILES_RESPONSE_BYTES + 1 - content = NeverReadContent() - closed = False - - async def __aenter__(self): - return self - - async def __aexit__(self, *_exc): - self.closed = True - - class Session: - def __init__(self, response): - self.response = response - self.payload = None - self.headers = None - - def post(self, _url, *, json, headers): - self.payload = json - self.headers = headers - return self.response - - async def scenario() -> None: - response = Response() - session = Session(response) - bridge = harness._ResponsesBridge( - "http://127.0.0.1:1", - "task", - {"max_tokens": 64}, - harness.legacy.AgentMetrics(), - max_seq_len=512, - ) - bridge._session = session - with pytest.raises(harness.CodexHarnessError, match="HTTP 503"): - await bridge._sample_miles() - assert session.payload["stream"] is True - assert session.headers == {} - assert response.content.calls == 0 - assert response.closed is True - - asyncio.run(scenario()) - - -def test_miles_sample_accepts_content_type_parameters_and_enforces_cap(monkeypatch): - _set_bridge_env(monkeypatch) - - async def accepted() -> None: - completion = _completion( - "terminal.exec", - '{"command":"id"}', - "call-media-type", - "inspect", - 32, - ) - runner, miles_url = await _fake_miles_wire( - _fake_miles_sse(completion), - content_type="text/event-stream; charset=utf-8", - ) - try: - async with harness._ResponsesBridge( - miles_url, - "task", - {"max_tokens": 64}, - harness.legacy.AgentMetrics(), - max_seq_len=512, - ) as bridge: - assert await bridge._sample_miles() == harness._parse_miles_sse( - _fake_miles_sse(completion) - ) - finally: - await runner.cleanup() - - async def rejected(body: bytes, content_type: str, match: str) -> None: - runner, miles_url = await _fake_miles_wire( - body, - content_type=content_type, - ) - try: - async with harness._ResponsesBridge( - miles_url, - "task", - {"max_tokens": 64}, - harness.legacy.AgentMetrics(), - max_seq_len=512, - ) as bridge: - with pytest.raises(harness.CodexHarnessError, match=match): - await bridge._sample_miles() - finally: - await runner.cleanup() - - asyncio.run(accepted()) - asyncio.run(rejected(b"{}", "application/json", "media type")) - asyncio.run( - rejected( - b"x" * (harness.MAX_MILES_RESPONSE_BYTES + 1), - "text/event-stream", - "oversized", - ) - ) - - -def test_fake_codex_bridge_preserves_history_and_filters_only_update_plan(monkeypatch): - _set_bridge_env(monkeypatch) - - async def scenario() -> None: - raw_terminal = '{ "command": "id" }' - runner, miles_url, miles_requests = await _fake_miles( - [ - _completion( - "terminal.exec", raw_terminal, "call-terminal", "inspect", 25 - ), - _completion( - "submit", - '{"evidence":"uid=1000"}', - "call-submit", - "finish", - 40, - ), - ] - ) - metrics = harness.legacy.AgentMetrics() - try: - async with harness._ResponsesBridge( - miles_url, - "solve target", - { - "temperature": 0.7, - "top_p": 0.9, - "max_tokens": 128, - "stream": True, - "n": 1, - }, - metrics, - max_seq_len=512, - ) as bridge: - headers = {"Authorization": f"Bearer {bridge.token}"} - initial = [ - { - "type": "message", - "role": "user", - "content": [ - {"type": "input_text", "text": "solve target"} - ], - } - ] - async with aiohttp.ClientSession() as session: - async with session.post( - f"{bridge.url}/v1/responses", - json=_codex_body(initial), - headers=headers, - ) as response: - assert response.status == 200 - first_events = _sse_events(await response.text()) - completed = first_events[-1]["response"] - function_call = completed["output"][1] - assert function_call["name"] == "terminal_exec" - assert function_call["arguments"] == raw_terminal - assert function_call["call_id"] == "call-terminal" - - rendered = '{"exit_code":0,"output":"uid=1000","timed_out":false,"truncated":false}' - bridge.expect_tool_output("call-terminal", rendered) - second_input = copy.deepcopy(bridge._expected_input) - assert second_input is not None - async with session.post( - f"{bridge.url}/v1/responses", - json=_codex_body(second_input), - headers=headers, - ) as response: - assert response.status == 200 - second_events = _sse_events(await response.text()) - submit = second_events[-1]["response"]["output"][1] - assert submit["name"] == "submit" - assert submit["call_id"] == "call-submit" - bridge.mark_terminal() - finally: - await runner.cleanup() - - assert len(miles_requests) == 2 - assert [tool["function"]["name"] for tool in miles_requests[0]["tools"]] == [ - "terminal.exec", - "submit", - ] - assert miles_requests[0]["max_tokens"] == 128 - assert miles_requests[0]["temperature"] == 0.7 - assert miles_requests[0]["top_p"] == 0.9 - assert miles_requests[0]["stream"] is True - assert miles_requests[0]["parallel_tool_calls"] is False - assert miles_requests[0]["reasoning_effort"] == "max" - assert miles_requests[0]["thinking"] == {"type": "enabled"} - assert miles_requests[0]["chat_template_kwargs"] == ( - harness.BACKEND_CHAT_TEMPLATE_KWARGS - ) - assert miles_requests[1]["messages"][2] == { - "role": "assistant", - "content": "", - "reasoning_content": "inspect", - "tool_calls": [ - { - "id": "call-terminal", - "type": "function", - "function": { - "name": "terminal.exec", - "arguments": raw_terminal, - }, - } - ], - } - assert miles_requests[1]["messages"][3] == { - "role": "tool", - "tool_call_id": "call-terminal", - "content": rendered, - } - assert metrics.turns == 2 - assert metrics.tool_calls == 2 - - asyncio.run(scenario()) - - -def test_bridge_compaction_is_same_session_trainable_and_keeps_atomic_tail(monkeypatch): - _set_bridge_env(monkeypatch, max_tokens=128) - monkeypatch.setenv("YETO_CODEX_COMPACTION_ENABLED", "1") - monkeypatch.setenv("YETO_CODEX_COMPACTION_TRIGGER_TOKENS", "300") - monkeypatch.setenv("YETO_CODEX_COMPACTION_SUMMARY_MAX_TOKENS", "128") - monkeypatch.setenv("YETO_CODEX_MAX_COMPACTIONS", "3") - - summary = { - "id": "chatcmpl-summary", - "object": "chat.completion", - "created": 1, - "model": harness.BACKEND_MODEL, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Goal: solve target. Evidence: uid=1000. Next: submit.", - "reasoning_content": "compress state", - }, - } - ], - "usage": {"prompt_tokens": 50, "completion_tokens": 10, "total_tokens": 60}, - } - - async def scenario() -> None: - requests: list[tuple[dict[str, Any], dict[str, str]]] = [] - queue = [ - _completion("terminal.exec", '{"command":"id"}', "call-1", "inspect", 25), - summary, - _completion("submit", '{"evidence":"uid=1000"}', "call-2", "finish", 40), - ] - - async def chat(request: web.Request) -> web.Response: - requests.append((await request.json(), dict(request.headers))) - return web.Response( - body=_fake_miles_sse(queue.pop(0)), - content_type="text/event-stream", - ) - - app = web.Application() - app.router.add_post("/v1/chat/completions", chat) - runner = web.AppRunner(app, access_log=None) - await runner.setup() - site = web.TCPSite(runner, "127.0.0.1", 0) - await site.start() - sockets = site._server.sockets if site._server is not None else [] - miles_url = f"http://127.0.0.1:{sockets[0].getsockname()[1]}" - try: - async with harness._ResponsesBridge( - miles_url, - "solve target", - {"max_tokens": 128}, - harness.legacy.AgentMetrics(), - max_seq_len=2048, - ) as bridge: - headers = {"Authorization": f"Bearer {bridge.token}"} - initial = [ - { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "solve target"}], - } - ] - async with aiohttp.ClientSession() as session: - async with session.post( - f"{bridge.url}/v1/responses", - json=_codex_body(initial), - headers=headers, - ) as response: - assert response.status == 200 - bridge.expect_tool_output("call-1", '{"output":"uid=1000"}') - second_input = copy.deepcopy(bridge._expected_input) - async with session.post( - f"{bridge.url}/v1/responses", - json=_codex_body(second_input), - headers=headers, - ) as response: - assert response.status == 200 - finally: - await runner.cleanup() - - assert len(requests) == 3 - _first, first_headers = requests[0] - summary_request, summary_headers = requests[1] - resumed, resumed_headers = requests[2] - assert first_headers["X-Miles-Compaction-Segment-Type"] == "execution" - assert summary_headers["X-Miles-Compaction-Segment-Type"] == "summary" - assert summary_headers["X-Miles-Compaction-Segment-Index"] == "1" - assert summary_headers["X-Miles-Compaction-Context-Budget"] == "2048" - assert summary_request["tools"] == [] - assert summary_request["max_tokens"] == 128 - assert summary_request["messages"][-1]["content"] == harness.COMPACTION_SUMMARY_INSTRUCTION - assert resumed_headers["X-Miles-Compaction-Context-Window"] == "1" - assert resumed_headers["X-Miles-Compaction-Segment-Index"] == "2" - assert resumed_headers["X-Miles-Compaction-Segment-Type"] == "execution" - assert resumed["messages"][0] == {"role": "system", "content": harness.BASE_INSTRUCTIONS} - assert "uid=1000" in resumed["messages"][1]["content"] - assert resumed["messages"][2:] == summary_request["messages"][2:-1] - - asyncio.run(scenario()) - - -def test_bridge_compaction_trigger_is_consumed_context_boundary(monkeypatch): - _set_bridge_env(monkeypatch, max_tokens=1024) - monkeypatch.setenv("YETO_CODEX_COMPACTION_ENABLED", "1") - monkeypatch.setenv("YETO_CODEX_COMPACTION_TRIGGER_TOKENS", "6144") - monkeypatch.setenv("YETO_CODEX_COMPACTION_SUMMARY_MAX_TOKENS", "1024") - monkeypatch.setenv("YETO_CODEX_MAX_COMPACTIONS", "3") - - bridge = harness._ResponsesBridge( - "http://127.0.0.1:1", - "boundary", - {"max_tokens": 1024}, - harness.legacy.AgentMetrics(), - max_seq_len=8192, - ) - - assert not bridge._should_compact(6143) - assert bridge._should_compact(6144) - assert ( - 6144 - + harness.COMPACTION_SUMMARY_INSTRUCTION_TOKEN_RESERVE - + bridge._compaction_summary_max_tokens - <= 8192 - ) - - -def test_bridge_compaction_rejects_trigger_without_summary_reserve(monkeypatch): - _set_bridge_env(monkeypatch, max_tokens=512) - monkeypatch.setenv("YETO_CODEX_COMPACTION_ENABLED", "1") - monkeypatch.setenv("YETO_CODEX_COMPACTION_TRIGGER_TOKENS", "500") - monkeypatch.setenv("YETO_CODEX_COMPACTION_SUMMARY_MAX_TOKENS", "512") - - with pytest.raises(harness.CodexHarnessError, match="insufficient room"): - harness._ResponsesBridge( - "http://127.0.0.1:1", - "reserve", - {"max_tokens": 512}, - harness.legacy.AgentMetrics(), - max_seq_len=1024, - ) - - -def test_bridge_rejects_history_mutation_without_an_extra_sample(monkeypatch): - _set_bridge_env(monkeypatch) - - async def scenario() -> None: - runner, miles_url, miles_requests = await _fake_miles( - [_completion("terminal.exec", '{"command":"id"}', "call-1", "think", 10)] - ) - try: - async with harness._ResponsesBridge( - miles_url, - "immutable task", - {"max_tokens": 64}, - harness.legacy.AgentMetrics(), - max_seq_len=None, - ) as bridge: - headers = {"Authorization": f"Bearer {bridge.token}"} - initial = [ - { - "type": "message", - "role": "user", - "content": [ - {"type": "input_text", "text": "immutable task"} - ], - } - ] - async with aiohttp.ClientSession() as session: - async with session.post( - f"{bridge.url}/v1/responses", - json=_codex_body(initial), - headers=headers, - ) as response: - assert response.status == 200 - bridge.expect_tool_output("call-1", "tool result") - mutated = copy.deepcopy(bridge._expected_input) - assert mutated is not None - mutated[0]["content"][0]["text"] = "changed task" - async with session.post( - f"{bridge.url}/v1/responses", - json=_codex_body(mutated), - headers=headers, - ) as response: - assert response.status == 400 - assert isinstance(bridge.fatal.result(), harness.CodexHarnessError) - finally: - await runner.cleanup() - assert len(miles_requests) == 1 - - asyncio.run(scenario()) - - -class _FakeEpisodeClient: - def __init__(self, *, output: str = "uid=1000", truncated: bool = False) -> None: - self.output = output - self.truncated = truncated - self.exec_calls: list[tuple[str, str]] = [] - self.submissions: list[tuple[str, dict[str, Any]]] = [] - - async def execute(self, episode_id: str, command: str, **_kwargs: Any): - self.exec_calls.append((episode_id, command)) - return { - "exit_code": 0, - "output": self.output, - "timed_out": False, - "truncated": self.truncated, - } - - async def submit(self, episode_id: str, value: dict[str, Any]): - self.submissions.append((episode_id, value)) - return {"accepted": True} - - -class _RecoveringEpisodeClient(_FakeEpisodeClient): - def __init__(self) -> None: - super().__init__() - self.execute_attempts = 0 - self.submit_attempts = 0 - - async def execute(self, episode_id: str, command: str, **kwargs: Any): - self.execute_attempts += 1 - if self.execute_attempts == 1: - self.exec_calls.append((episode_id, command)) - raise harness.EpisodeAPIError(400, "invalid_command", "command rejected") - return await super().execute(episode_id, command, **kwargs) - - async def submit(self, episode_id: str, value: dict[str, Any]): - self.submit_attempts += 1 - if self.submit_attempts == 1: - self.submissions.append((episode_id, value)) - raise harness.EpisodeAPIError(400, "invalid_submission", "evidence rejected") - return await super().submit(episode_id, value) - - -def test_codex_exhausted_precreate_503_returns_signed_admission(monkeypatch): - attempts = [] - - class AdmissionClient: - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - async def create(self, task_id, tier): - attempts.append((task_id, tier)) - raise harness.EpisodeAPIError( - 503, "infrastructure_error", "unavailable" - ) - - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "30") - monkeypatch.setattr(harness, "EpisodeClient", AdmissionClient) - monkeypatch.setattr(harness, "_attest_runtime", lambda: Path("/codex")) - monkeypatch.setattr(harness.legacy.random, "uniform", lambda *_args: 0.0) - - result = asyncio.run( - harness.run( - "http://127.0.0.1:1", - None, - metadata={"task_id": "CVE-2024-1234", "prompt_tier": "l2"}, - ) - ) - - assert result is not None - outcome, value = secrlenv_reward._verified_outcome(result) - assert len(attempts) == 2 - assert outcome["schema"] == 2 - assert outcome["episode_id"] is None - assert outcome[harness.INFRASTRUCTURE_503_RETRIES_KEY] == 1 - assert value == 0.0 - harness.legacy.require_no_episode_residue() - - -def test_codex_precreate_infrastructure_retry_timeout_is_not_a_verdict( - monkeypatch, -): - attempts = 0 - - class AdmissionTimeoutClient: - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - async def create(self, _task_id, _tier): - nonlocal attempts - attempts += 1 - if attempts == 1: - raise harness.EpisodeAPIError( - 503, "infrastructure_error", "unavailable" - ) - await asyncio.Event().wait() - - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "0.01") - monkeypatch.setattr(harness, "EpisodeClient", AdmissionTimeoutClient) - monkeypatch.setattr(harness, "_attest_runtime", lambda: Path("/codex")) - monkeypatch.setattr(harness.legacy.random, "uniform", lambda *_args: 0.0) - - result = asyncio.run( - harness.run( - "http://127.0.0.1:1", - None, - metadata={"task_id": "CVE-2024-1234", "prompt_tier": "l2"}, - ) - ) - - assert result is None - assert attempts == 2 - harness.legacy.require_no_episode_residue() - - -def test_run_does_not_start_codex_before_capacity_retry_creates_episode( - monkeypatch, -): - episode_id = "e" * 24 - events: list[str] = [] - - class CapacityThenReadyClient: - def __init__(self) -> None: - self.create_calls = 0 - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - async def create(self, task_id: str, tier: str): - assert (task_id, tier) == ("CVE-2024-1234", "l2") - self.create_calls += 1 - events.append("create") - if self.create_calls == 1: - raise harness.EpisodeAPIError( - 503, "capacity_reached", "busy" - ) - return { - "episode_id": episode_id, - "prompt": "immutable task", - } - - async def evaluate(self, value: str): - assert value == episode_id - events.append("evaluate") - return { - "task_id": "CVE-2024-1234", - "episode_id": episode_id, - "reward": 1.0, - "passed": True, - } - - async def close(self, value: str): - assert value == episode_id - events.append("close") - return {"closed": True} - - async def drive_after_create( - _binary, - _base_url, - _client, - episode, - _request_kwargs, - _metrics, - *, - max_seq_len, - ): - assert episode["episode_id"] == episode_id - assert max_seq_len is None - assert events == ["create", "create"] - events.append("drive") - return "completed" - - monkeypatch.setenv("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", "30") - monkeypatch.setattr(harness, "EpisodeClient", CapacityThenReadyClient) - monkeypatch.setattr(harness, "_attest_runtime", lambda: Path("/codex")) - monkeypatch.setattr(harness, "_drive_codex", drive_after_create) - monkeypatch.setattr(harness.legacy.random, "uniform", lambda *_args: 0.0) - try: - result = asyncio.run( - harness.run( - "http://127.0.0.1:1", - None, - metadata={"task_id": "CVE-2024-1234", "prompt_tier": "l2"}, - ) - ) - finally: - harness.legacy._release_episode(episode_id) - - assert result is not None - assert result["exit_status"] == "completed" - assert events == ["create", "create", "drive", "evaluate", "close"] - - -def test_run_close_failure_returns_terminal_cleanup_evidence(monkeypatch): - episode_id = "d" * 24 - close_calls = [] - - class CloseFailingClient: - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - async def create(self, _task_id: str, _tier: str): - return {"episode_id": episode_id, "prompt": "immutable task"} - - async def evaluate(self, value: str): - assert value == episode_id - return { - "task_id": "CVE-2024-1234", - "episode_id": episode_id, - "reward": 1.0, - "passed": True, - } - - async def close(self, value: str): - assert value == episode_id - close_calls.append(value) - raise harness.EpisodeClientError("close failed") - - async def complete_policy(*_args, **_kwargs): - return "completed" - - monkeypatch.setattr(harness, "EpisodeClient", CloseFailingClient) - monkeypatch.setattr(harness, "_attest_runtime", lambda: Path("/codex")) - monkeypatch.setattr(harness, "_drive_codex", complete_policy) - try: - result = asyncio.run( - harness.run( - "http://127.0.0.1:1", - None, - metadata={"task_id": "CVE-2024-1234", "prompt_tier": "l2"}, - ) - ) - assert result is not None - assert result["exit_status"] == harness.legacy.CLEANUP_ERROR_STATUS - assert ( - result[harness.OUTCOME_KEY]["status"] - == harness.legacy.CLEANUP_ERROR_STATUS - ) - assert len(close_calls) == 2 - assert harness.legacy._EPISODE_PHASES[episode_id] == "cleanup_pending" - finally: - harness.legacy._release_episode(episode_id) - - -def test_driver_reaps_isolated_process_group_before_home_cleanup(monkeypatch): - events: list[object] = [] - - class FakeProcess: - pid = 4321 - returncode = None - - async def wait(self): - events.append("wait-parent") - self.returncode = 0 - return 0 - - def terminate(self): - raise AssertionError("process-group TERM should be used") - - def kill(self): - raise AssertionError("process-group KILL should be used") - - async def scenario() -> None: - driver = harness._AppServerDriver( - Path("/unused"), - SimpleNamespace(), # type: ignore[arg-type] - _FakeEpisodeClient(), - "episode", - "task", - harness.legacy.AgentMetrics(), - ) - driver._process = FakeProcess() # type: ignore[assignment] - driver._process_group_id = 4321 - monkeypatch.setattr(harness, "_CODEX_DESCENDANT_TERM_GRACE_SECONDS", 0.0) - monkeypatch.setattr( - harness.os, - "killpg", - lambda process_group_id, sig: events.append((process_group_id, sig)), - ) - - await driver._stop() - - assert events == [ - (4321, harness.signal.SIGTERM), - "wait-parent", - (4321, harness.signal.SIGKILL), - ] - assert driver._process_group_id is None - - asyncio.run(scenario()) - - -def test_codex_argv_disables_background_plugin_sync(): - bridge = SimpleNamespace(url="http://127.0.0.1:1234") - - argv = harness._codex_argv(Path("/codex"), bridge) # type: ignore[arg-type] - - assert argv.count("features.plugins=false") == 1 - - -def test_driver_retries_transient_nonempty_isolated_home_cleanup(monkeypatch, tmp_path): - home_path = tmp_path / "codex-home" - home_path.mkdir() - - class RacingHome: - name = str(home_path) - attempts = 0 - - def cleanup(self): - self.attempts += 1 - if self.attempts < 3: - raise OSError(errno.ENOTEMPTY, "Directory not empty", "plugins") - home_path.rmdir() - - async def scenario() -> None: - home = RacingHome() - driver = harness._AppServerDriver( - Path("/unused"), - SimpleNamespace(), # type: ignore[arg-type] - _FakeEpisodeClient(), - "episode", - "task", - harness.legacy.AgentMetrics(), - ) - driver._isolated_home = home # type: ignore[assignment] - monkeypatch.setattr( - harness, - "_CODEX_HOME_CLEANUP_RETRY_DELAYS_SECONDS", - (0.0, 0.0, 0.0), - ) - - await driver.__aexit__(None, None, None) - - assert home.attempts == 3 - assert not home_path.exists() - - asyncio.run(scenario()) - - -def test_driver_home_cleanup_exhaustion_is_best_effort_and_nonfatal( - monkeypatch, tmp_path -): - home_path = tmp_path / "codex-home" - (home_path / "plugins").mkdir(parents=True) - - class PersistentlyRacingHome: - name = str(home_path) - attempts = 0 - - def cleanup(self): - self.attempts += 1 - raise OSError(errno.ENOTEMPTY, "Directory not empty", "plugins") - - async def scenario() -> None: - home = PersistentlyRacingHome() - driver = harness._AppServerDriver( - Path("/unused"), - SimpleNamespace(), # type: ignore[arg-type] - _FakeEpisodeClient(), - "episode", - "task", - harness.legacy.AgentMetrics(), - ) - driver._isolated_home = home # type: ignore[assignment] - monkeypatch.setattr( - harness, - "_CODEX_HOME_CLEANUP_RETRY_DELAYS_SECONDS", - (0.0, 0.0), - ) - - # A resource-hygiene race must not change a successful episode result. - await driver.__aexit__(None, None, None) - - assert home.attempts == 2 - assert not home_path.exists() - - asyncio.run(scenario()) - - -def test_driver_teardown_failures_do_not_mask_active_episode_error( - monkeypatch, tmp_path -): - class NoLaunchDriver(harness._AppServerDriver): - async def __aenter__(self): - return self - - home_path = tmp_path / "codex-home" - home_path.mkdir() - - class FailedHome: - name = str(home_path) - - def cleanup(self): - raise RuntimeError("secondary home cleanup failure") - - async def scenario() -> None: - driver = NoLaunchDriver( - Path("/unused"), - SimpleNamespace(), # type: ignore[arg-type] - _FakeEpisodeClient(), - "episode", - "task", - harness.legacy.AgentMetrics(), - ) - driver._isolated_home = FailedHome() # type: ignore[assignment] - - async def failed_stop(): - raise RuntimeError("secondary shutdown failure") - - monkeypatch.setattr(driver, "_stop", failed_stop) - with pytest.raises(ValueError, match="primary episode failure"): - async with driver: - raise ValueError("primary episode failure") - assert not home_path.exists() - - asyncio.run(scenario()) - - -def test_driver_rejects_mutated_arguments_but_preserves_bounded_truncation(monkeypatch): - _set_bridge_env(monkeypatch) - - async def scenario() -> None: - metrics = harness.legacy.AgentMetrics() - bridge = harness._ResponsesBridge( - "http://127.0.0.1:1", "task", {"max_tokens": 64}, metrics, max_seq_len=None - ) - bridge._pending = harness._PendingTool( - "call-1", "terminal_exec", "terminal.exec", '{"command":"id"}' - ) - bridge._history_after_model = [] - client = _FakeEpisodeClient(truncated=True) - driver = harness._AppServerDriver( - Path("/unused"), bridge, client, "episode", "task", metrics - ) - driver._thread_id = "thread" - driver._turn_id = "turn" - with pytest.raises(harness.CodexHarnessError, match="mutated raw"): - await driver._handle_tool_request( - { - "id": 1, - "method": "item/tool/call", - "params": { - "threadId": "thread", - "turnId": "turn", - "callId": "call-1", - "namespace": None, - "tool": "terminal_exec", - "arguments": {"command": "whoami"}, - }, - } - ) - observation = await driver._execute_terminal({"command": "id"}) - assert observation == { - "exit_code": 0, - "output": "uid=1000", - "timed_out": False, - "truncated": True, - } - client.output = "x" * (harness.MAX_TOOL_OUTPUT_BYTES + 1) - with pytest.raises(harness.CodexHarnessError, match="signed boundary"): - await driver._execute_terminal({"command": "id"}) - - asyncio.run(scenario()) - - -def test_driver_prefers_policy_boundary_set_while_stdout_read_is_pending(monkeypatch): - _set_bridge_env(monkeypatch) - - async def scenario() -> None: - bridge = harness._ResponsesBridge( - "http://127.0.0.1:1", - "task", - {"max_tokens": 64}, - harness.legacy.AgentMetrics(), - max_seq_len=None, - ) - bridge._fatal = asyncio.get_running_loop().create_future() - boundary = harness.CodexModelFailure("sampled policy failure") - - class RaceStdout: - async def readline(self) -> bytes: - bridge._fatal.set_result(boundary) - return b"" - - class RaceProcess: - stdout = RaceStdout() - returncode = None - - driver = harness._AppServerDriver( - Path("/unused"), - bridge, - _FakeEpisodeClient(), - "episode", - "task", - harness.legacy.AgentMetrics(), - ) - driver._process = RaceProcess() # type: ignore[assignment] - with pytest.raises(harness.CodexModelFailure) as caught: - await driver._read_message() - assert caught.value is boundary - - asyncio.run(scenario()) - - -def _stock_codex_binary() -> Path: - # YETO_CODEX_BINARY_PATH (the fetched bundle's binary, see - # scripts/fetch_codex_bundle.py) wins over the PATH install. - binary_name = os.environ.get("YETO_CODEX_BINARY_PATH") or shutil.which("codex") - if not binary_name or not Path(binary_name).is_file(): - pytest.skip("stock Codex is not installed on this controller") - version = subprocess.run( - [binary_name, "--version"], capture_output=True, text=True, check=True - ).stdout.strip() - if version != harness.CODEX_CLI_VERSION: - pytest.skip(f"requires {harness.CODEX_CLI_VERSION}, found {version}") - return Path(binary_name) - - -def test_stock_codex_0145_app_server_fake_miles_integration(monkeypatch): - binary = _stock_codex_binary() - _set_bridge_env(monkeypatch, max_tokens=32_768) - - async def scenario() -> None: - long_reasoning = "r" * 100_000 - maximum_output = "x" * harness.MAX_TOOL_OUTPUT_BYTES - runner, miles_url, miles_requests = await _fake_miles( - [ - _completion( - "terminal.exec", - '{"command":"id"}', - "call-1", - long_reasoning, - 25_000, - ), - _completion( - "submit", - '{"evidence":"uid=1000"}', - "call-2", - "submit", - 60_000, - ), - ] - ) - metrics = harness.legacy.AgentMetrics() - client = _FakeEpisodeClient(output=maximum_output, truncated=True) - try: - async with harness._ResponsesBridge( - miles_url, - "solve target", - {"temperature": 0.7, "max_tokens": 32_768}, - metrics, - max_seq_len=200_000, - ) as bridge, harness._AppServerDriver( - binary, - bridge, - client, - "episode-stock", - "solve target", - metrics, - ) as driver: - status = await asyncio.wait_for(driver.drive(), timeout=20.0) - assert status == "completed" - finally: - await runner.cleanup() - - assert client.exec_calls == [("episode-stock", "id")] - assert client.submissions == [ - ("episode-stock", {"evidence": "uid=1000"}) - ] - assert len(miles_requests) == 2 - assert [tool["function"]["name"] for tool in miles_requests[0]["tools"]] == [ - "terminal.exec", - "submit", - ] - tool_message = miles_requests[1]["messages"][-1] - assert tool_message["role"] == "tool" - assert json.loads(tool_message["content"])["output"] == maximum_output - assert json.loads(tool_message["content"])["truncated"] is True - assert len(json.loads(tool_message["content"])["output"]) == ( - harness.MAX_TOOL_OUTPUT_BYTES - ) - - asyncio.run(scenario()) - - -def test_stock_codex_turn_limit_is_a_normal_policy_boundary(monkeypatch): - binary = _stock_codex_binary() - _set_bridge_env(monkeypatch) - monkeypatch.setenv("SECRLENV_MAX_TURNS", "1") - - async def scenario() -> None: - runner, miles_url, miles_requests = await _fake_miles( - [ - _completion( - "terminal.exec", '{"command":"id"}', "call-1", "inspect", 20 - ) - ] - ) - metrics = harness.legacy.AgentMetrics() - client = _FakeEpisodeClient() - try: - async with harness._ResponsesBridge( - miles_url, - "solve target", - {"max_tokens": 128}, - metrics, - max_seq_len=512, - ) as bridge, harness._AppServerDriver( - binary, - bridge, - client, - "episode-limit", - "solve target", - metrics, - ) as driver: - assert await asyncio.wait_for(driver.drive(), timeout=20.0) == "max_turns" - finally: - await runner.cleanup() - - assert client.exec_calls == [("episode-limit", "id")] - assert client.submissions == [] - assert len(miles_requests) == 1 - assert metrics.turns == 1 - - asyncio.run(scenario()) - - -def test_stock_codex_model_format_failure_retains_sample_as_normal_failure(monkeypatch): - binary = _stock_codex_binary() - _set_bridge_env(monkeypatch) - - async def scenario() -> None: - completion = { - "id": "chatcmpl-policy-failure", - "object": "chat.completion", - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "I will only describe what to do.", - "reasoning_content": "decline tool use", - "tool_calls": [], - }, - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5, - "total_tokens": 15, - }, - } - runner, miles_url, miles_requests = await _fake_miles([completion]) - metrics = harness.legacy.AgentMetrics() - client = _FakeEpisodeClient() - try: - async with harness._ResponsesBridge( - miles_url, - "solve target", - {"max_tokens": 128}, - metrics, - max_seq_len=512, - ) as bridge, harness._AppServerDriver( - binary, - bridge, - client, - "episode-policy-failure", - "solve target", - metrics, - ) as driver: - assert await asyncio.wait_for(driver.drive(), timeout=20.0) == "max_turns" - finally: - await runner.cleanup() - - assert client.exec_calls == [] - assert client.submissions == [] - assert len(miles_requests) == 1 - assert metrics.turns == 1 - assert metrics.parse_failures == 1 - - asyncio.run(scenario()) - - -def test_stock_codex_daemon_400_tool_errors_recover_without_dropping_sample(monkeypatch): - binary = _stock_codex_binary() - _set_bridge_env(monkeypatch) - - async def scenario() -> None: - runner, miles_url, miles_requests = await _fake_miles( - [ - _completion( - "terminal.exec", '{"command":"bad"}', "call-1", "try", 15 - ), - _completion( - "terminal.exec", '{"command":"id"}', "call-2", "recover", 30 - ), - _completion( - "submit", - '{"evidence":"too vague"}', - "call-3", - "submit", - 45, - ), - _completion( - "submit", - '{"evidence":"uid=1000"}', - "call-4", - "repair", - 60, - ), - ] - ) - metrics = harness.legacy.AgentMetrics() - client = _RecoveringEpisodeClient() - try: - async with harness._ResponsesBridge( - miles_url, - "solve target", - {"max_tokens": 128}, - metrics, - max_seq_len=1024, - ) as bridge, harness._AppServerDriver( - binary, - bridge, - client, - "episode-recovery", - "solve target", - metrics, - ) as driver: - assert await asyncio.wait_for(driver.drive(), timeout=20.0) == "completed" - finally: - await runner.cleanup() - - assert client.exec_calls == [ - ("episode-recovery", "bad"), - ("episode-recovery", "id"), - ] - assert client.submissions == [ - ("episode-recovery", {"evidence": "too vague"}), - ("episode-recovery", {"evidence": "uid=1000"}), - ] - assert len(miles_requests) == 4 - assert json.loads(miles_requests[1]["messages"][-1]["content"]) == { - "error": "command rejected" - } - assert json.loads(miles_requests[3]["messages"][-1]["content"]) == { - "error": "evidence rejected" - } - assert metrics.tool_calls == 4 - assert metrics.terminal_calls == 1 - assert metrics.submit_calls == 1 - - asyncio.run(scenario()) diff --git a/tests/test_tbench21_sao_diloco_plan.py b/tests/test_tbench21_sao_diloco_plan.py deleted file mode 100644 index 21aa2a8be..000000000 --- a/tests/test_tbench21_sao_diloco_plan.py +++ /dev/null @@ -1,234 +0,0 @@ -from __future__ import annotations - -import importlib.util -import json -from pathlib import Path - -import pytest - -from yeto.rl import CODEX_OPENENV_IDENTITY_ENV, SECRLENV_AGENTS - -TOOL = ( - Path(__file__).resolve().parents[1] - / "tools" - / "probes" - / "build_tbench21_sao_diloco_plan.py" -) -SPEC = importlib.util.spec_from_file_location("build_tbench21_sao_diloco_plan", TOOL) -assert SPEC is not None and SPEC.loader is not None -plan = importlib.util.module_from_spec(SPEC) -SPEC.loader.exec_module(plan) - - -def _task_tree(root: Path, count: int = 89) -> Path: - tasks = root / "tasks" - for index in range(count): - task = tasks / f"task-{index:03d}" - (task / "environment").mkdir(parents=True) - (task / "instruction.md").write_text("untrusted task text\n") - (task / "environment" / "Dockerfile").write_text("FROM scratch\n") - (task / "task.toml").write_text( - f"""schema_version = "1.1" -[task] -name = "terminal-bench/task-{index:03d}" -[metadata] -category = "category-{index % 7}" -[environment] -docker_image = "example/task-{index:03d}:fixed" -cpus = 1 -memory_mb = 2048 -storage_mb = 10240 -""" - ) - return tasks - - -def _jsonl(path: Path) -> list[dict]: - return [json.loads(line) for line in path.read_text().splitlines() if line] - - -def test_builds_exact_full_suite_plan(tmp_path): - tasks = _task_tree(tmp_path / "source") - output = tmp_path / "plan" - - manifest_path = plan.build_plan( - tasks, - output, - seed="fixed-seed", - run_root="/run/tbench21", - ) - manifest = json.loads(manifest_path.read_text()) - - assert manifest["schema"] == plan.SCHEMA - assert manifest["split"]["train_task_count"] == 44 - assert manifest["split"]["eval_task_count"] == 45 - train_ids = set(manifest["split"]["train_task_ids"]) - eval_ids = set(manifest["split"]["eval_task_ids"]) - assert not train_ids & eval_ids - assert len(train_ids | eval_ids) == 89 - assert manifest["rollouts"] == { - "per_task": 4, - "baseline": 356, - "train": 176, - "eval": 180, - "episode_timeout_seconds": 1800, - "target_concurrency": 304, - "per_island_concurrency": 38, - "seed_base": 82621, - "per_island_seeds": list(range(82621, 82629)), - "baseline_achievable_concurrency": 304, - "train_cardinality_limited_concurrency": 176, - "eval_cardinality_limited_concurrency": 180, - } - - baseline = _jsonl(output / "baseline" / "all.jsonl") - train = _jsonl(output / "train" / "all.jsonl") - evaluation = _jsonl(output / "eval" / "all.jsonl") - assert (len(baseline), len(train), len(evaluation)) == (356, 176, 180) - assert {row["metadata"]["task_id"] for row in train} == train_ids - assert {row["metadata"]["task_id"] for row in evaluation} == eval_ids - assert all(row["metadata"]["episode_timeout_seconds"] == 1800 for row in baseline) - assert all( - row["metadata"]["rollout_seed"] - == 82621 + row["metadata"]["island_id"] - for row in baseline + train + evaluation - ) - assert all( - len([row for row in baseline if row["metadata"]["task_id"] == task_id]) == 4 - for task_id in train_ids | eval_ids - ) - - assert [len(_jsonl(output / "train" / f"island-{i}.jsonl")) for i in range(8)] == [ - 22 - ] * 8 - baseline_counts = [ - len(_jsonl(output / "baseline" / f"island-{i}.jsonl")) for i in range(8) - ] - assert sum(baseline_counts) == 356 - assert max(baseline_counts) - min(baseline_counts) == 1 - - value = json.loads((output / "value-pretraining-selection.json").read_text()) - assert value["selection"] == "all" - assert value["sample_count"] == 356 - assert value["actor_eval_task_count_exposed_to_critic"] == 45 - assert value["actor_eval_trajectory_count_exposed_to_critic"] == 180 - assert value["critic_pretraining_exposes_actor_eval_tasks"] is True - assert manifest["compaction"] == { - "enabled": True, - "trainer_objective": "sao", - "max_seq_len": 8192, - "trigger_semantics": "consumed-context-tokens-at-least", - "trigger_tokens": 6144, - "summary_max_tokens": 1024, - "max_compactions_per_episode": 3, - } - - topology = manifest["topology"] - assert topology["islands"] == topology["physical_gpus"] == 8 - assert topology["one_physical_gpu_per_island"] is True - assert topology["cross_island_collectives"] is False - for island in topology["island_contracts"]: - assert island["rollout_seed"] == 82621 + island["island_id"] - assert island["entrypoint"].endswith("train_sao_streaming_secrlenv.py") - assert island["runtime_contracts"]["actor_syncer_port"] == 29400 - assert island["runtime_contracts"]["critic_syncer_port"] == 29401 - assert island["openenv_agent_contract"] == { - "custom_agent_function_path": ( - "yeto.rl.harness.codex." - "codex_openenv_subprocess_agent_function.run" - ), - "custom_rm_path": "openenv_generate.reward_func", - "dynamic_sampling_filter_path": ( - "openenv_generate.check_terminal_bench_episode" - ), - "input_key": "messages", - "tito_model": "qwen35", - "codex_backend_profile": "qwen35_08b", - "model": "Qwen/Qwen3.5-0.8B", - "model_revision": "2fc06364715b967f1860aea9cf38778875588b17", - "max_seq_len": 8192, - "must_not_enter_secrlenv_retry_path": True, - } - assert ( - island["openenv_agent_contract"]["custom_agent_function_path"] - not in SECRLENV_AGENTS - ) - assert island["ray_fractional_gpu"] == { - "actor": 0.4, - "critic": 0.4, - "sglang": 0.2, - } - args = island["miles_topology_args"] - assert "--sao-one-gpu-island" in args - assert "--sao-compaction" in args - assert args[args.index("--input-key") + 1] == "messages" - assert "--colocate" in args - assert args[args.index("--sglang-mem-fraction-static") + 1] == "0.15" - assert args[args.index("--sglang-max-total-tokens") + 1] == "393216" - assert args[args.index("--sglang-max-mamba-cache-size") + 1] == "256" - environment = island["environment"] - assert environment["SECRLENV_MAX_ROLLOUT_TIME_SECONDS"] == "1800" - assert environment["OPENENV_MAX_ROLLOUT_TIME_SECONDS"] == "1800" - assert environment["YETO_CODEX_COMPACTION_ENABLED"] == "1" - assert environment["YETO_CODEX_COMPACTION_TRIGGER_TOKENS"] == "6144" - assert environment["YETO_CODEX_COMPACTION_SUMMARY_MAX_TOKENS"] == "1024" - assert environment["YETO_CODEX_MAX_COMPACTIONS"] == "3" - assert all( - environment[name] == value - for name, value in CODEX_OPENENV_IDENTITY_ENV.items() - ) - assert "PYTORCH_CUDA_ALLOC_CONF" not in environment - assert ( - manifest["required_gates"]["openenv_tbench_max_active_episodes_at_least"] == 304 - ) - assert ( - manifest["required_gates"]["openenv_compaction_trajectory_evidence_v2"] is True - ) - assert "terminal_bench_signed_outcome_hmac_key_file" not in manifest["required_gates"] - assert manifest["required_gates"]["direct_codex_qwen35_08b_preflight"] is True - assert manifest["launch_readiness"]["ready"] is False - assert [ - issue["code"] for issue in manifest["launch_readiness"]["blocking_issues"] - ] == [ - "task-images-not-attested", - "openenv-capacity-not-attested", - "one-gpu-island-runtime-not-smoked", - "final-layout-attestations-not-generated", - ] - assert manifest["integration_notes"] == [ - { - "code": "streaming-entrypoint-retains-legacy-cli-name", - "blocking": False, - "detail": manifest["integration_notes"][0]["detail"], - } - ] - - -def test_split_is_reproducible(tmp_path): - tasks = _task_tree(tmp_path / "source") - first = plan.build_plan( - tasks, - tmp_path / "first", - seed="fixed-seed", - run_root="/run/one", - ) - second = plan.build_plan( - tasks, - tmp_path / "second", - seed="fixed-seed", - run_root="/run/two", - ) - one = json.loads(first.read_text())["split"] - two = json.loads(second.read_text())["split"] - assert one == two - - -def test_rejects_incomplete_suite(tmp_path): - tasks = _task_tree(tmp_path / "source", count=88) - with pytest.raises(plan.PlanError, match="exactly 89"): - plan.build_plan( - tasks, - tmp_path / "output", - seed="fixed-seed", - run_root="/run/tbench21", - ) diff --git a/tools/build_eval_data.py b/tools/build_eval_data.py deleted file mode 100644 index f88270dfe..000000000 --- a/tools/build_eval_data.py +++ /dev/null @@ -1,124 +0,0 @@ -#!/usr/bin/env python3 -"""TB2 eval / train jsonl from the hold-out list (rl-eval-difficulty-buckets 1.1 second half, 1.3). - - python tools/build_eval_data.py --tb2-tasks-dir ~/work/tb2-data \ - --holdout data/eval/tb2-holdout.json --out-dir data - -Writes -* ``/eval/tb2-holdout-eval.jsonl`` -- the 30 hold-out tasks (eval only); -* ``/tb2/tb2-train.jsonl`` -- every other TB2 task except the unusable table - ``data/eval/tb2-unusable.json`` (S17 N13; 46 at tb2@2fd12b8 = 89 - 30 hold-out - 13 - unusable; the S15 smoke-6 included: they were excluded from the eval pool, not from training); -* ``/eval/tb2-data.sha256.json`` -- sha256 of the hold-out list and both files - (pin these in the run config). - -Row format = design D6.a: the smoke-6 system prompt (the task statement itself -comes from the task's ``instruction.md`` at rollout time, ``task_prompt``) and -``metadata`` with ``task_id``, ``benchmark``, ``benchmark_version``, -``difficulty``, ``difficulty_source`` and, on eval rows, ``eval_bucket``. -Deterministic: rows sorted by ``task_id``; the same inputs give the same bytes. -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import sys -from pathlib import Path -from typing import Any - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -SYSTEM_PROMPT = ( - "You are an autonomous terminal agent solving a Terminal-Bench task. You will be given the task " - "instruction, then interact with a real Linux shell. On each turn respond with EXACTLY ONE shell " - "command inside a single ```bash code block and nothing else. Inspect the environment, make the " - "required changes, and verify your work. When you are confident the task is fully complete, reply " - "with TASK_COMPLETE (with no code block)." -) -EVAL_JSONL = "eval/tb2-holdout-eval.jsonl" -TRAIN_JSONL = "tb2/tb2-train.jsonl" -SHA_JSON = "eval/tb2-data.sha256.json" - - -def _row(spec: Any, *, with_bucket: bool) -> dict[str, Any]: - meta = {"task_id": spec.task_id, "benchmark": spec.benchmark, "benchmark_version": spec.benchmark_version, - "difficulty": spec.difficulty, "difficulty_source": spec.difficulty_source} - if with_bucket: - if not spec.eval_bucket: - raise ValueError(f"{spec.task_id}: hold-out task without an eval bucket") - meta["eval_bucket"] = spec.eval_bucket - return {"prompt": [{"role": "system", "content": SYSTEM_PROMPT}], "metadata": meta} - - -def _jsonl(rows: list[dict[str, Any]]) -> bytes: - return "".join(json.dumps(r, sort_keys=True, ensure_ascii=False) + "\n" for r in rows).encode() - - -def build(tasks_dir: str | Path, holdout: dict[str, Any], *, - unusable: dict[str, str] | None = None) -> dict[str, bytes]: - """{relative path: bytes} for the eval jsonl, the train jsonl and the sha256 record.""" - from yeto.rl.harness.reward_env import benchmark as bm - from yeto.rl.harness.reward_env import tb2 - - adapter = tb2.Tb2Benchmark(tasks_dir) - if holdout.get("benchmark_version") != adapter.version: - raise ValueError(f"hold-out list is for {holdout.get('benchmark_version')}, checkout is {adapter.version}") - held = bm.holdout_ids(holdout) - every = adapter.task_ids() - missing = sorted(set(held) - set(every)) - if missing: - raise ValueError(f"hold-out tasks not in the checkout: {missing}") - by_item = {i["task_id"]: i for i in holdout["items"]} - eval_rows = [] - for tid in sorted(held): - spec = adapter.task_spec(tid) - if spec.eval_bucket != by_item[tid]["eval_bucket"]: - raise ValueError(f"{tid}: bucket {spec.eval_bucket} != hold-out list {by_item[tid]['eval_bucket']}") - eval_rows.append(_row(spec, with_bucket=True)) - unusable = {} if unusable is None else unusable - unknown = sorted(set(unusable) - set(every)) - if unknown: - raise ValueError(f"unusable tasks not in the checkout: {unknown}") - clash = sorted(set(unusable) & set(held)) - if clash: - raise ValueError(f"unusable tasks in the hold-out list (rebuild it with tools/reward_env/holdout.py): {clash}") - train_rows = [_row(adapter.task_spec(t), with_bucket=False) for t in every - if t not in set(held) and t not in unusable] - bm.assert_disjoint([r["metadata"]["task_id"] for r in train_rows], held) - eval_bytes, train_bytes = _jsonl(eval_rows), _jsonl(train_rows) - record = { - "benchmark_version": adapter.version, - "holdout_sha256": bm.holdout_sha256(holdout), - "unusable_excluded": sorted(unusable), - EVAL_JSONL: {"rows": len(eval_rows), "sha256": hashlib.sha256(eval_bytes).hexdigest()}, - TRAIN_JSONL: {"rows": len(train_rows), "sha256": hashlib.sha256(train_bytes).hexdigest()}, - } - return {EVAL_JSONL: eval_bytes, TRAIN_JSONL: train_bytes, - SHA_JSON: (json.dumps(record, indent=1, sort_keys=True) + "\n").encode()} - - -def main(argv: list[str] | None = None) -> int: - p = argparse.ArgumentParser() - p.add_argument("--tb2-tasks-dir", required=True) - p.add_argument("--holdout", required=True) - p.add_argument("--out-dir", required=True) - p.add_argument("--tb2-unusable", default=None, - help="unusable-task table (default: data/eval/tb2-unusable.json); its tasks leave training") - p.add_argument("--tb2-no-unusable", action="store_true", help="do not apply the unusable-task table") - a = p.parse_args(argv) - from yeto.rl.harness.reward_env import tb2 - - unusable = {} if a.tb2_no_unusable else tb2.load_unusable(a.tb2_unusable) - out = Path(a.out_dir) - for rel, data in build(a.tb2_tasks_dir, json.loads(Path(a.holdout).read_text()), - unusable=unusable).items(): - (out / rel).parent.mkdir(parents=True, exist_ok=True) - (out / rel).write_bytes(data) - print(rel, len(data), hashlib.sha256(data).hexdigest()[:12]) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tools/probes/build_tbench21_sao_diloco_plan.py b/tools/probes/build_tbench21_sao_diloco_plan.py deleted file mode 100755 index fd9d7759a..000000000 --- a/tools/probes/build_tbench21_sao_diloco_plan.py +++ /dev/null @@ -1,520 +0,0 @@ -#!/usr/bin/env python3 -"""Build the immutable Terminal-Bench 2.1 SAO/DiLoCo run inputs. - -This tool only prepares data and a launch contract. It never starts Docker, -Ray, a model server, or a training process. The resulting plan intentionally -records that value pretraining consumes all four baseline trajectories from -all tasks, including the actor-evaluation half, while actor RL sees only the -44-task training split. -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import os -import subprocess -from collections import Counter -from collections.abc import Iterable -from pathlib import Path -from typing import Any - -import tomllib - -from yeto.rl import CODEX_OPENENV_AGENT, CODEX_OPENENV_IDENTITY_ENV -from yeto.rl.codex_backend import QWEN35_08B_MODEL, QWEN35_08B_REVISION - -SCHEMA = "yeto.tbench21-sao-diloco-plan.v1" -SPLIT_DOMAIN = b"yeto-tbench21-sao-split-v1\0" -ROLLOUTS_PER_TASK = 4 -EPISODE_TIMEOUT_SECONDS = 1800 -ISLANDS = 8 -TARGET_CONCURRENCY = 304 -PER_ISLAND_CONCURRENCY = 38 -SGLANG_MEM_FRACTION_STATIC = 0.15 -SGLANG_MAX_TOTAL_TOKENS = 393216 -SGLANG_MAX_MAMBA_CACHE_SIZE = 256 -MAX_SEQ_LEN = 8192 -COMPACTION_TRIGGER_TOKENS = 6144 -COMPACTION_SUMMARY_MAX_TOKENS = 1024 -MAX_COMPACTIONS = 3 -ROLLOUT_SEED_BASE = 82621 -MIN_H200_MEMORY_BYTES = 120 * 1024**3 -SYSTEM_MESSAGE = ( - "The authenticated Terminal-Bench daemon supplies the immutable task " - "instruction. Use the Codex terminal tools to solve it and submit exactly once." -) - - -class PlanError(RuntimeError): - pass - - -def _sha256_bytes(value: bytes) -> str: - return hashlib.sha256(value).hexdigest() - - -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - for block in iter(lambda: handle.read(1024 * 1024), b""): - digest.update(block) - return digest.hexdigest() - - -def _canonical_bytes(value: object) -> bytes: - return json.dumps( - value, sort_keys=True, separators=(",", ":"), ensure_ascii=False - ).encode("utf-8") - - -def _write_json(path: Path, value: object) -> None: - path.write_bytes(_canonical_bytes(value) + b"\n") - - -def _write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> None: - with path.open("wb") as handle: - for row in rows: - handle.write(_canonical_bytes(row) + b"\n") - - -def _git_revision(root: Path) -> str | None: - try: - value = subprocess.run( - [ - "git", - "-c", - f"safe.directory={root}", - "-C", - str(root), - "rev-parse", - "HEAD", - ], - check=True, - stdout=subprocess.PIPE, - stderr=subprocess.DEVNULL, - text=True, - ).stdout.strip() - except (OSError, subprocess.CalledProcessError): - return None - return value if len(value) == 40 else None - - -def _load_tasks(tasks_dir: Path) -> list[dict[str, Any]]: - if not tasks_dir.is_dir() or tasks_dir.is_symlink(): - raise PlanError("tasks directory must be a real directory") - tasks: list[dict[str, Any]] = [] - for task_dir in sorted(tasks_dir.iterdir(), key=lambda path: path.name): - if not task_dir.is_dir() or task_dir.is_symlink(): - continue - task_toml = task_dir / "task.toml" - instruction = task_dir / "instruction.md" - dockerfile = task_dir / "environment" / "Dockerfile" - if ( - not task_toml.is_file() - or not instruction.is_file() - or not dockerfile.is_file() - ): - raise PlanError(f"task {task_dir.name!r} is incomplete") - try: - raw = tomllib.loads(task_toml.read_text(encoding="utf-8")) - except (OSError, UnicodeError, tomllib.TOMLDecodeError) as error: - raise PlanError(f"task {task_dir.name!r} has invalid task.toml") from error - metadata = raw.get("metadata", {}) - environment = raw.get("environment", {}) - task = raw.get("task", {}) - if not all(isinstance(value, dict) for value in (metadata, environment, task)): - raise PlanError(f"task {task_dir.name!r} metadata is malformed") - docker_image = environment.get("docker_image") - if not isinstance(docker_image, str) or not docker_image: - raise PlanError(f"task {task_dir.name!r} has no Docker image identity") - category = metadata.get("category", "unknown") - if not isinstance(category, str) or not category: - raise PlanError(f"task {task_dir.name!r} has no category") - tasks.append( - { - "task_id": task_dir.name, - "task_name": task.get("name"), - "category": category, - "docker_image": docker_image, - "cpus": environment.get("cpus"), - "memory_mb": environment.get("memory_mb"), - "storage_mb": environment.get("storage_mb"), - "task_toml_sha256": _sha256_file(task_toml), - "instruction_sha256": _sha256_file(instruction), - "dockerfile_sha256": _sha256_file(dockerfile), - } - ) - if len(tasks) != 89: - raise PlanError( - f"Terminal-Bench 2.1 must contain exactly 89 tasks, got {len(tasks)}" - ) - ids = [task["task_id"] for task in tasks] - if len(set(ids)) != len(ids): - raise PlanError("Terminal-Bench task IDs are not unique") - return tasks - - -def _split_tasks( - tasks: list[dict[str, Any]], seed: str -) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - """Return an exact deterministic 44-train/45-eval partition.""" - - if not seed: - raise PlanError("split seed must not be empty") - - def key(task: dict[str, Any]) -> tuple[bytes, str]: - payload = ( - SPLIT_DOMAIN - + seed.encode("utf-8") - + b"\0" - + task["task_id"].encode("utf-8") - ) - return hashlib.sha256(payload).digest(), task["task_id"] - - ranked = sorted(tasks, key=key) - return ranked[:44], ranked[44:] - - -def _rollout_rows(tasks: list[dict[str, Any]], *, phase: str) -> list[dict[str, Any]]: - rows: list[dict[str, Any]] = [] - for task_index, task in enumerate(tasks): - for replica in range(ROLLOUTS_PER_TASK): - island_id = (task_index * ROLLOUTS_PER_TASK + replica) % ISLANDS - rollout_seed = ROLLOUT_SEED_BASE + island_id - sample_id = f"{phase}:{task['task_id']}:r{replica}" - rows.append( - { - "messages": [{"role": "system", "content": SYSTEM_MESSAGE}], - "metadata": { - "task_id": task["task_id"], - "prompt_tier": "l2", - "max_seq_len": MAX_SEQ_LEN, - "split": phase, - "rollout_replica": replica, - "sample_id": sample_id, - "island_id": island_id, - "rollout_seed": rollout_seed, - "episode_timeout_seconds": EPISODE_TIMEOUT_SECONDS, - }, - } - ) - return rows - - -def _island_contract(island_id: int, run_root: str) -> dict[str, Any]: - base = 32000 + island_id * 100 - return { - "island_id": island_id, - "physical_gpu_index": island_id, - "cuda_visible_devices": str(island_id), - "rollout_seed": ROLLOUT_SEED_BASE + island_id, - "run_dir": f"{run_root}/island-{island_id}", - "entrypoint": "/root/miles/tools/probes/train_sao_streaming_secrlenv.py", - "runtime_contracts": { - "sao_context": f"{run_root}/island-{island_id}/sao-context.json", - "streaming_context": f"{run_root}/island-{island_id}/sao-streaming-context.json", - "actor_syncer_port": 29400, - "critic_syncer_port": 29401, - }, - "openenv_agent_contract": { - "custom_agent_function_path": CODEX_OPENENV_AGENT, - "custom_rm_path": "openenv_generate.reward_func", - "dynamic_sampling_filter_path": ( - "openenv_generate.check_terminal_bench_episode" - ), - "input_key": "messages", - "tito_model": "qwen35", - "codex_backend_profile": "qwen35_08b", - "model": QWEN35_08B_MODEL, - "model_revision": QWEN35_08B_REVISION, - "max_seq_len": MAX_SEQ_LEN, - "must_not_enter_secrlenv_retry_path": True, - }, - "ray_gcs_port": 26000 + island_id, - "ports": { - "session_server": 31000 + island_id, - "rollout_engine_base": base, - "sglang_router": base + 20, - "sglang_router_prometheus": base + 21, - "train_master_base": base + 40, - }, - "environment": { - "SECRLENV_MAX_ROLLOUT_TIME_SECONDS": str(EPISODE_TIMEOUT_SECONDS), - "OPENENV_MAX_ROLLOUT_TIME_SECONDS": str(EPISODE_TIMEOUT_SECONDS), - "YETO_CODEX_COMPACTION_ENABLED": "1", - "YETO_CODEX_COMPACTION_TRIGGER_TOKENS": str(COMPACTION_TRIGGER_TOKENS), - "YETO_CODEX_COMPACTION_SUMMARY_MAX_TOKENS": str( - COMPACTION_SUMMARY_MAX_TOKENS - ), - "YETO_CODEX_MAX_COMPACTIONS": str(MAX_COMPACTIONS), - **CODEX_OPENENV_IDENTITY_ENV, - "NCCL_NVLS_ENABLE": "0", - "RAY_DEDUP_LOGS": "0", - }, - "miles_topology_args": [ - "--sao-one-gpu-island", - "--sao-compaction", - "--input-key", - "messages", - "--colocate", - "--num-gpus-per-node", - "1", - "--actor-num-nodes", - "1", - "--actor-num-gpus-per-node", - "1", - "--critic-num-nodes", - "1", - "--critic-num-gpus-per-node", - "1", - "--rollout-num-gpus", - "1", - "--rollout-num-gpus-per-engine", - "1", - "--tensor-model-parallel-size", - "1", - "--pipeline-model-parallel-size", - "1", - "--context-parallel-size", - "1", - "--expert-model-parallel-size", - "1", - "--sglang-server-concurrency", - str(PER_ISLAND_CONCURRENCY), - "--sglang-max-running-requests", - str(PER_ISLAND_CONCURRENCY), - "--async-max-concurrent-samples", - str(PER_ISLAND_CONCURRENCY), - "--sglang-mem-fraction-static", - str(SGLANG_MEM_FRACTION_STATIC), - "--sglang-max-total-tokens", - str(SGLANG_MAX_TOTAL_TOKENS), - "--sglang-max-mamba-cache-size", - str(SGLANG_MAX_MAMBA_CACHE_SIZE), - ], - "ray_fractional_gpu": {"actor": 0.4, "critic": 0.4, "sglang": 0.2}, - } - - -def _category_counts(tasks: list[dict[str, Any]]) -> dict[str, int]: - return dict(sorted(Counter(task["category"] for task in tasks).items())) - - -def build_plan(tasks_dir: Path, output_dir: Path, *, seed: str, run_root: str) -> Path: - if output_dir.exists() or output_dir.is_symlink(): - raise PlanError("output directory must be fresh") - if not os.path.isabs(run_root) or ".." in Path(run_root).parts: - raise PlanError("run root must be an absolute safe path") - - tasks = _load_tasks(tasks_dir) - train_tasks, eval_tasks = _split_tasks(tasks, seed) - baseline_tasks = sorted(tasks, key=lambda task: task["task_id"]) - rows_by_phase = { - "baseline": _rollout_rows(baseline_tasks, phase="baseline"), - "train": _rollout_rows(train_tasks, phase="train"), - "eval": _rollout_rows(eval_tasks, phase="eval"), - } - if [len(rows_by_phase[name]) for name in ("baseline", "train", "eval")] != [ - 356, - 176, - 180, - ]: - raise AssertionError("Terminal-Bench rollout cardinality changed") - - output_dir.mkdir(mode=0o755, parents=True) - files: dict[str, str] = {} - for phase, rows in rows_by_phase.items(): - phase_dir = output_dir / phase - phase_dir.mkdir() - all_path = phase_dir / "all.jsonl" - _write_jsonl(all_path, rows) - files[str(all_path.relative_to(output_dir))] = _sha256_file(all_path) - for island_id in range(ISLANDS): - island_path = phase_dir / f"island-{island_id}.jsonl" - _write_jsonl( - island_path, - (row for row in rows if row["metadata"]["island_id"] == island_id), - ) - files[str(island_path.relative_to(output_dir))] = _sha256_file(island_path) - - baseline_ids = [row["metadata"]["sample_id"] for row in rows_by_phase["baseline"]] - value_selection = { - "schema": "yeto.sao-value-pretraining-selection.v1", - "source_phase": "baseline", - "selection": "all", - "sample_count": len(baseline_ids), - "task_count": len(tasks), - "rollouts_per_task": ROLLOUTS_PER_TASK, - "baseline_sample_ids": baseline_ids, - "critic_pretraining_exposes_actor_eval_tasks": True, - "actor_eval_task_count_exposed_to_critic": len(eval_tasks), - "actor_eval_trajectory_count_exposed_to_critic": len(eval_tasks) - * ROLLOUTS_PER_TASK, - } - value_path = output_dir / "value-pretraining-selection.json" - _write_json(value_path, value_selection) - files[value_path.name] = _sha256_file(value_path) - - task_contracts = { - task["task_id"]: {key: value for key, value in task.items() if key != "task_id"} - for task in sorted(tasks, key=lambda item: item["task_id"]) - } - inventory_sha256 = _sha256_bytes(_canonical_bytes(task_contracts)) - split_payload = { - "seed": seed, - "algorithm": "sha256-domain-ranked-first-44", - "train_task_ids": [task["task_id"] for task in train_tasks], - "eval_task_ids": [task["task_id"] for task in eval_tasks], - } - split_sha256 = _sha256_bytes(_canonical_bytes(split_payload)) - - manifest = { - "schema": SCHEMA, - "terminal_bench": { - "version": "2.1", - "revision": _git_revision(tasks_dir.parent), - "task_count": len(tasks), - "task_inventory_sha256": inventory_sha256, - "task_contracts": task_contracts, - }, - "split": { - **split_payload, - "sha256": split_sha256, - "train_task_count": len(train_tasks), - "eval_task_count": len(eval_tasks), - "train_category_counts": _category_counts(train_tasks), - "eval_category_counts": _category_counts(eval_tasks), - }, - "rollouts": { - "per_task": ROLLOUTS_PER_TASK, - "baseline": 356, - "train": 176, - "eval": 180, - "episode_timeout_seconds": EPISODE_TIMEOUT_SECONDS, - "target_concurrency": TARGET_CONCURRENCY, - "per_island_concurrency": PER_ISLAND_CONCURRENCY, - "seed_base": ROLLOUT_SEED_BASE, - "per_island_seeds": [ - ROLLOUT_SEED_BASE + island_id for island_id in range(ISLANDS) - ], - "baseline_achievable_concurrency": TARGET_CONCURRENCY, - "train_cardinality_limited_concurrency": 176, - "eval_cardinality_limited_concurrency": 180, - }, - "value_pretraining": { - "selection_path": value_path.name, - "selection_sha256": files[value_path.name], - "uses_all_baseline_trajectories": True, - "sample_count": 356, - "critic_pretraining_exposes_actor_eval_tasks": True, - }, - "compaction": { - "enabled": True, - "trainer_objective": "sao", - "max_seq_len": MAX_SEQ_LEN, - "trigger_semantics": "consumed-context-tokens-at-least", - "trigger_tokens": COMPACTION_TRIGGER_TOKENS, - "summary_max_tokens": COMPACTION_SUMMARY_MAX_TOKENS, - "max_compactions_per_episode": MAX_COMPACTIONS, - }, - "topology": { - "islands": ISLANDS, - "physical_gpus": ISLANDS, - "one_physical_gpu_per_island": True, - "model": QWEN35_08B_MODEL, - "model_revision": QWEN35_08B_REVISION, - "minimum_gpu_memory_bytes": MIN_H200_MEMORY_BYTES, - "cross_island_collectives": False, - "synchronization": "dual-role streaming DiLoCo", - "syncer_sessions": {"actor_port": 29400, "critic_port": 29401}, - "island_contracts": [ - _island_contract(island_id, run_root) for island_id in range(ISLANDS) - ], - }, - "required_gates": { - "codex_harness": True, - "compaction": True, - "all_task_images_available": True, - "openenv_tbench_max_active_episodes_at_least": TARGET_CONCURRENCY, - "openenv_compaction_trajectory_evidence_v2": True, - "direct_codex_qwen35_08b_preflight": True, - "host_open_file_limit_at_least": 65536, - "fresh_per_island_ray_clusters": True, - "actor_and_critic_streaming_syncers_healthy": True, - "layout_attestation_per_role": True, - }, - "launch_readiness": { - "ready": False, - "blocking_issues": [ - { - "code": "task-images-not-attested", - "detail": ( - "All 89 immutable Terminal-Bench task images must be present " - "and capacity-preflighted before the baseline or online phase." - ), - }, - { - "code": "openenv-capacity-not-attested", - "detail": ( - "The trusted OpenEnv service has not yet attested at least 304 " - "simultaneous Terminal-Bench episodes on this host." - ), - }, - { - "code": "one-gpu-island-runtime-not-smoked", - "detail": ( - "Actor, critic, and SGLang fractional placement is validated " - "fail-closed but still needs a physical one-GPU memory/runtime " - "smoke before the eight-island launch." - ), - }, - { - "code": "final-layout-attestations-not-generated", - "detail": ( - "Each island still needs its final actor/critic layout probe " - "and hash-bound streaming runtime after fragment count P is " - "discovered." - ), - }, - ], - }, - "integration_notes": [ - { - "code": "streaming-entrypoint-retains-legacy-cli-name", - "blocking": False, - "detail": ( - "The retained entrypoint and flag names mention SecRLEnv for " - "compatibility, but the v2 base context binds the benchmark as " - "terminal-bench-2.1 without synthetic SecRLEnv fields. Direct " - "preflight attests the exact Codex/OpenEnv Qwen3.5-0.8B adapter " - "before Miles parses arguments or allocates a GPU." - ), - } - ], - "files": dict(sorted(files.items())), - } - manifest_path = output_dir / "manifest.json" - _write_json(manifest_path, manifest) - return manifest_path - - -def main(argv: list[str] | None = None) -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--tasks-dir", type=Path, required=True) - parser.add_argument("--output-dir", type=Path, required=True) - parser.add_argument("--split-seed", default="tbench21-sao-20260826") - parser.add_argument("--run-root", default="/data/sft/tbench21-sao-diloco-full") - args = parser.parse_args(argv) - manifest = build_plan( - args.tasks_dir.resolve(), - args.output_dir.resolve(), - seed=args.split_seed, - run_root=args.run_root, - ) - print(manifest) - - -if __name__ == "__main__": - main() diff --git a/tools/probes/s17_eval_island_gpu_smoke.py b/tools/probes/s17_eval_island_gpu_smoke.py deleted file mode 100644 index 953b662ae..000000000 --- a/tools/probes/s17_eval_island_gpu_smoke.py +++ /dev/null @@ -1,85 +0,0 @@ -"""(codex-env.json in OUT: yeto.launcher.codex_container_env(codex_bundle_contract(args, /codex)) for qwen35_08b / 4096 / xhigh.) -S17 N11 eval-island GPU smoke: real model (Qwen3.5-0.8B, codex qwen35_08b) on TB2 hold-out 3 tasks. - -seed (Modal CPU, same image): version 0 into Volume through the training export path -island (Modal H100!:1): SGLang + Miles session server + codex worker -> TB2 sandboxes -> judge. -Run with reward-env-venv python (modal 1.6.1). Raw output stays in this directory. -""" -import json, os, sys, time, tomllib, threading -from pathlib import Path - -REPO = Path("/home/michael/work/s17-eval-infer") -sys.path.insert(0, str(REPO)) -OUT = Path(os.environ.get("N11_OUT", "s1-runs/s17-n11-evalinfer")).resolve(); OUT.mkdir(parents=True, exist_ok=True) -APP = "yeto-s17-n11-evalinfer" -VOLUME = os.environ.get("N11_VOLUME", "yeto-eval-store-n11") -TASKS = ["prove-plus-comm", "constraints-scheduling", "largest-eigenval"] -MODEL, REV = "Qwen/Qwen3.5-0.8B", "2fc06364715b967f1860aea9cf38778875588b17" -HARD_S = int(os.environ.get("N11_HARD_S", "2400")) -SAMPLING = {"rollout_temperature": 1.0, "rollout_top_p": 1.0, "rollout_top_k": None, - "rollout_max_response_len": 4096, "rollout_max_context_len": 8192, - "yeto_codex_reasoning_effort": "xhigh", "yeto_codex_max_turns": "12"} - -import modal -from yeto.cloud import modal_eval_island as mei - -log = open(OUT / "smoke.log", "a") -def say(*a): - line = " ".join(str(x) for x in a) - print(time.strftime("%H:%M:%SZ", time.gmtime()), line, file=log, flush=True) - print(line, flush=True) - -prof = tomllib.loads(Path("~/.modal.toml").expanduser().read_text()) -prof = next(v for v in prof.values() if isinstance(v, dict) and v.get("active")) if any( - isinstance(v, dict) and v.get("active") for v in prof.values()) else next(iter(prof.values())) -envs = json.loads((OUT / "codex-env.json").read_text()) -envs.update({"SECRLENV_MAX_TURNS": "12", "YETO_HARNESS_TB2_LEASE_SECONDS": "600", - "YETO_CODEX_WORKER_STDERR": "1", - "MODAL_TOKEN_ID": prof["token_id"], "MODAL_TOKEN_SECRET": prof["token_secret"]}) -spec = mei.EvalFunctionSpec(app_name=APP, volume=VOLUME, workdir=str(REPO), - codex_dir="/home/michael/work/codex-bundle/codex", tb2_dir="/home/michael/work/tb2-data", - gpu="H100!", gpus=1, cpu=8.0, memory_mib=65536, timeout_s=HARD_S, envs=envs) -app, island_fn, seed_fn = mei.build_eval_app(spec) - -config = { - "store": mei.EVAL_STORE_MOUNT, "volume": VOLUME, "island_id": "eval-n11", - "holdout": f"{mei.WORKDIR_MOUNT}/data/eval/tb2-holdout.json", - "eval_data": f"{mei.WORKDIR_MOUNT}/data/eval/tb2-holdout-eval.jsonl", - "task_ids": TASKS, "trials_v0": int(os.environ.get("N11_TRIALS", "1")), "trials": 1, "assert_gpu_name": "H100", - "loader": "yeto.rl.eval.codex_infer:loader_factory", "attempt": "yeto.rl.eval.codex_infer:attempt_factory", - "infer": {"base_model": MODEL, "revision": REV, "tito_model": "qwen35", "context_length": 8192, - "max_tokens": 4096, "mem_fraction_static": 0.8}, -} -events = [] -def emit(event, **f): - events.append({"event": event, "ts": time.time(), **f}) - with open(OUT / "launcher-events.jsonl", "a") as fh: - fh.write(json.dumps(events[-1], default=str) + "\n") - -def watchdog(): - time.sleep(HARD_S + 300) - say("WATCHDOG: hard limit passed, stopping app") - os.system(f"/home/michael/work/reward-env-venv/bin/modal app stop {APP} >> {OUT}/watchdog.out 2>&1") - os._exit(3) -threading.Thread(target=watchdog, daemon=True).start() - -say("start", APP, "hard", HARD_S) -with modal.enable_output(): - with app.run(): - say("app run", app.app_id) - (OUT / "app_id.txt").write_text(str(app.app_id) + "\n") - t = time.time() - seed = seed_fn.remote({"volume": VOLUME, "store": mei.EVAL_STORE_MOUNT, "model": MODEL, "revision": REV, - "rank": 16, "targets": "attention", "sampling": SAMPLING, "seed": 0, - "training_finished": True}) - say("seed done", round(time.time() - t, 1), "s", json.dumps(seed, default=str)[:600]) - (OUT / "seed.json").write_text(json.dumps(seed, indent=1, default=str)) - launcher = mei.EvalIslandLauncher(client=mei.ModalFunctionClient(island_fn), emit=emit, config=config, - gpu="H100!", gpus=1, price_per_hour=None, max_starts=1) - try: - summary = launcher.run() - say("island done", json.dumps(summary, default=str)[:2000]) - (OUT / "summary.json").write_text(json.dumps(summary, indent=1, default=str)) - except Exception as exc: # noqa: BLE001 - say("island FAILED", type(exc).__name__, str(exc)[:2000]) -say("app exited") diff --git a/tools/reward_env/build_images.py b/tools/reward_env/build_images.py deleted file mode 100644 index 2988e7d06..000000000 --- a/tools/reward_env/build_images.py +++ /dev/null @@ -1,85 +0,0 @@ -#!/usr/bin/env python3 -"""Plan (default) or build prebaked reward-environment images on Modal (CPU only). - - # plan only: no network, no Modal; writes a JSON report - python tools/reward_env/build_images.py --benchmark tb2 --tasks-dir ~/work/tb2-data --out plan.json - - # build (needs a Modal token; costs CPU build time; needs approval first) - python tools/reward_env/build_images.py ... --build --app yeto-reward-env-build [--tasks a,b] - -``--build`` never targets the production app ``yeto-tbench2``. Each image is -built by Modal from ``Image.from_registry(base).run_commands()``; -Modal caches layers by content, so sandboxes later created from the same -definition (``yeto.cloud.modal_reward_env.PrebakedModalSandboxBackend``) reuse the build. -""" - -from __future__ import annotations - -import argparse -import json -import sys -import time -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - -from yeto.rl.harness.reward_env import benchmark as bm # noqa: E402 - - -def plan_report(adapter, task_ids): - rows = [] - for tid in task_ids: - plan = adapter.prebake_plan(tid) - rows.append({"task_id": tid, "base_image": plan.base_image, "digest": plan.digest(), - "complete": plan.complete, "commands": list(plan.commands), "skipped": list(plan.skipped)}) - return {"benchmark": adapter.name, "schema": bm.PREBAKE_SCHEMA, "tasks": rows, - "complete": sum(r["complete"] for r in rows), "total": len(rows)} - - -def build(adapter, task_ids, app_name): - import modal - - from yeto.cloud.modal_reward_env import modal_image - from yeto.rl.harness.reward_env.tb2 import check_build_app - - app = modal.App.lookup(check_build_app(app_name), create_if_missing=True) - results = [] - for tid in task_ids: - plan = adapter.prebake_plan(tid) - started = time.monotonic() - try: - with modal.enable_output(): - modal_image(plan).build(app) - results.append({"task_id": tid, "ok": True, "seconds": round(time.monotonic() - started, 1)}) - except Exception as exc: # noqa: BLE001 - report and continue - results.append({"task_id": tid, "ok": False, "error": f"{type(exc).__name__}: {exc}"[:500], - "seconds": round(time.monotonic() - started, 1)}) - print(json.dumps(results[-1]), flush=True) - return results - - -def main(argv=None) -> int: - p = argparse.ArgumentParser() - p.add_argument("--benchmark", default="tb2") - p.add_argument("--tasks-dir") - p.add_argument("--tasks", help="comma separated subset") - p.add_argument("--out") - p.add_argument("--build", action="store_true") - p.add_argument("--app") - a = p.parse_args(argv) - adapter = bm.get_adapter(a.benchmark, **({"tasks_dir": a.tasks_dir} if a.tasks_dir else {})) - ids = a.tasks.split(",") if a.tasks else adapter.task_ids() - report = plan_report(adapter, ids) - if a.build: - if not a.app: - p.error("--build needs an explicit --app (never yeto-tbench2)") - report["build"] = build(adapter, ids, a.app) - text = json.dumps(report, indent=2, ensure_ascii=False) - if a.out: - Path(a.out).write_text(text) - print(f"{report['complete']}/{report['total']} plans complete" if a.out else text) - return 0 if all(r.get("ok", True) for r in report.get("build", [])) else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tools/reward_env/holdout.py b/tools/reward_env/holdout.py index c9f5fcd9e..a271776c5 100644 --- a/tools/reward_env/holdout.py +++ b/tools/reward_env/holdout.py @@ -1,11 +1,11 @@ #!/usr/bin/env python3 """Write WP3 (rl-eval-difficulty-buckets D6b) hold-out lists for the built-in benchmarks. - python tools/reward_env/holdout.py --tb2-tasks-dir ~/work/tb2-data \ - --swev-data swev-78f471bf.parquet --out-dir data/eval + python tools/reward_env/holdout.py --swev-data swev-78f471bf.parquet --out-dir data/eval -TB2 writes tb2-holdout.json (eval) and tb2-train.json (training task list cut -against that hold-out). Tasks in data/eval/tb2-unusable.json leave both. +The TB2 lists (data/eval/tb2-holdout.json, tb2-train.json) were written by the +TB2 task loader of the self-written Codex harness, removed in S26 E0.5 (git tag +s26-pre-harness-removal has the generator). """ from __future__ import annotations @@ -22,56 +22,12 @@ def main(argv=None) -> int: p = argparse.ArgumentParser() - p.add_argument("--tb2-tasks-dir") p.add_argument("--swev-data") - p.add_argument("--tb2-exclude-jsonl", action="append", default=[], - help="dataset jsonl whose metadata.task_id rows leave the TB2 eval pool " - "(default: the built-in S15 smoke-6 list)") - p.add_argument("--tb2-exclude-reason", default=None, - help="reason recorded for --tb2-exclude-jsonl tasks (default: S15 smoke training)") - p.add_argument("--tb2-unusable", default=None, - help="unusable-task table (default: data/eval/tb2-unusable.json); its tasks leave " - "both the hold-out pool and the training list") - p.add_argument("--tb2-no-unusable", action="store_true", help="do not apply the unusable-task table") - p.add_argument("--tb2-no-exclude", action="store_true", - help="do not exclude any TB2 task (neither smoke-6 nor the unusable table)") p.add_argument("--out-dir", required=True) a = p.parse_args(argv) out = Path(a.out_dir) out.mkdir(parents=True, exist_ok=True) written = {} - if a.tb2_tasks_dir: - from yeto.rl.harness.reward_env import tb2 - - if a.tb2_no_exclude: - exclude = {} - elif a.tb2_exclude_jsonl: - reason = a.tb2_exclude_reason or tb2.SMOKE6_REASON - exclude = {} - for path in a.tb2_exclude_jsonl: - for line in Path(path).read_text().splitlines(): - if line.strip(): - tid = bm.task_id_of(json.loads(line)) - if tid is None: - raise SystemExit(f"{path}: row without metadata.task_id") - exclude[tid] = reason - else: - exclude = tb2.smoke6_exclusions() - unusable = {} if (a.tb2_no_exclude or a.tb2_no_unusable) else tb2.load_unusable(a.tb2_unusable) - clash = sorted(set(unusable) & set(exclude)) - if clash: - raise SystemExit(f"tasks both in the exclusion list and the unusable table: {clash}") - adapter = tb2.Tb2Benchmark(a.tb2_tasks_dir) - unknown = sorted(set(unusable) - set(adapter.task_ids())) - if unknown: - raise SystemExit(f"unusable tasks not in the TB2 checkout: {unknown}") - holdout = tb2.build_holdout(adapter, exclude={**exclude, **unusable}, - **({"rule": tb2.HOLDOUT_RULE_UNUSABLE} if unusable else {})) - unknown = sorted({e["task_id"] for e in holdout.get("excluded", [])} - set(adapter.task_ids())) - if unknown: - raise SystemExit(f"excluded tasks not in the TB2 checkout: {unknown}") - written["tb2-holdout.json"] = holdout - written["tb2-train.json"] = bm.build_train_split(adapter, holdout, exclude=unusable) if a.swev_data: from yeto.rl.harness.reward_env import swebench_verified as sv diff --git a/yeto/cli.py b/yeto/cli.py index 7c68497f8..d03d1b5a4 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -1256,10 +1256,6 @@ def _make_head_task(args, extra_mounts: dict | None = None): # The head's controller defines, deploys and spawns modal: islands # through the Modal SDK (yeto.modal_runner), so it needs the package. head_pip += ' && pip install -q "modal>=1.0"' - if getattr(args, "codex_bundle_dir", None): - # S17 M1: the head attests the Codex bundle and imports the harness - # preflight (aiohttp / dill) when it builds each island. - head_pip += " && pip install -q aiohttp dill" if getattr(args, "wandb", False): # The head tails the syncer's event tape into W&B (yeto.wandb_tape). head_pip += " && pip install -q wandb" @@ -1429,16 +1425,6 @@ def cmd_launch_head(args) -> int: args.rl_initial_adapter ) args.rl_initial_adapter = launcher.RL_HEAD_INITIAL_ADAPTER_PATH - # S17 M1: a signed Codex run needs its bundle and harness env on the head, - # which builds the islands (attested here first; fails before provisioning). - try: - codex_mounts, codex_secrets = launcher.codex_head_staging(args) - except ValueError as exc: - print(f"[yeto] {exc}", file=sys.stderr) - return 1 - if codex_mounts: - data_mounts.update(codex_mounts) - args.codex_bundle_dir = launcher.HEAD_CODEX_BUNDLE_PATH args_dict = _serializable_args(args) runs.create_run(name, args_dict) learner_names = launcher.learner_cluster_names(name, specs) @@ -1488,7 +1474,6 @@ def cmd_launch_head(args) -> int: # 0.20: the head's LocalSyncer and the islands it launches read the # island HMAC key from this secret env (elastic only; {} in legacy). secrets.update(launcher.island_hmac_secret(args)) - secrets.update(codex_secrets) # secret-handling-hardening: HF/W&B/CyberGym tokens ride as secrets too. envs, env_secrets = launcher.split_secret_envs(envs) secrets.update(env_secrets) diff --git a/yeto/cloud/modal_ckpt_store.py b/yeto/cloud/modal_ckpt_store.py index f7a4bf3fa..f98140865 100644 --- a/yeto/cloud/modal_ckpt_store.py +++ b/yeto/cloud/modal_ckpt_store.py @@ -2,7 +2,7 @@ (rl-resume-from-checkpoint). Moved out of ``yeto/rl/engine/resume.py`` (the neutral core must not import -``modal``; same pattern as ``yeto.cloud.modal_reward_env``). The core only knows +``modal``; same pattern as the other ``yeto.cloud.modal_*`` modules). The core only knows ``CheckpointStore`` and the ``YETO_RL_STORE_IMPL`` hook; the launcher exports ``YETO_RL_STORE_IMPL=yeto.cloud.modal_ckpt_store:store_from_env`` plus ``YETO_RL_STORE_MODAL_VOLUME=`` on a Modal island with a diff --git a/yeto/cloud/modal_eval_island.py b/yeto/cloud/modal_eval_island.py index 2e9fe3972..08845ab4e 100644 --- a/yeto/cloud/modal_eval_island.py +++ b/yeto/cloud/modal_eval_island.py @@ -253,7 +253,6 @@ def timed_commit() -> None: # --- the Modal GPU function (5.10) ------------------------------------------------- EVAL_STORE_MOUNT = "/mnt/yeto-eval-store" # the eval store Volume inside the eval island -CODEX_MOUNT = "/opt/yeto/codex" # yeto.rl.CODEX_CONTAINER_DIR (signed codex bundle) TB2_MOUNT = "/root/tb2-data" # terminal-bench-2 checkout (YETO_HARNESS_TB2_TASKS_DIR) WORKDIR_MOUNT = "/root/sky_workdir" # yeto source, same place as the training islands IMAGE_PYTHONPATH = "/root/miles:/root/sglang/python:/root/sky_workdir" @@ -270,13 +269,12 @@ def default_image_ref() -> str: @dataclass class EvalFunctionSpec: """Shape of the eval-island Modal function. ``gpu`` is a Modal GPU string - (``H100!`` = exactly H100); ``envs`` become one Modal secret (codex - contract env, Modal token for the judge sandboxes, TB2 knobs).""" + (``H100!`` = exactly H100); ``envs`` become one Modal secret (Modal + token for the judge sandboxes, TB2 knobs).""" app_name: str volume: str workdir: str - codex_dir: str tb2_dir: str gpu: str = "H100!" gpus: int = 1 @@ -315,7 +313,7 @@ def _gpu_sampler(path: Path, stop: Any, interval_s: float = 5.0) -> None: def _image_paths(environ: Any = None) -> None: """Miles / SGLang / yeto source on ``sys.path``; subprocesses (SGLang, session - server, codex worker) get exactly the training island's ``PYTHONPATH``, without + server) get exactly the training island's ``PYTHONPATH``, without Modal's runtime copies (their protobuf/grpclib would shadow the image's).""" import sys @@ -416,7 +414,6 @@ def build_eval_app(spec: EvalFunctionSpec) -> tuple[Any, Any, Any]: .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"}) .add_local_dir(spec.workdir, WORKDIR_MOUNT, copy=False, ignore=[".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs"]) - .add_local_dir(spec.codex_dir, CODEX_MOUNT, copy=False) .add_local_dir(spec.tb2_dir, TB2_MOUNT, copy=False, ignore=[".git"])) app = modal.App(spec.app_name) volume = modal.Volume.from_name(spec.volume, create_if_missing=True) diff --git a/yeto/cloud/modal_reward_env.py b/yeto/cloud/modal_reward_env.py deleted file mode 100644 index 104fbc687..000000000 --- a/yeto/cloud/modal_reward_env.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Cloud layer: Modal sandbox backend for the reward environment's TB2 adapter. - -Moved out of ``yeto/rl/harness/reward_env/tb2.py`` (neutral core must not -import ``modal``; spec rl-framework-neutral-core). The neutral adapter only -knows ``tb2_provider.SandboxBackend`` / ``Tb2EnvironmentProvider``; this module -is the Modal implementation, injected at launch with -``YETO_HARNESS_ENVIRONMENT_PROVIDER=yeto.cloud.modal_reward_env:modal_provider``. -Behaviour is unchanged from the original location. -""" - -from __future__ import annotations - -import base64 -import os -from typing import Any - -from yeto.rl.harness.codex import tb2_provider as tb2 -from yeto.rl.harness.reward_env import tb2 as rt -from yeto.rl.harness.reward_env.benchmark import PrebakePlan -from yeto.rl.rlf.modal import RlfCtfConfig, RlfModalProvider, RlfModalVmConfig - - - -def modal_image(plan: PrebakePlan, *, prebake: bool = True) -> Any: - from modal import Image - - image = Image.from_registry(plan.base_image) - if prebake and not plan.empty: - image = image.run_commands(prebake_run_command(plan)) - return image - - -def prebake_run_command(plan: PrebakePlan) -> str: - """One Dockerfile ``RUN`` line for the multi-line prebake script. - - A Dockerfile instruction cannot span raw newlines (S17 G3: Modal rejected - ``bash -lc ''`` with "could not parse Dockerfile"), so the - script travels base64-encoded and runs as a login shell from a temp file. - """ - b64 = base64.b64encode(rt.prebake_script(plan).encode()).decode() - return (f"echo {b64} | base64 -d > /tmp/yeto-prebake.sh && bash -l /tmp/yeto-prebake.sh" - " && rm -f /tmp/yeto-prebake.sh") - - -def _prebake_enabled() -> bool: - return os.environ.get(rt.PREBAKE_ENV, "1") != "0" - - -class PrebakedModalSandboxBackend(tb2.ModalSandboxBackend): - """``tb2_provider.ModalSandboxBackend`` with the task's prebaked image.""" - - def __init__(self, *, prebake: bool = True, **kwargs: Any) -> None: - kwargs.setdefault("app_name", rt.DEFAULT_APP) - super().__init__(**kwargs) - self.prebake = prebake - - def create(self, task: tb2.Tb2Task, trajectory_id: str) -> tb2.ModalSandbox: - from modal import Sandbox - - plan = rt.prebake_from_test_sh((task.tests_dir / "test.sh").read_text(), task.docker_image) - tags = {"yeto-tb2-task": task.task_id, "yeto-trajectory": trajectory_id[:63], - "yeto-prebake": plan.digest()[:16] if self.prebake and not plan.empty else "none"} - if self.run_id: - tags["yeto-run-id"] = self.run_id[:63] - sandbox = Sandbox.create( - "sleep", "infinity", app=self._get_app(), image=modal_image(plan, prebake=self.prebake), - timeout=self.ttl_s, idle_timeout=self.idle_timeout_s, cpu=float(task.cpus), - memory=int(task.memory_mb), workdir=task.workdir, tags=tags, - **self._client_kwargs(), - **self.network_policy.grant(task.task_id).modal_kwargs(), - ) - return tb2.ModalSandbox(sandbox, task.workdir) - - -def modal_provider(miles_args: Any = None) -> tb2.Tb2EnvironmentProvider: - """``YETO_HARNESS_ENVIRONMENT_PROVIDER=yeto.cloud.modal_reward_env:modal_provider``.""" - del miles_args - tb2.require_modal_client() - backend = PrebakedModalSandboxBackend( - prebake=_prebake_enabled(), - app_name=os.environ.get(rt.APP_ENV, rt.DEFAULT_APP), - ttl_s=int(tb2._env_float(tb2.SANDBOX_TTL_ENV, tb2.DEFAULT_SANDBOX_TTL_S) or tb2.DEFAULT_SANDBOX_TTL_S), - idle_timeout_s=int(tb2._env_float(tb2.IDLE_TIMEOUT_ENV, tb2.DEFAULT_IDLE_TIMEOUT_S) - or tb2.DEFAULT_IDLE_TIMEOUT_S), - network_policy=tb2.NetworkPolicy.from_env(), - ) - return tb2._provider(backend) - - -RLF_CTF_TASKS_DIR_ENV = "YETO_RLF_CTF_TASKS_DIR" -RLF_CTF_FLAGS_ENV = "YETO_RLF_CTF_FLAGS_JSON" # {"task_id": "flag{...}"} (trusted only) -RLF_CTF_VM_IMAGE_ENV = "YETO_RLF_CTF_VM_IMAGE" -RLF_CTF_LEASE_SECONDS_ENV = "YETO_RLF_CTF_LEASE_SECONDS" - - -def rlf_ctf_config_from_env(environ: Any = None) -> RlfCtfConfig | None: - """CTF packaging config, or None when the launch did not configure it. - - Fail closed: setting ``YETO_RLF_CTF_TASKS_DIR`` without flags (or with - invalid JSON) raises instead of silently degrading to the VM-only lease. - """ - import json - from pathlib import Path - - environ = os.environ if environ is None else environ - tasks_dir = environ.get(RLF_CTF_TASKS_DIR_ENV) - if not tasks_dir: - return None - if not Path(tasks_dir).is_dir(): - raise RuntimeError(f"{RLF_CTF_TASKS_DIR_ENV}={tasks_dir}: not a directory") - raw_flags = environ.get(RLF_CTF_FLAGS_ENV) - if not raw_flags: - raise RuntimeError(f"{RLF_CTF_TASKS_DIR_ENV} is set but {RLF_CTF_FLAGS_ENV} is not " - "(the trusted flags never live in the task package)") - flags = json.loads(raw_flags) - if not isinstance(flags, dict) or not all(isinstance(k, str) and isinstance(v, str) and v - for k, v in flags.items()): - raise RuntimeError(f"{RLF_CTF_FLAGS_ENV} must be a JSON object of task_id -> flag") - lease = float(environ.get(RLF_CTF_LEASE_SECONDS_ENV) or 1800.0) - return RlfCtfConfig(tasks_dir=tasks_dir, flags=flags, - vm_image=environ.get(RLF_CTF_VM_IMAGE_ENV) or "docker:27-dind", - lease_seconds=lease) - - -def rlf_modal_provider(miles_args: Any = None, *, - config: RlfModalVmConfig | None = None, - ctf: RlfCtfConfig | None = None) -> RlfModalProvider: - """``YETO_HARNESS_ENVIRONMENT_PROVIDER=yeto.cloud.modal_reward_env:rlf_modal_provider``. - - Same factory contract as ``modal_provider`` (preflight calls - ``factory(miles_args)`` and requires ``.acquire`` on the product), but each - acquire opens one isolated Modal VM carrying RLF task/episode/trajectory - identity instead of a TB2 sandbox. - - When CTF packaging is configured (``ctf=`` or ``YETO_RLF_CTF_TASKS_DIR`` + - ``YETO_RLF_CTF_FLAGS_JSON``), the provider is a full ``EnvironmentProvider``: - each acquire additionally runs the packaged nested-Docker CTF orchestration - and the lease exposes env_url/env_token/verifier/destroy/describe/deadline - (audit breakpoint 3). Without it, the S20 VM-only lease is returned. - """ - del miles_args - return RlfModalProvider(config=config, ctf=ctf if ctf is not None else rlf_ctf_config_from_env()) diff --git a/yeto/cloud/modal_rlf_baseline.py b/yeto/cloud/modal_rlf_baseline.py deleted file mode 100644 index 6ede1d4e3..000000000 --- a/yeto/cloud/modal_rlf_baseline.py +++ /dev/null @@ -1,193 +0,0 @@ -"""Modal deployment shape for the RLF Qwen3-30B-A3B baseline (rlf-codex-rl 5.6). - -1xA100-80GB container running ``yeto.rl.eval.rlf_baseline.run_from_config``: -SGLang bf16 from the ``yeto-fn-models`` Volume -(``hf/Qwen3-30B-A3B/ad44e777/``) + Miles session server + the Codex harness -against per-episode RLF CTF VMs (created by the in-container Modal client). - -Mirrors ``modal_eval_island.build_eval_app``: same pinned Miles-ports image, -same workdir/codex mounts. Launch (like the S17 N11 smoke): - - python tools/probes/... style driver, or: - - modal run yeto.cloud.modal_rlf_baseline --config-json run.json - -Local entrypoint inputs (never baked into the image): - -* ``YETO_CODEX_BUNDLE_DIR`` — the attested Codex bundle dir - (``scripts/fetch_codex_bundle.py``), mounted at ``/opt/yeto/codex``; -* ``YETO_RLF_CTF_TASKS_DIR`` — local dir of admitted task packages, mounted - at ``/root/rlf-tasks`` (the container-side value overrides it); -* ``YETO_RLF_BASELINE_ENVS_JSON`` — JSON file with the container env: - the Codex contract env (``yeto.launcher.codex_container_env( - codex_bundle_contract(...))``, same as the eval island's codex-env.json), - ``YETO_RLF_CTF_FLAGS_JSON`` (trusted flags; - live only in the Modal Secret, never in the image or the traces) and the - sandbox-only ``MODAL_TOKEN_ID`` / ``MODAL_TOKEN_SECRET`` the RLF provider - needs to open episode VMs. -""" - -from __future__ import annotations - -import json -import os -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Mapping - -APP_NAME = "yeto-rlf-baseline" -BASELINE_GPU = "A100-80GB" -MODELS_VOLUME = "yeto-fn-models" # hf/// + MANIFEST.json -RESULTS_VOLUME = "yeto-rlf-baseline-runs" -MODELS_MOUNT = "/mnt/yeto-models" # yeto.launcher.MODEL_STORE_MOUNT -RESULTS_MOUNT = "/mnt/yeto-rlf-runs" -WORKDIR_MOUNT = "/root/sky_workdir" # same as the training/eval islands -CODEX_MOUNT = "/opt/yeto/codex" # yeto.launcher.CODEX_CONTAINER_DIR -TASKS_MOUNT = "/root/rlf-tasks" -ENVS_JSON_ENV = "YETO_RLF_BASELINE_ENVS_JSON" - -_SECRET_ENV_KEYS = ("YETO_RLF_CTF_FLAGS_JSON", "MODAL_TOKEN_ID", "MODAL_TOKEN_SECRET") - - -def default_image_ref() -> str: - from yeto.rl import MILES_NEXT_IMAGE # re-exported pin (same route as the eval island) - - return MILES_NEXT_IMAGE.removeprefix("docker:") - - -@dataclass -class BaselineFunctionSpec: - """Shape of the baseline Modal function; ``gpu`` is a Modal GPU string.""" - - app_name: str - workdir: str - codex_dir: str - tasks_dir: str - models_volume: str = MODELS_VOLUME - results_volume: str = RESULTS_VOLUME - gpu: str = BASELINE_GPU - cpu: float = 8.0 - memory_mib: int = 65536 - timeout_s: int = 6 * 3600 - image_ref: str = field(default_factory=default_image_ref) - envs: Mapping[str, str] = field(default_factory=dict) - - def container_env(self) -> dict[str, str]: - # No PYTHONPATH override (S17 N11: it breaks Modal's container entrypoint). - return {"HOME": "/root", "PYTHONUNBUFFERED": "1", "YETO_RLF_CTF_TASKS_DIR": TASKS_MOUNT, - **{str(k): str(v) for k, v in self.envs.items()}} - - -def build_baseline_app(spec: BaselineFunctionSpec) -> tuple[Any, Any]: - """(app, baseline GPU function). Models Volume read-only; results Volume committed.""" - import modal - - image = (modal.Image.from_registry(spec.image_ref) - .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"}) - .add_local_dir(spec.workdir, WORKDIR_MOUNT, copy=False, - ignore=[".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs"]) - .add_local_dir(spec.codex_dir, CODEX_MOUNT, copy=False) - .add_local_dir(spec.tasks_dir, TASKS_MOUNT, copy=False)) - app = modal.App(spec.app_name) - models = modal.Volume.from_name(spec.models_volume, create_if_missing=True) - results = modal.Volume.from_name(spec.results_volume, create_if_missing=True) - secret = modal.Secret.from_dict(spec.container_env()) - run = app.function(image=image, gpu=spec.gpu, cpu=spec.cpu, memory=spec.memory_mib, - timeout=spec.timeout_s, secrets=[secret], - volumes={MODELS_MOUNT: models, RESULTS_MOUNT: results}, - name="rlf_baseline")(baseline_body) - return app, run - - -def baseline_body(config: Mapping[str, Any]) -> dict[str, Any]: - """In-container entry: run the baseline, then copy the evidence onto the results Volume.""" - import shutil - import sys - import time - - t0 = time.time() - sys.path.insert(0, WORKDIR_MOUNT) - from yeto.rl.eval import rlf_baseline - - config = dict(config) - local_out = "/tmp/rlf-baseline-out" - config["out_dir"] = local_out - config.setdefault("tasks_dir", TASKS_MOUNT) - # capture sglang/session-server logs into the results volume copy - config.setdefault("baseline", {}) - config["baseline"].setdefault("log_dir", local_out + "/logs") - # Trusted flags arrive via the Modal Secret env, never via run.json on disk. - if not config.get("flags"): - config["flags"] = json.loads(os.environ.get("YETO_RLF_CTF_FLAGS_JSON") or "{}") - run_name = str(config.get("run_name") or time.strftime("rlf-%Y%m%d-%H%M%S", time.gmtime())) - result = rlf_baseline.run_from_config(config) - final_dir = Path(RESULTS_MOUNT) / run_name - shutil.copytree(local_out, final_dir, dirs_exist_ok=True) - import modal - - modal.Volume.from_name(str(config.get("results_volume") or RESULTS_VOLUME), - create_if_missing=True).commit() - return {"run_name": run_name, "results": str(final_dir), "wall_s": round(time.time() - t0, 3), - **result} - - -def spec_from_env(environ: Mapping[str, str] | None = None) -> BaselineFunctionSpec: - """Local-side spec from the environment (see the module docstring).""" - environ = os.environ if environ is None else environ - repo = Path(__file__).resolve().parents[2] - codex_dir = environ.get("YETO_CODEX_BUNDLE_DIR") - tasks_dir = environ.get("YETO_RLF_CTF_TASKS_DIR") - if not codex_dir or not Path(codex_dir).is_dir(): - raise RuntimeError("YETO_CODEX_BUNDLE_DIR must point at the attested Codex bundle dir") - if not tasks_dir or not Path(tasks_dir).is_dir(): - raise RuntimeError("YETO_RLF_CTF_TASKS_DIR must point at the local admitted task packages") - envs: dict[str, str] = {} - envs_json = environ.get(ENVS_JSON_ENV) - if envs_json: - envs.update({str(k): str(v) for k, v in json.loads(Path(envs_json).read_text()).items()}) - for key in _SECRET_ENV_KEYS: - if environ.get(key) and key not in envs: - envs[key] = str(environ[key]) - return BaselineFunctionSpec(app_name=environ.get("YETO_RLF_BASELINE_APP", APP_NAME), - workdir=str(repo), codex_dir=codex_dir, tasks_dir=tasks_dir, - envs=envs) - - -try: # the Modal entry point (only usable with the modal client installed) - import modal as _modal - - try: - app, baseline_run = build_baseline_app(spec_from_env()) - except RuntimeError: - # In-container hydration re-imports this module; the local-only spec - # env (bundle dir, tasks dir) is absent there. The container finds - # the function body (``baseline_body``) by qualname, not via ``app``. - app = None - baseline_run = None - - def main(argv: "list[str] | None" = None) -> int: - """Local driver: ``python -m yeto.cloud.modal_rlf_baseline run.json``. - - Driver pattern (same as the S17 eval-island probes): the app is built - locally and the run is spawned inside ``app.run()``; ``modal run -m`` - is NOT usable here because the spec needs local-only paths. - """ - import sys - - args = sys.argv[1:] if argv is None else argv - if app is None or baseline_run is None: - raise RuntimeError("baseline app is only built locally (spec env missing)") - payload = json.loads(Path(args[0]).read_text()) - with app.run(): - result = baseline_run.remote(payload) - print(json.dumps(result, indent=1, default=str)) - return 0 - -except ImportError: # pragma: no cover - CPU tests have no modal client - _modal = None - app = None - baseline_run = None - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/yeto/cloud/modal_tb2_baseline.py b/yeto/cloud/modal_tb2_baseline.py deleted file mode 100644 index 04eec3b89..000000000 --- a/yeto/cloud/modal_tb2_baseline.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Modal deployment shape for the TB2 base-model baseline (``yeto.rl.eval.tb2_baseline``). - -One GPU container: the pinned Miles-ports image (same mounts and env as the -eval island), the base model read-only from the models Volume -(``hf///``), TB2 sandboxes created by the in-container Modal -client. Evidence (traces, summary, SGLang / session-server logs, nvidia-smi -samples) is copied to the results Volume every ``sync_s`` seconds and once at -the end, so a container killed by its timeout still leaves its traces. - -The driver builds the app locally (the spec needs local-only paths) and must -``spawn`` the call inside ``app.run(detach=True)`` and return: killing a client -that blocks on ``.remote()`` cancels the input. -""" - -from __future__ import annotations - -import json -import os -import shutil -import threading -import time -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Mapping - -from .modal_eval_island import (CODEX_MOUNT, TB2_MOUNT, WORKDIR_MOUNT, _gpu_sampler, _image_paths, - default_image_ref) - -MODELS_VOLUME = "yeto-fn-models" -MODELS_MOUNT = "/mnt/yeto-models" -RESULTS_VOLUME = "yeto-tb2-baseline-runs" -RESULTS_MOUNT = "/mnt/yeto-tb2-baseline" -FUNCTION_NAME = "tb2_baseline" - - -@dataclass -class BaselineFunctionSpec: - """``gpu`` is a Modal GPU string pinned to one type (``H200``; ``H100!`` keeps Modal from upgrading H100 to H200); ``envs`` become one secret - (codex contract env, Modal token for the TB2 sandboxes, TB2 / worker knobs).""" - - app_name: str - workdir: str - codex_dir: str - tb2_dir: str - gpu: str = "H200" - gpus: int = 1 - cpu: float = 16.0 - memory_mib: int = 131072 - timeout_s: int = 4 * 3600 - image_ref: str = field(default_factory=default_image_ref) - envs: Mapping[str, str] = field(default_factory=dict) - - def container_env(self) -> dict[str, str]: - # No PYTHONPATH here (see EvalFunctionSpec.container_env). - return {"HOME": "/root", "PYTHONUNBUFFERED": "1", "YETO_HARNESS_TB2_TASKS_DIR": TB2_MOUNT, - **{str(k): str(v) for k, v in self.envs.items()}} - - -def _sync(src: Path, dst: Path, volume: Any) -> None: - shutil.copytree(src, dst, dirs_exist_ok=True) - volume.commit() - - -def baseline_body(config: Mapping[str, Any]) -> dict[str, Any]: - """Inside the container: assert the GPU type, run the baseline, keep the evidence on the results Volume.""" - import subprocess - - import modal - - _image_paths() - from yeto.rl.eval import tb2_baseline - - names = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], - capture_output=True, text=True, timeout=30).stdout.split("\n") - names = [n.strip() for n in names if n.strip()] - want = str(config["assert_gpu_name"]) - if not names or not all(want in n for n in names): - raise RuntimeError(f"baseline asked for {want}, got {names}") - run_name = str(config["run_name"]) - local = Path("/tmp/tb2-baseline") / run_name - final = Path(RESULTS_MOUNT) / run_name - local.mkdir(parents=True, exist_ok=True) - (local / "body.json").write_text(json.dumps({"gpu_names": names, "container": os.environ.get("MODAL_TASK_ID"), - "start_ts": time.time()}, sort_keys=True)) - results = modal.Volume.from_name(RESULTS_VOLUME, create_if_missing=True) - stop = threading.Event() - sampler = threading.Thread(target=_gpu_sampler, args=(local / "gpu.jsonl", stop), daemon=True) - sampler.start() - - def syncer() -> None: - while not stop.wait(float(config.get("sync_s", 120))): - try: - _sync(local, final, results) - except Exception as exc: # noqa: BLE001 - the next tick retries - print(f"[tb2-baseline] sync failed: {exc}", flush=True) - - threading.Thread(target=syncer, daemon=True).start() - infer = dict(config.get("infer") or {}) - infer.setdefault("log_dir", str(local / "logs")) - try: - summary = tb2_baseline.run_from_config({**config, "infer": infer, "out_dir": str(local)}) - finally: - stop.set() - sampler.join(timeout=10) - _sync(local, final, results) - return {"run_name": run_name, "results": str(final), **{k: summary.get(k) for k in - ("attempts", "judged", "passed", "pass_rate", "infra_errors", "tasks_mixed", "attempts_wall_s")}} - - -def build_baseline_app(spec: BaselineFunctionSpec) -> tuple[Any, Any]: - """(app, baseline GPU function). Models Volume read-only.""" - import modal - - image = (modal.Image.from_registry(spec.image_ref) - .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"}) - .add_local_dir(spec.workdir, WORKDIR_MOUNT, copy=False, - ignore=[".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs", - "syncer/target"]) - .add_local_dir(spec.codex_dir, CODEX_MOUNT, copy=False) - .add_local_dir(spec.tb2_dir, TB2_MOUNT, copy=False, ignore=[".git"])) - app = modal.App(spec.app_name) - models = modal.Volume.from_name(MODELS_VOLUME) - results = modal.Volume.from_name(RESULTS_VOLUME, create_if_missing=True) - gpu = spec.gpu if spec.gpus == 1 else f"{spec.gpu}:{spec.gpus}" - run = app.function(image=image, gpu=gpu, cpu=spec.cpu, memory=spec.memory_mib, timeout=spec.timeout_s, - secrets=[modal.Secret.from_dict(spec.container_env())], retries=0, - volumes={MODELS_MOUNT: models.read_only(), RESULTS_MOUNT: results}, - name=FUNCTION_NAME)(baseline_body) - return app, run diff --git a/yeto/launch/args.py b/yeto/launch/args.py index 81e0de97d..cac76db7e 100644 --- a/yeto/launch/args.py +++ b/yeto/launch/args.py @@ -10,7 +10,7 @@ from .loss import pickled_loss_path, resolve_loss_function from .storage import modal_checkpoint_store from .rl_args import _prepare_rl_args -from .codex import modal_sandbox_secret_name +from .harness_env import modal_sandbox_secret_name def prepare_launch_args( diff --git a/yeto/launch/codex.py b/yeto/launch/codex.py deleted file mode 100644 index 18c63f1e3..000000000 --- a/yeto/launch/codex.py +++ /dev/null @@ -1,314 +0,0 @@ -"""Codex harness and reward-sandbox environment for RL islands.""" - -from __future__ import annotations - -import json -import os -import re -import shlex -from pathlib import Path - - -# --- Stock Codex run bundle (ports engine; design R-SCOPE / T3 G2) --------------- -# -# The legacy direct SSH harness freezes the attested Codex artifact into the run -# bundle and ``--volume``-mounts it at /opt/yeto/codex (ssh_harness.py). The -# sky / Modal launch paths do the same from a local bundle directory -# (``scripts/fetch_codex_bundle.py``): the launcher attests it with the same -# ``_codex_harness_contract`` the SSH harness uses, mounts the directory at -# ``CODEX_CONTAINER_DIR`` (sky file_mounts / Modal add_local_dir) and injects -# the ``YETO_CODEX_*`` environment the container preflights check -# (``yeto.rl.adapters.miles.island_entry._preflight_codex_harness`` expected_env and -# ``codex_harness_agent._attest_runtime``); every value comes from the contract. -CODEX_BUNDLE_DIR_ENV = "YETO_CODEX_BUNDLE_DIR" -CODEX_CONTAINER_DIR = "/opt/yeto/codex" -CODEX_BUNDLE_FILES = ( - "codex_harness_binary", - "codex_package_manifest", - "codex_app_server_schema", -) -# Harness hook / provider env the ports island preflight resolves -# (entry.HARNESS_PREFLIGHT_ENV, preflight.ENVIRONMENT_PROVIDER_ENV); the -# provider must come from the launching environment (no in-tree default). -HARNESS_PREFLIGHT_ENV = "YETO_HARNESS_PREFLIGHT" -HARNESS_ENVIRONMENT_PROVIDER_ENV = "YETO_HARNESS_ENVIRONMENT_PROVIDER" -HARNESS_PASSTHROUGH_ENV = ( - HARNESS_ENVIRONMENT_PROVIDER_ENV, - # tb2_provider Modal Sandbox backend: the island creates task sandboxes - # itself. secret-handling-hardening: it gets a separate sandbox-only Modal - # token (YETO_SANDBOX_MODAL_TOKEN_*), never the main MODAL_TOKEN_*. - "YETO_SANDBOX_MODAL_TOKEN_ID", - "YETO_SANDBOX_MODAL_TOKEN_SECRET", - "OPENENV_RUN_ID", - "SECRLENV_MAX_TURNS", -) -HARNESS_PASSTHROUGH_ENV_PREFIXES = ("YETO_HARNESS_TB2_", "YETO_RLF_") -# tb2_provider's Modal Sandbox backend runs inside the island (rollout -# workers create the task sandboxes), so the island python needs the Modal -# client. The ports image does not ship it; install it in setup without -# touching the image's protobuf 7 (modal pins <7 but works; verified in-image -# on ghcr 12fcd9e5: Sandbox create/exec/terminate, ray/sglang/miles import). -MODAL_SANDBOX_PROVIDER = "yeto.rl.harness.codex.tb2_provider:modal_provider" -# rl-agentic-reward-env: same TB2 provider, sandboxes from prebaked images. -PREBAKED_MODAL_SANDBOX_PROVIDER = "yeto.cloud.modal_reward_env:modal_provider" -RLF_MODAL_SANDBOX_PROVIDER = "yeto.cloud.modal_reward_env:rlf_modal_provider" -MODAL_SANDBOX_PROVIDERS = frozenset({MODAL_SANDBOX_PROVIDER, PREBAKED_MODAL_SANDBOX_PROVIDER, RLF_MODAL_SANDBOX_PROVIDER}) -SANDBOX_MODAL_TOKEN_ENVS = ("YETO_SANDBOX_MODAL_TOKEN_ID", "YETO_SANDBOX_MODAL_TOKEN_SECRET") - - -def require_sandbox_modal_token(envs) -> None: - """secret-handling-hardening: a Modal Sandbox provider on the island needs - the sandbox-only token; the main MODAL_TOKEN_* never goes to an island.""" - missing = [n for n in SANDBOX_MODAL_TOKEN_ENVS if not envs.get(n)] - if missing: - raise ValueError( - f"the Modal Sandbox provider needs {', '.join(missing)} in the launching " - "environment: a separate Modal token used only for task sandboxes (the main " - "MODAL_TOKEN_ID/SECRET is never shipped to a learner island)") -# --modal-sandbox-secret: in-island fail-closed check (names only, never values). -SANDBOX_TOKEN_IN_ISLAND_CHECK = ( - "for _v in " + " ".join(SANDBOX_MODAL_TOKEN_ENVS) + "; do " - '[ -n "$(printenv "$_v")" ] || { echo "[yeto-setup] $_v missing in the island: ' - 'the --modal-sandbox-secret Modal Secret must hold it" >&2; exit 1; }; done' -) - - -def modal_sandbox_secret_name(args, spec) -> str | None: - """--modal-sandbox-secret NAME, or None when unset. Only Modal islands can - attach a named Modal Secret: any other cloud is refused, never ignored.""" - name = getattr(args, "modal_sandbox_secret", None) - if not name: - return None - if getattr(spec, "cloud", None) != "modal": - raise ValueError( - f"--modal-sandbox-secret attaches a Modal Secret to a Modal island; " - f"island cloud {getattr(spec, 'cloud', None)!r} cannot use it") - if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,63}", str(name)): - raise ValueError(f"--modal-sandbox-secret {name!r} is not a Modal Secret name") - return str(name) - - -MODAL_CLIENT_SETUP = ( - # --ignore-installed: in a Modal Function container the setup shell - # already sees Modal's runtime copies (/pkg, /__modal/deps), so a plain - # install reports "already satisfied" and writes nothing into the island - # python's site-packages, which is all its own processes can import - # (codex-smoke-20261003-8/-9). - "python3 -m pip install -q --no-deps --ignore-installed " - "'modal==1.5.5' 'grpclib>=0.4.7,<0.4.10' 'synchronicity~=0.12.5' cbor2 toml " - "types-certifi types-toml watchfiles\n" - # Verify the packages physically landed in site-packages and import under - # the run script's PYTHONPATH (what the learner and Ray workers see). - 'PYTHONPATH="$HOME/miles:$HOME/sglang/python:$HOME/sky_workdir${PYTHONPATH:+:$PYTHONPATH}" ' - "python3 -c 'import os, sys, sysconfig; import modal, grpclib, synchronicity; " - "sp = sysconfig.get_paths()[\"purelib\"]; " - "missing = [m for m in (\"modal\", \"grpclib\", \"synchronicity\") if not os.path.isdir(os.path.join(sp, m))]; " - "sys.exit(f\"not in {sp}: {missing}\" if missing else 0)' " - "|| { echo '[yeto-setup] Modal client unusable in the island python' >&2; exit 1; }" -) -CODEX_COMPACTION_ENV = "YETO_CODEX_COMPACTION_ENABLED" - - -def codex_bundle_paths(bundle_dir: str) -> dict[str, str]: - """The three attested files inside a fetched bundle dir, keyed like the SSH - harness CLI (``--codex-harness-binary`` ...).""" - from ..rl import CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH, CODEX_CONTAINER_BINARY_PATH - - root = Path(bundle_dir).expanduser() - return { - "codex_harness_binary": str(root / os.path.basename(CODEX_CONTAINER_BINARY_PATH)), - "codex_package_manifest": str(root / "codex-package.json"), - "codex_app_server_schema": str( - root / os.path.basename(CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH) - ), - } - - -def codex_bundle_contract(args, bundle_dir: str) -> dict: - """Attest ``bundle_dir`` against the Yeto pins (same checks as the SSH harness).""" - from argparse import Namespace - - from ..rl.ssh_harness import HarnessError, _codex_harness_contract - - if not os.path.isdir(os.path.expanduser(bundle_dir)): - raise ValueError(f"{CODEX_BUNDLE_DIR_ENV}={bundle_dir!r} is not a directory") - try: - return _codex_harness_contract(Namespace(**codex_bundle_paths(bundle_dir)), args) - except HarnessError as exc: - raise ValueError(f"Codex bundle {bundle_dir}: {exc}") from exc - - -def codex_container_env(contract: dict) -> dict[str, str]: - """``YETO_CODEX_*`` for the island container, mirroring the SSH harness - ``codex_env`` template and the learner's expected_env, from the contract.""" - from ..rl.ssh_harness import _canonical_json, _plan_digest - - backend = contract["backend"] - env = { - "YETO_CODEX_BINARY_PATH": contract["container_binary_path"], - "YETO_CODEX_BINARY_SHA256": contract["binary_sha256"], - "YETO_CODEX_BINARY_SIZE_BYTES": str(contract["binary_size_bytes"]), - "YETO_CODEX_VERSION": contract["cli_version"], - "YETO_CODEX_APP_SERVER_PROTOCOL_REVISION": contract["app_server_protocol_revision"], - "YETO_CODEX_APP_SERVER_SCHEMA_SHA256": contract["app_server_schema_sha256"], - "YETO_CODEX_BASE_INSTRUCTIONS_SHA256": contract["base_instructions_sha256"], - "YETO_CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256": contract["terminal_exec_tool_schema_sha256"], - "YETO_CODEX_SUBMIT_TOOL_SCHEMA_SHA256": contract["submit_tool_schema_sha256"], - "YETO_CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256": contract["dynamic_tools_schema_sha256"], - "YETO_CODEX_REASONING_EFFORT": contract["reasoning_effort"], - "YETO_CODEX_BACKEND_MAX_TOKENS": str(backend["max_tokens"]), - "YETO_CODEX_BACKEND_REASONING_EFFORT": backend["reasoning_effort"], - "YETO_CODEX_BACKEND_THINKING": backend["thinking"]["type"], - "YETO_CODEX_CHAT_TEMPLATE": backend["chat_template"], - "YETO_CODEX_CHAT_TEMPLATE_KWARGS": _canonical_json(backend["chat_template_kwargs"]), - "YETO_CODEX_TITO_ALLOWED_APPEND_ROLES": "tool,user", - } - env.update({str(k): str(v) for k, v in contract.get("openenv_identity_env", {}).items()}) - env["YETO_CODEX_HARNESS_CONTRACT_SHA256"] = _plan_digest(contract) - return env - - -def rl_gae_variant(args) -> str | None: - """``advantage.gae_variant`` of the prepared RL algorithm spec (None if absent).""" - raw = getattr(args, "rl_algorithm_spec_json", None) - if getattr(args, "training_mode", "sft") != "rl" or not raw: - return None - try: - payload = json.loads(raw) - except (TypeError, ValueError): - return None - value = ((payload or {}).get("advantage") or {}).get("gae_variant") - return None if value is None else str(value) - - -def compactionrl_launch_env(args, environ, harness_compacts: bool) -> dict[str, str]: - """Rollout-worker env for CompactionRL (design D8; progress "S13 Codex 压缩接线"). - - A ``cross_segment_per_sample`` spec switches ``YETO_CODEX_COMPACTIONRL=1`` on - (and forwards ``YETO_CODEX_COMPACTIONRL_T_COMP`` from the launching - environment when set); the island preflight forwards every ``YETO_CODEX_*`` - to the Ray rollout workers. Any other spec returns {} (env unchanged). - Raises when spec, harness and the launching environment disagree - (``compactionrl.check_rollout_compaction``).""" - from ..rl.algos import compactionrl as crl - - variant = rl_gae_variant(args) - raw_switch = environ.get(crl.COMPACTION_SWITCH_ENV) - t_comp = environ.get(crl.COMPACTION_T_COMP_ENV) - if variant is None and not raw_switch and not t_comp: - return {} - try: - explicit = crl.compaction_switch(raw_switch) - wanted = variant == crl.COMPACTED_GAE_VARIANT - if raw_switch not in (None, "") and explicit != wanted: - if wanted: - raise ValueError( - f"{crl.COMPACTION_SWITCH_ENV}={raw_switch!r} in the launching environment " - f"contradicts the CompactionRL spec ({crl.COMPACTED_GAE_VARIANT})" - ) - crl.check_rollout_compaction( - variant, wanted or explicit, harness_compacts=harness_compacts, t_comp=t_comp - ) - except ValueError as error: - raise ValueError(f"CompactionRL launch: {error}") from error - if not wanted: - return {} - generate = getattr(args, "custom_generate_function_path", None) - if getattr(args, "custom_agent_function_path", None) and generate != crl.CODEX_OPENENV_GENERATE: - # s19-compaction-g1-20261010b: with upstream agentic_tool_call.generate the - # segment sessions are never collected, samples carry no tokens_after and - # the trainer fails at round 0 (cross_segment_per_sample needs tokens_after). - raise ValueError( - f"CompactionRL launch: {crl.COMPACTED_GAE_VARIANT} needs " - f"--custom-generate-function-path {crl.CODEX_OPENENV_GENERATE} (got {generate!r})") - env = {crl.COMPACTION_SWITCH_ENV: "1"} - if t_comp: - env[crl.COMPACTION_T_COMP_ENV] = str(t_comp).strip() - return env - - -def codex_harness_launch(args, environ=None) -> tuple[str, dict[str, str], dict[str, str]] | None: - """(learner flags, container envs, file_mounts) for a signed Codex agent on - the ports engine, or None when the run is not a Codex run. - - Fails closed: no bundle dir, a bundle that does not match the pins, no - environment provider, or trainable compaction switched on in the launching - environment (R-D5a) all raise before any cloud resource is requested.""" - from ..rl import CODEX_OPENENV_AGENT, SIGNED_CODEX_AGENTS - from ..rl.harness.codex.preflight import HARNESS_PREFLIGHT_SPEC, _FALSE - - environ = os.environ if environ is None else environ - custom_agent = getattr(args, "custom_agent_function_path", None) - from ..rl.synthetic_segments import SYNTHETIC_SEGMENTS_REWARD - - # TEST ONLY (task 6.4 G1): the synthetic two-segment reward stands in for a - # compacting harness; it is accepted only with cross_segment_per_sample. - synthetic = getattr(args, "reward_function", None) == SYNTHETIC_SEGMENTS_REWARD - if synthetic and rl_gae_variant(args) != "cross_segment_per_sample": - raise ValueError(f"{SYNTHETIC_SEGMENTS_REWARD} is only valid with gae_variant cross_segment_per_sample") - compaction_env = compactionrl_launch_env( - args, environ, custom_agent == CODEX_OPENENV_AGENT or (synthetic and custom_agent is None)) - if custom_agent not in SIGNED_CODEX_AGENTS: - return None - bundle_dir = getattr(args, "codex_bundle_dir", None) or environ.get(CODEX_BUNDLE_DIR_ENV) - if not bundle_dir: - raise ValueError( - f"the signed Codex harness on the ports engine needs {CODEX_BUNDLE_DIR_ENV} " - "(a bundle from scripts/fetch_codex_bundle.py) to mount at " - f"{CODEX_CONTAINER_DIR}" - ) - if environ.get(CODEX_COMPACTION_ENV, "").strip().lower() not in _FALSE: - raise ValueError( - f"{CODEX_COMPACTION_ENV} is set in the launching environment; the fork pin " - "session server has no trainable compaction (R-D5a)" - ) - contract = codex_bundle_contract(args, bundle_dir) - flags = f" --codex-harness-contract {shlex.quote(_json_compact(contract))}" - envs = codex_container_env(contract) - if custom_agent == CODEX_OPENENV_AGENT: - envs[HARNESS_PREFLIGHT_ENV] = environ.get(HARNESS_PREFLIGHT_ENV) or HARNESS_PREFLIGHT_SPEC - if not environ.get(HARNESS_ENVIRONMENT_PROVIDER_ENV): - raise ValueError( - f"{HARNESS_ENVIRONMENT_PROVIDER_ENV}=module:callable is required: the " - "Codex Terminal-Bench preflight fails closed without an environment provider" - ) - for name in HARNESS_PASSTHROUGH_ENV: - if environ.get(name): - envs[name] = environ[name] - for name, value in environ.items(): - if name.startswith(HARNESS_PASSTHROUGH_ENV_PREFIXES) and value: - envs[name] = value - envs.update(compaction_env) # {} unless the spec is CompactionRL - mounts = {CODEX_CONTAINER_DIR: str(Path(bundle_dir).expanduser().resolve())} - return flags, envs, mounts - - -HEAD_CODEX_BUNDLE_PATH = "~/yeto-codex-bundle" - - -def codex_head_staging(args, environ=None) -> tuple[dict[str, str], dict[str, str]]: - """(head file_mounts, head-job secret envs) for a signed Codex run under - ``--controller head`` (S17 M1). - - The head VM replays the launch and builds every island itself, so it needs - what ``codex_harness_launch`` reads on the submitter: the bundle directory - (staged at ``HEAD_CODEX_BUNDLE_PATH``; the caller points - ``args.codex_bundle_dir`` there) and the harness env (provider, TB2 knobs, - reward HMAC key, Modal token). The env goes as sky secrets (redacted), - never as plain envs. Returns ({}, {}) for a non-Codex run; fails closed - exactly like ``codex_harness_launch`` (attests the local bundle first). - """ - environ = os.environ if environ is None else environ - if codex_harness_launch(args, environ) is None: - return {}, {} - bundle_dir = getattr(args, "codex_bundle_dir", None) or environ.get(CODEX_BUNDLE_DIR_ENV) - mounts = {HEAD_CODEX_BUNDLE_PATH: str(Path(bundle_dir).expanduser().resolve())} - secrets = {name: environ[name] for name in (HARNESS_PREFLIGHT_ENV, *HARNESS_PASSTHROUGH_ENV) - if environ.get(name)} - secrets.update({name: value for name, value in environ.items() - if name.startswith(HARNESS_PASSTHROUGH_ENV_PREFIXES) and value}) - return mounts, secrets - - -def _json_compact(value) -> str: - return json.dumps(value, sort_keys=True, separators=(",", ":")) diff --git a/yeto/launch/harness_env.py b/yeto/launch/harness_env.py new file mode 100644 index 000000000..6821a9478 --- /dev/null +++ b/yeto/launch/harness_env.py @@ -0,0 +1,90 @@ +"""Launch-side island environment: the Modal sandbox secret and the CompactionRL switch.""" + +from __future__ import annotations + +import json +import os +import re + + +def modal_sandbox_secret_name(args, spec) -> str | None: + """--modal-sandbox-secret NAME, or None when unset. Only Modal islands can + attach a named Modal Secret: any other cloud is refused, never ignored.""" + name = getattr(args, "modal_sandbox_secret", None) + if not name: + return None + if getattr(spec, "cloud", None) != "modal": + raise ValueError( + f"--modal-sandbox-secret attaches a Modal Secret to a Modal island; " + f"island cloud {getattr(spec, 'cloud', None)!r} cannot use it") + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,63}", str(name)): + raise ValueError(f"--modal-sandbox-secret {name!r} is not a Modal Secret name") + return str(name) + + +def rl_gae_variant(args) -> str | None: + """``advantage.gae_variant`` of the prepared RL algorithm spec (None if absent).""" + raw = getattr(args, "rl_algorithm_spec_json", None) + if getattr(args, "training_mode", "sft") != "rl" or not raw: + return None + try: + payload = json.loads(raw) + except (TypeError, ValueError): + return None + value = ((payload or {}).get("advantage") or {}).get("gae_variant") + return None if value is None else str(value) + + +def compactionrl_launch_env(args, environ, harness_compacts: bool) -> dict[str, str]: + """Rollout-worker env for CompactionRL (design D8; progress "S13 Codex 压缩接线"). + + A ``cross_segment_per_sample`` spec switches ``YETO_CODEX_COMPACTIONRL=1`` on + (and forwards ``YETO_CODEX_COMPACTIONRL_T_COMP`` from the launching + environment when set). Any other spec returns {} (env unchanged). + Raises when spec, harness and the launching environment disagree + (``compactionrl.check_rollout_compaction``).""" + from ..rl.algos import compactionrl as crl + + variant = rl_gae_variant(args) + raw_switch = environ.get(crl.COMPACTION_SWITCH_ENV) + t_comp = environ.get(crl.COMPACTION_T_COMP_ENV) + if variant is None and not raw_switch and not t_comp: + return {} + try: + explicit = crl.compaction_switch(raw_switch) + wanted = variant == crl.COMPACTED_GAE_VARIANT + if raw_switch not in (None, "") and explicit != wanted: + if wanted: + raise ValueError( + f"{crl.COMPACTION_SWITCH_ENV}={raw_switch!r} in the launching environment " + f"contradicts the CompactionRL spec ({crl.COMPACTED_GAE_VARIANT})" + ) + crl.check_rollout_compaction( + variant, wanted or explicit, harness_compacts=harness_compacts, t_comp=t_comp + ) + except ValueError as error: + raise ValueError(f"CompactionRL launch: {error}") from error + if not wanted: + return {} + env = {crl.COMPACTION_SWITCH_ENV: "1"} + if t_comp: + env[crl.COMPACTION_T_COMP_ENV] = str(t_comp).strip() + return env + + +def check_compaction_launch(args, environ=None) -> None: + """Refuse a CompactionRL launch that cannot work, before any cloud resource. + + No agent harness compacts rollouts (the Codex OpenEnv bridge was removed in + S26 E0.5). TEST ONLY (task 6.4 G1): the synthetic two-segment reward stands + in for a compacting harness; it is accepted only with + ``cross_segment_per_sample`` and without a custom agent. + """ + from ..rl.synthetic_segments import SYNTHETIC_SEGMENTS_REWARD + + environ = os.environ if environ is None else environ + synthetic = getattr(args, "reward_function", None) == SYNTHETIC_SEGMENTS_REWARD + if synthetic and rl_gae_variant(args) != "cross_segment_per_sample": + raise ValueError(f"{SYNTHETIC_SEGMENTS_REWARD} is only valid with gae_variant cross_segment_per_sample") + harness_compacts = synthetic and getattr(args, "custom_agent_function_path", None) is None + compactionrl_launch_env(args, environ, harness_compacts) diff --git a/yeto/launch/modal_islands.py b/yeto/launch/modal_islands.py index a7f12af6f..b96808845 100644 --- a/yeto/launch/modal_islands.py +++ b/yeto/launch/modal_islands.py @@ -21,7 +21,7 @@ _rl_checkpoint_storage_name, droppable_rejoins_in_new_container, modal_checkpoint_store, resumes_in_new_container, ) -from .codex import CODEX_CONTAINER_DIR, modal_sandbox_secret_name +from .harness_env import modal_sandbox_secret_name from .fleet import TEARDOWN_UNVERIFIED_EXIT @@ -74,16 +74,13 @@ def build_modal_island_config(args, spec: ClusterSpec, learner_id: int, task, sy if rl and island_keeps_spot_volume(args): volume_name = _rl_checkpoint_storage_name(args.cluster_prefix, learner_id) volume_mount = _rl_checkpoint_mount(args.rl_completed_groups_path).replace("~", "/root", 1) - # Codex run bundle: the sky task mounts it at CODEX_CONTAINER_DIR; Modal - # mounts the same directory through the image (add_local_dir). + # Every sky file_mount (local prompt file, initial adapter, HF token file) + # is mounted into the container the same way; `~` is /root in Modal. all_mounts = dict(getattr(task, "file_mounts", None) or {}) - codex_dir = all_mounts.get(CODEX_CONTAINER_DIR) - # Every other sky file_mount (local prompt file, initial adapter, HF token - # file) is mounted into the container the same way; `~` is /root in Modal. extra_mounts = { ("/root" + target[1:] if target.startswith("~") else target): str(source) for target, source in all_mounts.items() - if target != CODEX_CONTAINER_DIR and not str(source).startswith(("s3://", "gs://", "r2://")) + if not str(source).startswith(("s3://", "gs://", "r2://")) } requirements: tuple[str, ...] = () if not rl: @@ -111,8 +108,6 @@ def build_modal_island_config(args, spec: ClusterSpec, learner_id: int, task, sy pip_requirements=requirements, volume_name=volume_name, volume_mount=volume_mount, - codex_dir=codex_dir, - codex_mount=CODEX_CONTAINER_DIR if codex_dir else None, extra_mounts=extra_mounts, checkpoint_store_volume_name=None if store_volume is None else store_volume[0], checkpoint_store_mount=None if store_volume is None else MODAL_CHECKPOINT_STORE_MOUNT, diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 0c21669b9..53df042ee 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -343,16 +343,6 @@ def _prepare_rl_args( _rl_miles_function(args.custom_generate_function_path) custom_agent = getattr(args, "custom_agent_function_path", None) _rl_miles_function(custom_agent, "--custom-agent-function-path") - from ..rl import ( - SECRLENV_AGENTS, - SECRLENV_GENERATE, - SECRLENV_GROUP_FILTER, - SECRLENV_INFRASTRUCTURE_REPLACEMENTS, - SECRLENV_REWARD, - SECRLENV_ZERO_VARIANCE_REPLACEMENTS, - SIGNED_CODEX_AGENTS, - ) - if custom_agent is not None: if args.custom_generate_function_path is None: raise ValueError( @@ -367,120 +357,23 @@ def _prepare_rl_args( raise ValueError( "--custom-agent-function-path requires --tito-model" ) - if ( - args.apply_chat_template_kwargs is not None - and custom_agent not in SIGNED_CODEX_AGENTS - ): + if args.apply_chat_template_kwargs is not None: raise ValueError( "agentic session rollouts preserve raw messages and do not accept " "--apply-chat-template-kwargs" ) - codex_reasoning_effort = getattr(args, "codex_reasoning_effort", None) - codex_backend_profile = getattr(args, "codex_backend_profile", None) - if custom_agent in SIGNED_CODEX_AGENTS: - from ..rl.codex_backend import ( - stock_codex_backend_profile, - validate_stock_codex_fields, - ) - - codex_backend_profile = codex_backend_profile or str(args.tito_model) - args.codex_backend_profile = codex_backend_profile - codex_profile = stock_codex_backend_profile(codex_backend_profile) - codex_chat_template_kwargs = codex_profile["chat_template_kwargs"] - if args.apply_chat_template_kwargs is None: - args.apply_chat_template_kwargs = codex_chat_template_kwargs - validate_stock_codex_fields( - tito_model=str(args.tito_model), - codex_backend_profile=codex_backend_profile, - rl_model_recipe=getattr(args, "rl_model_recipe", "generic"), - model=str(args.model), - model_revision=str(args.model_revision), - rollout_model=getattr(args, "rollout_model", None), - rollout_model_revision=getattr(args, "rollout_model_revision", None), - apply_chat_template_kwargs=args.apply_chat_template_kwargs, - tito_allowed_append_roles=list( - getattr(args, "tito_allowed_append_roles", None) or () - ), - codex_reasoning_effort=codex_reasoning_effort, - lora_targets=str(args.lora_targets), - expert_full_count=int(getattr(args, "expert_full_count", 0)), - lora_expert_rank=int(getattr(args, "rl_lora_expert_rank", None) or 0), - ) - elif codex_reasoning_effort is not None or codex_backend_profile is not None: - raise ValueError( - "--codex-reasoning-effort/--codex-backend-profile requires a signed Codex agent" - ) - if getattr(args, "codex_exclude_reasoning_from_loss", False): - if custom_agent not in SIGNED_CODEX_AGENTS: - raise ValueError( - "--codex-exclude-reasoning-from-loss requires a signed Codex agent" - ) - from ..rl.codex_backend import stock_codex_reasoning_markers - - if stock_codex_reasoning_markers(codex_backend_profile) is None: - raise ValueError( - f"Codex profile {codex_backend_profile!r} declares no reasoning " - "markers; --codex-exclude-reasoning-from-loss is unavailable" - ) - secrlenv_infrastructure_replacements = getattr( - args, "secrlenv_max_infrastructure_replacements", None - ) - if custom_agent in SECRLENV_AGENTS: - if args.custom_generate_function_path != SECRLENV_GENERATE: - raise ValueError( - "the SecRLEnv agents require the signed SecRLEnv generate wrapper" - ) - if args.reward_function != SECRLENV_REWARD: - raise ValueError( - "the SecRLEnv agents require the signed SecRLEnv reward function" - ) - configured_filter = getattr(args, "dynamic_sampling_filter_path", None) - if configured_filter not in {None, SECRLENV_GROUP_FILTER}: - raise ValueError( - "the SecRLEnv agents require the exact signed group filter" - ) - zero_variance_replacements = getattr( - args, "dynamic_sampling_max_replacements", None - ) - if zero_variance_replacements not in { - None, - SECRLENV_ZERO_VARIANCE_REPLACEMENTS, - } or isinstance(zero_variance_replacements, bool): - raise ValueError( - "the SecRLEnv agents require zero variance replacements" - ) - if secrlenv_infrastructure_replacements not in { - None, - SECRLENV_INFRASTRUCTURE_REPLACEMENTS, - } or isinstance(secrlenv_infrastructure_replacements, bool): - raise ValueError( - "the SecRLEnv agents require exactly one infrastructure replacement" - ) - expected_oversampling = args.rollout_batch_size + 1 - if args.over_sampling_batch_size == args.rollout_batch_size: - args.over_sampling_batch_size = expected_oversampling - elif args.over_sampling_batch_size != expected_oversampling: - # agentic-rollout-utilization 1.1 (10-09 ruling): the SecRLEnv +1 - # is a reserved same-task retry slot (_SecRLEnvRetryDataSource), not - # over-sampling; over-sample/cut-off applies to the other agents. - raise ValueError( - "the SecRLEnv agents require oversampling equal to the training " - "batch plus one (a reserved same-task retry slot; over-sampling " - "with cut-off is not available to SecRLEnv agents)" - ) - args.dynamic_sampling_filter_path = SECRLENV_GROUP_FILTER - args.dynamic_sampling_max_replacements = ( - SECRLENV_ZERO_VARIANCE_REPLACEMENTS - ) - args.secrlenv_max_infrastructure_replacements = ( - SECRLENV_INFRASTRUCTURE_REPLACEMENTS - ) - elif secrlenv_infrastructure_replacements is not None: - raise ValueError( - "--secrlenv-max-infrastructure-replacements requires a signed " - "SecRLEnv agent" - ) + # The signed stock-Codex and SecRLEnv agents these flags configured were + # removed in S26 E0.5 (tag s26-pre-harness-removal); refuse them. + for flag, attr in ( + ("--codex-reasoning-effort", "codex_reasoning_effort"), + ("--codex-backend-profile", "codex_backend_profile"), + ("--codex-exclude-reasoning-from-loss", "codex_exclude_reasoning_from_loss"), + ("--secrlenv-max-infrastructure-replacements", "secrlenv_max_infrastructure_replacements"), + ): + if getattr(args, attr, None) not in (None, False): + raise ValueError(f"{flag} requires a signed Codex agent; the signed Codex and " + "SecRLEnv agents were removed (S26 E0.5)") agent_max_seq_len = getattr(args, "agent_max_seq_len", None) if agent_max_seq_len is not None: if custom_agent is None: @@ -517,16 +410,13 @@ def _prepare_rl_args( bounded_filter = "yeto.rl.filters.bounded_nonzero_reward_std" if dynamic_filter == stock_filter: args.dynamic_sampling_filter_path = bounded_filter - elif dynamic_filter not in { - bounded_filter, - SECRLENV_GROUP_FILTER, - }: + elif dynamic_filter != bounded_filter: raise ValueError( "--dynamic-sampling-max-replacements is only supported with " "an approved bounded nonzero-variance filter" ) - # After the agent contract binds its filter (SecRLEnv) and the over-sampling - # count: the expected spec hash must be the one the island computes. + # After the over-sampling count is final: the expected spec hash must be + # the one the island computes. _prepare_ports_algorithm(args) if not args.use_session_server and ( args.session_server_ip is not None diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 1185f4c9a..ce862f1dd 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -38,11 +38,7 @@ from .in_vm_docker import ( DOCKER_LOGIN_UNSET, IN_VM_DOCKER_CLOUDS, in_vm_docker_run, in_vm_docker_setup, ) -from .codex import ( - HARNESS_ENVIRONMENT_PROVIDER_ENV, MODAL_CLIENT_SETUP, MODAL_SANDBOX_PROVIDERS, - SANDBOX_TOKEN_IN_ISLAND_CHECK, codex_harness_launch, modal_sandbox_secret_name, - require_sandbox_modal_token, -) +from .harness_env import check_compaction_launch def submitted_unix() -> float: @@ -96,7 +92,7 @@ def make_miles_island_task( ): raise ValueError("RL task requires prepared source and reward provenance") island_data, island_data_sha256, data_mounts = island_task_data(args, learner_id) - codex_launch = codex_harness_launch(args) # None unless a signed Codex agent + check_compaction_launch(args) # --rl-island-use-gpus-per-node: island rules see M GPUs/node; resources stay physical. island = rl_island_spec(args, spec) @@ -195,11 +191,6 @@ def make_miles_island_task( " --dynamic-sampling-max-replacements " f"{args.dynamic_sampling_max_replacements}" ) - if getattr(args, "secrlenv_max_infrastructure_replacements", None) is not None: - flags += ( - " --secrlenv-max-infrastructure-replacements " - f"{args.secrlenv_max_infrastructure_replacements}" - ) if args.rl_offload_train: flags += " --rl-offload-train" algorithm_prelude, algorithm_flags = _ports_algorithm_flags(args) @@ -254,20 +245,6 @@ def make_miles_island_task( flags += f" --session-server-workers {int(args.session_server_workers)}" if args.tito_model: flags += f" --tito-model {shlex.quote(args.tito_model)}" - if getattr(args, "codex_backend_profile", None): - flags += ( - " --codex-backend-profile " - f"{shlex.quote(args.codex_backend_profile)}" - ) - if getattr(args, "codex_exclude_reasoning_from_loss", False): - flags += " --codex-exclude-reasoning-from-loss" - if getattr(args, "codex_reasoning_effort", None): - # The learner's _preflight_codex_harness re-validates the stock - # profile (xhigh) and fails closed without it. - flags += ( - " --codex-reasoning-effort " - f"{shlex.quote(args.codex_reasoning_effort)}" - ) if getattr(args, "tito_allowed_append_roles", None): roles = " ".join( shlex.quote(role) for role in args.tito_allowed_append_roles @@ -319,10 +296,6 @@ def make_miles_island_task( file_mounts[HF_TOKEN_PATH] = local_token if rl_sao_streaming(args): file_mounts.update(sao_contract_mounts(args, learner_id)) - if codex_launch is not None: - codex_flags, codex_envs, codex_mounts = codex_launch - flags += codex_flags - file_mounts.update(codex_mounts) envs = { "SYNCER_ADDR": syncer_addr, **({"CRITIC_SYNCER_ADDR": critic_syncer_address(syncer_addr)} @@ -378,8 +351,6 @@ def make_miles_island_task( for name in ("CYBERGYM_REWARD_SCHEME", "CYBERGYM_REWARD_VIEW"): if os.environ.get(name): envs[name] = os.environ[name] - if codex_launch is not None: - envs.update(codex_envs) envs.update(elastic_debug_delay_env(args)) # 0.25 ({} unless the test switch is given) envs.update(elastic_debug_pause_env(args)) # link-pause test switch ({} by default) if getattr(args, "wandb", False): @@ -401,15 +372,6 @@ def make_miles_island_task( store_envs = store_object_envs(args, spec.cloud, spec.num_nodes) if store_envs: # the store's object API (yeto.cloud.s3_store) needs boto3 in the image setup_steps.append(STORE_CLIENT_SETUP) - sandbox_secret = modal_sandbox_secret_name(args, spec) - if codex_launch is not None and envs.get(HARNESS_ENVIRONMENT_PROVIDER_ENV) in MODAL_SANDBOX_PROVIDERS: - if sandbox_secret is None: - require_sandbox_modal_token(envs) - else: - # The token comes from the named Modal Secret at container start; - # the island re-checks it before anything else runs. - setup_steps.append(SANDBOX_TOKEN_IN_ISLAND_CHECK) - setup_steps.append(MODAL_CLIENT_SETUP) store_fs = model_store_filesystem(getattr(args, "model_store", None), spec.cloud, spec.region) store_env = "" if (spec.cloud == "modal" and getattr(args, "modal_model_volume", None) diff --git a/yeto/launcher.py b/yeto/launcher.py index 0abd2f4b5..f104fef1e 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -125,13 +125,8 @@ from .launch.in_vm_docker import ( # noqa: F401 DOCKER_LOGIN_UNSET, IN_VM_DOCKER_CLOUDS, in_vm_docker_image, in_vm_docker_run, ) -from .launch.codex import ( # noqa: F401 - CODEX_BUNDLE_DIR_ENV, CODEX_COMPACTION_ENV, CODEX_CONTAINER_DIR, - HARNESS_ENVIRONMENT_PROVIDER_ENV, HARNESS_PASSTHROUGH_ENV, HARNESS_PREFLIGHT_ENV, - HEAD_CODEX_BUNDLE_PATH, MODAL_CLIENT_SETUP, MODAL_SANDBOX_PROVIDER, MODAL_SANDBOX_PROVIDERS, - PREBAKED_MODAL_SANDBOX_PROVIDER, SANDBOX_MODAL_TOKEN_ENVS, SANDBOX_TOKEN_IN_ISLAND_CHECK, - codex_bundle_contract, codex_harness_launch, codex_head_staging, modal_sandbox_secret_name, - require_sandbox_modal_token, +from .launch.harness_env import ( # noqa: F401 + check_compaction_launch, compactionrl_launch_env, modal_sandbox_secret_name, ) from .launch.rl_island_task import _worker_node_ip_flag # noqa: F401 from .launch.diffusion_sample import make_diffusion_sample_task, run_diffusion_sample # noqa: F401 diff --git a/yeto/modal_runner.py b/yeto/modal_runner.py index b25be3c4f..ce3d6fefc 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -216,12 +216,6 @@ class ModalIslandConfig: # links CONTAINER_WORKDIR to it). workdir: str | None = str(REPO_ROOT) python_version: str = "3.12" - # Stock Codex run bundle (ports engine, signed Codex agent): a local dir - # holding the attested binary / package manifest / app-server schema, - # mounted read-only at ``codex_mount`` (= ``/opt/yeto/codex``, the - # CODEX_CONTAINER_BINARY_PATH parent) the same way sky file_mounts does. - codex_dir: str | None = None - codex_mount: str | None = None # Other sky file_mounts (container path -> local file or dir), mounted # read-only at start-up like the workdir. extra_mounts: dict[str, str] = field(default_factory=dict) @@ -294,10 +288,6 @@ def validate(self) -> None: raise ValueError("model_volume_name and model_volume_mount go together") if self.model_volume_name and self.model_volume_name in (self.volume_name, self.tape_volume_name): raise ValueError("the model volume must differ from the checkpoint/tape volumes") - if (self.codex_dir is None) != (self.codex_mount is None): - raise ValueError("codex_dir and codex_mount go together") - if self.codex_dir is not None and not os.path.isdir(self.codex_dir): - raise ValueError(f"codex_dir {self.codex_dir} is not a directory") for target, source in self.extra_mounts.items(): if not target.startswith("/"): raise ValueError(f"extra_mounts target {target} must be absolute") @@ -1022,9 +1012,6 @@ def build_image(self, cfg: ModalIslandConfig): copy=False, ignore=MODAL_WORKDIR_IGNORE, ) - if cfg.codex_dir and cfg.codex_mount: - # Codex run bundle (sky: file_mounts[codex_mount] = codex_dir). - image = image.add_local_dir(cfg.codex_dir, cfg.codex_mount, copy=False) for target, source in sorted(cfg.extra_mounts.items()): source = os.path.expanduser(source) if os.path.isdir(source): diff --git a/yeto/rl/__init__.py b/yeto/rl/__init__.py index f83714f06..4c920a88e 100644 --- a/yeto/rl/__init__.py +++ b/yeto/rl/__init__.py @@ -35,96 +35,3 @@ def __getattr__(name: str): return getattr(pins, name) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") -SECRLENV_AGENT_PATH = "yeto/rl/harness/codex/agent.py" -SECRLENV_AGENT_SHA256 = ( - "84c6c708cba2142e3d043ae047c90f529b90b953b739261466b8a6e28b2c7411" -) -SECRLENV_AGENT = "yeto.rl.harness.codex.agent.run" -SECRLENV_REWARD = "yeto.rl.harness.codex.reward:reward_func" -SECRLENV_GROUP_FILTER = "yeto.rl.harness.codex.reward.check_group" -# Moved to the Miles adapter (yeto-framework-decoupling 5.7); was yeto.rl.harness.codex.generate. -SECRLENV_GENERATE = "yeto.rl.adapters.miles.harness_glue.codex_generate.generate" -SECRLENV_GENERATE_SHA256 = ( - "1630dc8a3aae3cbbd4708a81c9803920791ee58d6725c3fcaa4736373f8f9e48" -) -SECRLENV_ZERO_VARIANCE_REPLACEMENTS = 0 -SECRLENV_INFRASTRUCTURE_REPLACEMENTS = 1 - -# Stock Codex is part of the signed Yeto security-environment harness. These -# pins identify the official Linux artifact, not the controller's host binary. -CODEX_HARNESS_AGENT = "yeto.rl.harness.codex.codex_harness_agent.run" -CODEX_HARNESS_AGENT_PATH = "yeto/rl/harness/codex/codex_harness_agent.py" -CODEX_HARNESS_AGENT_SHA256 = ( - "459855e8845ab3b3cc1a4e2f47e5a9e8cc6cad79197208e43aa1d822a7c300ad" -) -CODEX_BASE_INSTRUCTIONS_SHA256 = ( - "1c183656ca1319142cba9e76baa199b7ab59f770a51a76660622a087e74ba846" -) -# 2026-10-07 (S15 stage 2): signed TB2 system prompt (codex_harness_agent. -# TB2_BASE_INSTRUCTIONS); the legacy CODEX_BASE_INSTRUCTIONS_SHA256 is unchanged. -CODEX_TB2_BASE_INSTRUCTIONS_SHA256 = ( - "707e144f4c41dd791f5875ea6693be4738d21e7dec73a3ea09ccbf92106bdeb3" -) -CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256 = ( - "868dbbff9fe2f5a57573826cae1ae1f4ceac04eff8689a522d9af7ef1b589c5a" -) -CODEX_SUBMIT_TOOL_SCHEMA_SHA256 = ( - "162980cf1de2346e6a246a739c10a31d0f5bd30c62b27c080887b298ecad1a6f" -) -CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256 = ( - "06142f7664a668c11149b9410af6438423654ac7ee85fa382226bc7fbbf101af" -) -CODEX_CLI_VERSION = "codex-cli 0.145.0" -CODEX_NPM_PACKAGE = "@openai/codex@0.145.0-linux-x64" -CODEX_LINUX_TARGET = "x86_64-unknown-linux-musl" -CODEX_LINUX_BINARY_SHA256 = ( - "a2a05dafaa1acb002a45eaec0a462de5b13694fcfcd7bc43305f14781ce7be14" -) -CODEX_LINUX_BINARY_SIZE_BYTES = 310_730_800 -CODEX_NPM_TARBALL_SHA256 = ( - "11239480f8e3efd1430f23bbe91c1a397856b8bbe6185ccbaee2382d25e03df2" -) -CODEX_PACKAGE_MANIFEST_SHA256 = ( - "8da5349aa5a4242f5e11c5ca8ff4a16d8f9f912cb8accebea4def94edbf30aee" -) -CODEX_APP_SERVER_PROTOCOL_REVISION = "v2" -CODEX_APP_SERVER_SCHEMA_SHA256 = ( - "a88d865c3ca41fc63672baf423e28b3bfdb85b989e1c95e6e268932daaac91a0" -) -CODEX_CONTAINER_BINARY_PATH = "/opt/yeto/codex/codex-x86_64-unknown-linux-musl" -CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH = ( - "/opt/yeto/codex/codex_app_server_protocol.v2.schemas.json" -) - -# Terminal-Bench uses a thin isolated-process wrapper around the same attested -# stock Codex runtime. It is intentionally not a SecRLEnv agent: its reward, -# retry, and cleanup evidence contracts are Terminal-Bench-specific. -CODEX_OPENENV_AGENT = "yeto.rl.harness.codex.codex_openenv_subprocess_agent_function.run" -CODEX_OPENENV_AGENT_MODULES = ( - "codex_openenv_subprocess_agent_function.py", - "codex_openenv_agent_worker.py", - "codex_openenv_agent_function.py", - # CompactionRL bridge (opt-in, YETO_CODEX_COMPACTIONRL); imported by - # codex_openenv_agent_function. Name list only: no per-module hash pin. - "compaction_bridge.py", -) -CODEX_OPENENV_IDENTITY_ENV = { - "YETO_CODEX_OPENENV_BACKEND_PROFILE": "qwen35_08b", - "YETO_CODEX_OPENENV_MODEL_ID": "Qwen/Qwen3.5-0.8B", - "YETO_CODEX_OPENENV_MODEL_REVISION": ("2fc06364715b967f1860aea9cf38778875588b17"), - "YETO_CODEX_OPENENV_BASE_INSTRUCTIONS_SHA256": ( - "1c183656ca1319142cba9e76baa199b7ab59f770a51a76660622a087e74ba846" - ), - "YETO_CODEX_OPENENV_TERMINAL_EXEC_TOOL_SCHEMA_SHA256": ( - "868dbbff9fe2f5a57573826cae1ae1f4ceac04eff8689a522d9af7ef1b589c5a" - ), - "YETO_CODEX_OPENENV_SUBMIT_TOOL_SCHEMA_SHA256": ( - "162980cf1de2346e6a246a739c10a31d0f5bd30c62b27c080887b298ecad1a6f" - ), - "YETO_CODEX_OPENENV_DYNAMIC_TOOLS_SCHEMA_SHA256": ( - "06142f7664a668c11149b9410af6438423654ac7ee85fa382226bc7fbbf101af" - ), -} -SIGNED_CODEX_AGENTS = frozenset((CODEX_HARNESS_AGENT, CODEX_OPENENV_AGENT)) - -SECRLENV_AGENTS = frozenset((SECRLENV_AGENT, CODEX_HARNESS_AGENT)) diff --git a/yeto/rl/adapters/miles/binding.py b/yeto/rl/adapters/miles/binding.py index adc7b90d6..f638d0fb0 100644 --- a/yeto/rl/adapters/miles/binding.py +++ b/yeto/rl/adapters/miles/binding.py @@ -16,8 +16,7 @@ 3. call -- one fixed CPU input through the mapped loading path gives a result equal to the pre-rename path (skipped, and reported as skipped, when the implementation cannot be imported here, e.g. a ``miles.*`` built-in on a - machine without Miles; not run for a filter that needs environment - evidence the fixed input cannot carry, see ``EVIDENCE_FILTERS``). + machine without Miles). The result is a JSON-able dict for the runtime manifest. """ @@ -31,13 +30,10 @@ from types import SimpleNamespace from typing import Any -from yeto.rl import SECRLENV_GROUP_FILTER from yeto.rl.engine.algorithm import ( BOUNDED_NONZERO_STD_FILTER, LEGACY_FILTER_NAMES, - SECRLENV_GROUP_CHECK_FILTER, STOCK_NONZERO_STD_FILTER, - TBENCH_GROUP_CHECK_FILTER, ) BINDING_SCHEMA = "yeto-backend-binding-v1" @@ -47,15 +43,8 @@ BOUNDED_NONZERO_STD_FILTER: "yeto.rl.filters.bounded_nonzero_reward_std", STOCK_NONZERO_STD_FILTER: "miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std", - SECRLENV_GROUP_CHECK_FILTER: SECRLENV_GROUP_FILTER, - TBENCH_GROUP_CHECK_FILTER: "yeto.rl.harness.codex.tbench_reward.check_group", } -# Filters whose input must carry signed environment evidence (the SecRLEnv -# group check refuses any group without it), so the fixed CPU input of check -# ③ cannot exercise them; checks ① and ② still run. -EVIDENCE_FILTERS = frozenset({SECRLENV_GROUP_CHECK_FILTER, TBENCH_GROUP_CHECK_FILTER}) - class BindingError(RuntimeError): """A neutral name does not bind to the same implementation as before the rename.""" @@ -134,9 +123,6 @@ def check_filter(name: str, *, paths: Mapping[str, str] | None = None, if ident != ident_old: raise BindingError(f"{name!r} ② identity: {ident} != pre-rename {ident_old}") result["identity"] = ident - if name in EVIDENCE_FILTERS: - result["call"] = "skipped: needs environment evidence" - return result try: new_fn, old_fn = loader(mapped), loader(old) except ImportError as exc: diff --git a/yeto/rl/adapters/miles/carry_over.py b/yeto/rl/adapters/miles/carry_over.py index 4e4eca4a0..6905fda98 100644 --- a/yeto/rl/adapters/miles/carry_over.py +++ b/yeto/rl/adapters/miles/carry_over.py @@ -462,11 +462,11 @@ def trained_token_versions(sample: Any, current_version: int) -> list[int | None RESUME_STATS_ATTR = "rollout_resume_stats" # Miles fork, set at rollout start SUSPEND_EXPIRED_PREFIX = "CodexSuspendExpired" POLICY_AGE_EXCEEDED_PREFIX = "policy_age_exceeded" +# Sample metadata key of an infrastructure failure reason (set by the agent function). +INFRASTRUCTURE_KEY = "tbench_infrastructure_error" def _infrastructure_reason(sample: Any) -> str: - from yeto.rl.harness.codex.tbench_reward import INFRASTRUCTURE_KEY - meta = getattr(sample, "metadata", None) reason = meta.get(INFRASTRUCTURE_KEY) if isinstance(meta, dict) else None return str(reason) if reason else "" diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index dadbb341d..7b1b627eb 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -95,7 +95,7 @@ "--rollout-all-samples-process-path", "--rollout-sample-filter-path", "--rollout-function-path", - "--data-source-path", # SecRLEnv same-task retry (secrlenv_retry) + "--data-source-path", "--buffer-filter-path", "--fully-async", "--debug-train-only", @@ -189,48 +189,6 @@ def check(value, _config): return check -def _stock_codex_append_roles_or_none(reason: str) -> _Check: - """Append roles are decided by the ``--tito-model`` template upstream, so - the option is refused -- except for a signed stock Codex agent, whose - launcher-side profile check (``validate_stock_codex_fields``) requires the - roles to be stated and equal to the profile's. Those are accepted when they - match the profile and are NOT emitted to argv (the template decides).""" - - def check(value, config): - if value is None: - return None - agent = getattr(config, "agent", None) - try: - from yeto.rl import SIGNED_CODEX_AGENTS - from yeto.rl.codex_backend import ( - stock_codex_backend_profile, - stock_codex_profiles_for_tito_model, - ) - except ImportError: # pragma: no cover - defensive - return reason - if getattr(agent, "custom_agent_function_path", None) not in SIGNED_CODEX_AGENTS: - return reason - if not getattr(agent, "tito_model", None): - return reason - # ``tito_model`` is the Miles tokenizer family, not the profile name - # (``qwen35`` happened to be both; ``qwen4exp`` serves the Flash-Next - # profiles, rl-fn-codex-rollout 1.0): the roles are fixed by every - # profile of that family, which the launcher already matched to the - # declared ``--codex-backend-profile`` (``validate_stock_codex_fields``). - names = stock_codex_profiles_for_tito_model(str(agent.tito_model)) - if not names: - return reason - fixed = {tuple(stock_codex_backend_profile(n)["tito_allowed_append_roles"]) for n in names} - if len(fixed) != 1 or tuple(value) != next(iter(fixed)): - return ( - f"stock Codex profiles {list(names)} (tito-model {agent.tito_model!r}) fix the " - f"append roles to {sorted(list(r) for r in fixed)} (got {list(value)})" - ) - return None - - return check - - def _must_be_false(reason: str) -> _Check: def check(value, _config): return reason if value else None @@ -249,14 +207,7 @@ def check(value, _config): ) -# yeto's Codex wrapper calls AGENTIC_TOOL_CALL_GENERATE with the same args (so it -# reads the same agent flags) and turns CompactionRL segment sessions into one -# sample per segment; without it a compacted rollout trains as one sample. -from yeto.rl.algos.compactionrl import CODEX_OPENENV_GENERATE # noqa: E402 (single source) -from yeto.rl import SECRLENV_GENERATE # noqa: E402 (wraps AGENTIC_TOOL_CALL_GENERATE) -AGENT_FLAG_READERS = frozenset( - {AGENTIC_TOOL_CALL_GENERATE, CODEX_OPENENV_GENERATE, SECRLENV_GENERATE} -) +AGENT_FLAG_READERS = frozenset({AGENTIC_TOOL_CALL_GENERATE}) def _requires_agentic_generate(flag: str) -> _Check: @@ -452,7 +403,7 @@ def _check_sglang_tp(value, config): "agent.session_server_port": _ok, "agent.session_server_workers": _ok, "agent.tito_model": _ok, - "agent.tito_allowed_append_roles": _stock_codex_append_roles_or_none( + "agent.tito_allowed_append_roles": _must_be_none( "upstream Miles has no --tito-allowed-append-roles; the append roles are " "decided by the --tito-model template (allowed_append_roles)" ), @@ -526,22 +477,6 @@ def _algorithm_filter(config, algorithm: AlgorithmSpec) -> str | None: return algorithm.dynamic_sampling_filter -def _secrlenv_retry_argv(agent, batch, dynamic_filter: str | None, - max_policy_age: int) -> tuple[str, ...]: - """P0-1.7b.4-Q1: the SecRLEnv agents bind the same-task retry classes.""" - - from .secrlenv_retry import rollout_argv - - try: - return rollout_argv(custom_agent=agent.custom_agent_function_path, - dynamic_filter=dynamic_filter, - rollout_batch_size=batch.groups_per_round, - over_sampling_batch_size=batch.over_sampling_batch_size, - max_policy_age=max_policy_age) - except ValueError as exc: - raise UnmappedConfigError("agent.custom_agent_function_path", str(exc)) from None - - def placement_request(config) -> PlacementRequest: parallel = config.parallel trainer = parallel.actor_num_nodes * parallel.actor_num_gpus_per_node @@ -1075,7 +1010,6 @@ def translate_run_config( values.extend(("--custom-generate-function-path", agent.custom_generate_function_path)) if agent.custom_agent_function_path: # IR-1 (leaf policy proved agentic generate) values.extend(("--custom-agent-function-path", agent.custom_agent_function_path)) - values.extend(_secrlenv_retry_argv(agent, batch, dynamic_filter, max_policy_age)) if agent.agent_max_seq_len is not None: values.extend(("--max-seq-len", str(int(agent.agent_max_seq_len)))) if agent.use_session_server: diff --git a/yeto/rl/adapters/miles/entry.py b/yeto/rl/adapters/miles/entry.py index f05f8ac4e..b9a1683a0 100644 --- a/yeto/rl/adapters/miles/entry.py +++ b/yeto/rl/adapters/miles/entry.py @@ -28,8 +28,6 @@ from yeto.rl.engine.algorithm import ( BOUNDED_NONZERO_STD_FILTER, - SECRLENV_GROUP_CHECK_FILTER, - TBENCH_GROUP_CHECK_FILTER, STOCK_NONZERO_STD_FILTER, AlgorithmSpec, ) @@ -273,10 +271,7 @@ def miles_capabilities( runtime_fingerprint=fingerprint, parameter_layouts={"lora"}, placements={"colocated"}, - dynamic_sampling_filters={ - BOUNDED_NONZERO_STD_FILTER, STOCK_NONZERO_STD_FILTER, SECRLENV_GROUP_CHECK_FILTER, - TBENCH_GROUP_CHECK_FILTER, - }, + dynamic_sampling_filters={BOUNDED_NONZERO_STD_FILTER, STOCK_NONZERO_STD_FILTER}, execution_modes={"colocated-serial"}, execution=ExecutionCapabilities( critic=False, max_policy_staleness=0, rollout_logprobs=True @@ -432,21 +427,13 @@ def load_tool_wait_source(miles_args: Any, elastic: Any = None) -> Any: The elastic drain board when wired; the island's named board when the tool-wait workload generate is configured (it counts on that board); 0 (``TOOL_WAIT_NO_BOARD_STOCK``) for Miles' stock generate, which makes no - tool calls; the island's named board for the Codex OpenEnv agent, which - enters/exits it around every tool call (agentic-rollout-utilization 1.3); - None (unknown) for any other custom generate. + tool calls; None (unknown) for any other custom generate. """ - from yeto.rl import CODEX_OPENENV_AGENT - from .rollout import TOOL_WAIT_NO_BOARD_STOCK board = getattr(elastic, "tool_wait_board", None) if elastic is not None else None if board is not None: return board - if getattr(miles_args, "custom_agent_function_path", None) == CODEX_OPENENV_AGENT: - from .elastic_wiring import LazyBoardActor - - return LazyBoardActor(int(getattr(miles_args, "yeto_rl_learner_id", 0) or 0)) custom = getattr(miles_args, "custom_generate_function_path", None) if not custom: return TOOL_WAIT_NO_BOARD_STOCK @@ -531,10 +518,8 @@ def publish_member_cell(miles_args: Any, learner_id: int | None = None, environ: Source, in order: an explicit ``miles_args.yeto_rl_cell_id`` (INFRA may pre-set it), ``learner_id`` (the island number the launcher already assigns; ``island_id`` in the - ``rl_engine_selected`` event), ``miles_args.yeto_rl_learner_id``, else 0. The env - export reaches rollout-worker subprocesses through ``worker_runtime_env`` and - ``preflight.configure_rollout_worker``; ``preflight.resolve_member`` then yields - ``engine:`` (``engine:0`` on a single island). + ``rl_engine_selected`` event), ``miles_args.yeto_rl_learner_id``, else 0. Rollout-side + code reads the env export (``engine:``; ``engine:0`` on a single island). """ environ = os.environ if environ is None else environ cell = getattr(miles_args, "yeto_rl_cell_id", None) @@ -1280,7 +1265,7 @@ def selection_event( return event -def connect_island_ray(*, environ=None, ray_module=None, miles_args=None) -> str | None: +def connect_island_ray(*, environ=None, ray_module=None) -> str | None: """Connect the driver to the island's own Ray and pin every actor to it. A SkyPilot machine runs two Ray instances: the island's (6379, started by @@ -1355,12 +1340,6 @@ def connect_island_ray(*, environ=None, ray_module=None, miles_args=None) -> str # --rl-elastic: the rollout metadata hook runs inside Ray workers, which # inherit the raylet's environment, not the driver's. env_vars[ELASTIC_METADATA_ENV] = "1" - if miles_args is not None: - # Codex harness: the agent function runs in RolloutExecutor actors and - # configures itself from this env (preflight.configure_rollout_worker). - from yeto.rl.harness.codex.preflight import worker_runtime_env - - env_vars.update(worker_runtime_env(miles_args, environ)) ray_module.init(address=address, runtime_env={"env_vars": env_vars}) return address @@ -1974,7 +1953,7 @@ def run_ports_island( start_phases = StartPhases( lambda event, **fields: append_tape_event(miles_args, {"event": event, **fields})) start_phases.replay() - connect_island_ray(miles_args=miles_args) + connect_island_ray() topology = getattr(launch.placement, "topology", None) if topology is not None and topology.nodes > 1: # tasks 3.3: a restarted learner on a partial island (worker node DEAD in the GCS) diff --git a/yeto/rl/adapters/miles/harness_glue/codex_generate.py b/yeto/rl/adapters/miles/harness_glue/codex_generate.py deleted file mode 100644 index 0ab0ba00e..000000000 --- a/yeto/rl/adapters/miles/harness_glue/codex_generate.py +++ /dev/null @@ -1,129 +0,0 @@ -"""SecRLEnv wrapper for pinned Miles' agentic generate function. - -Moved from ``yeto/rl/harness/codex/generate.py`` (yeto-framework-decoupling 5.7): -Miles loads ``generate`` by path (``--custom-generate-function-path``). The -agent-metadata capture it reads stays neutral in ``yeto.rl.harness.codex.generate``. -This file verifies itself against ``yeto.rl.harness.codex.pins.GENERATE_SHA256``. -""" - -from __future__ import annotations - -import copy -import hashlib -import hmac -from pathlib import Path -from typing import Any - -from yeto.rl.harness.codex.generate import _CAPTURED_AGENT_METADATA -from yeto.rl.harness.codex.pins import GENERATE_SHA256 as SECRLENV_GENERATE_SHA256 -from yeto.rl.harness.codex.reward import ( - OUTCOME_KEY, - UntrustedOutcome, - _verified_outcome, - remember_terminal_error, -) - - -def _verify_pinned_source() -> None: - actual = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() - if not hmac.compare_digest(actual, SECRLENV_GENERATE_SHA256): - raise RuntimeError("secrlenv generate wrapper does not match its Yeto pin") - - -_verify_pinned_source() - - -def _generated_samples(output: Any) -> list[Any]: - samples = getattr(output, "samples", None) - if isinstance(samples, list): - flattened: list[Any] = [] - pending = list(samples) - while pending: - sample = pending.pop(0) - if isinstance(sample, list): - pending[0:0] = sample - else: - flattened.append(sample) - return flattened - return [] if samples is None else [samples] - - -def _status_value(sample: Any) -> str | None: - status = getattr(sample, "status", None) - value = getattr(status, "value", status) - return value if isinstance(value, str) else None - - -async def generate(input: Any) -> Any: - """Preserve verified evidence across Miles' no-record aborted fallback. - - Pinned Miles intentionally returns an aborted copy when a session emitted no - model records or all records were truncated, but it does so before merging - the custom agent's metadata. SecRLEnv may restore that metadata only after - validating the complete outcome; every other missing/invalid - evidence path remains fatal. - - The restored sample is marked ``failed``, not ``aborted``: Miles drops a - group with an aborted sample before the dynamic filter runs, so the - SecRLEnv group check would never see it and could not schedule the - same-task retry. A ``failed`` sample still never trains (the group check - rejects any status other than completed or truncated). - """ - - from miles.rollout.generate_hub.agentic_tool_call import generate as upstream - - token = _CAPTURED_AGENT_METADATA.set(None) - try: - return await _generate(upstream, input) - except UntrustedOutcome as error: - remember_terminal_error(getattr(getattr(input, "state", None), "args", None), error) - raise - finally: - _CAPTURED_AGENT_METADATA.reset(token) - - -async def _generate(upstream: Any, input: Any) -> Any: - """Validate the upstream output; restore evidence on the no-record fallback.""" - - output = await upstream(input) - samples = _generated_samples(output) - if not samples: - raise UntrustedOutcome("SecRLEnv generation returned no samples") - - has_outcome = [ - isinstance(getattr(sample, "metadata", None), dict) - and OUTCOME_KEY in sample.metadata - for sample in samples - ] - if all(has_outcome): - return output - if ( - any(has_outcome) - or len(samples) != 1 - or _status_value(samples[0]) != "aborted" - or getattr(samples[0], "response", None) not in {None, ""} - ): - raise UntrustedOutcome("SecRLEnv generation lost outcome metadata") - - captured = _CAPTURED_AGENT_METADATA.get() - _verified_outcome(captured) - if not isinstance(captured, dict): # Kept explicit for type narrowing. - raise UntrustedOutcome("SecRLEnv generation has no outcome") - evidence = {OUTCOME_KEY: copy.deepcopy(captured[OUTCOME_KEY])} - metadata = getattr(samples[0], "metadata", None) - if not isinstance(metadata, dict): - raise UntrustedOutcome("SecRLEnv sample metadata is not an object") - metadata.update(evidence) - from miles.utils.types import Sample - - samples[0].status = Sample.Status.FAILED - return output - - -def _add_arguments(parser: Any) -> None: - from miles.rollout.generate_hub.agentic_tool_call import generate as upstream - - upstream.add_arguments(parser) - - -generate.add_arguments = _add_arguments diff --git a/yeto/rl/adapters/miles/harness_glue/codex_openenv.py b/yeto/rl/adapters/miles/harness_glue/codex_openenv.py deleted file mode 100644 index a71e42b9a..000000000 --- a/yeto/rl/adapters/miles/harness_glue/codex_openenv.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Miles pieces of the codex OpenEnv generate wrapper (yeto-framework-decoupling 5.7). - -``yeto.rl.harness.codex.codex_openenv_generate`` keeps the trajectory -bookkeeping (neutral) and reaches these through the backend registry -(role ``harness_glue``): upstream ``agentic_tool_call.generate``, the -``Sample.Status.ABORTED`` value and the per-session sample collection. -""" - -from __future__ import annotations - -from typing import Any, Awaitable, Callable - - -def load_agentic_upstream() -> Callable[[Any], Awaitable[Any]]: - from miles.rollout.generate_hub.agentic_tool_call import generate as upstream - - return upstream - - -def aborted_status() -> Any: - from miles.utils.types import Sample - - return Sample.Status.ABORTED - - -async def collect_segment_session(input: Any, router: str, session_id: str) -> tuple[list[Any], dict[str, Any]]: - """Collect (and delete) one extra session exactly like upstream's tracer does.""" - from miles.rollout.generate_utils.openai_endpoint_utils import ( - COMPUTED_FIELDS, - ROLLOUT_SAMPLING_MASK_FIELDS, - OpenAIEndpointTracer, - should_return_sampling_mask, - ) - - fields = COMPUTED_FIELDS - if should_return_sampling_mask(input.args, input.sampling_params, evaluation=input.evaluation): - fields += ROLLOUT_SAMPLING_MASK_FIELDS - tracer = OpenAIEndpointTracer(router_url=router, session_id=session_id, samples_wire_fields=fields) - reply = await tracer.collect_samples(input.sample, max_seq_len=getattr(input.args, "max_seq_len", None)) - return list(reply.samples), dict(reply.session_metadata or {}) diff --git a/yeto/rl/adapters/miles/in_flight.py b/yeto/rl/adapters/miles/in_flight.py index 44a32207b..ffa7adc9e 100644 --- a/yeto/rl/adapters/miles/in_flight.py +++ b/yeto/rl/adapters/miles/in_flight.py @@ -104,32 +104,6 @@ def buffered_entries(buffer: Iterable[Sequence[Any]]) -> list[dict[str, Any]]: return entries -def agentic_entries(gates: Mapping[str, Any], inflight: Mapping[str, Mapping[str, Any]]) -> list[dict[str, Any]]: - """One entry per suspended agentic trajectory (reference only).""" - entries = [] - for trajectory_id in sorted(gates): - stage = (inflight.get(trajectory_id) or {}).get("stage") - entry = InFlightTrajectory( - trajectory_id=str(trajectory_id), group_id=f"agentic:{trajectory_id}", - task_id=str(trajectory_id), provenance=TokenProvenance((), ()), - session_ref=f"{AGENTIC_SESSION_PREFIX}{trajectory_id}", - ).to_dict() - if stage: - entry["stage"] = str(stage) - entries.append(entry) - return entries - - -def _codex_suspension() -> tuple[Mapping[str, Any], Mapping[str, Any]]: - import sys - - module = sys.modules.get("yeto.rl.harness.codex.codex_openenv_subprocess_agent_function") - if module is None: - return {}, {} - state = getattr(module, "_SUSPEND_STATE", {}) or {} - return dict(state.get("gates") or {}), dict(getattr(module, "_INFLIGHT", {}) or {}) - - def export_in_flight(executor: Any) -> dict[str, Any]: """Runs INSIDE the rollout executor actor: the in-flight entries and the data buffer length they account for.""" @@ -143,8 +117,7 @@ def export_in_flight(executor: Any) -> dict[str, Any]: source = getattr(executor, "data_source", None) source = getattr(source, "__self__", source) buffer = list(getattr(source, "buffer", None) or []) - gates, inflight = _codex_suspension() - entries = buffered_entries(buffer) + agentic_entries(gates, inflight) + entries = buffered_entries(buffer) return {"entries": entries, "buffer_groups": len(buffer)} diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 48585224f..103d8ca24 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -12,7 +12,6 @@ import tempfile import time from collections.abc import Mapping, Sequence -from decimal import Decimal, DecimalException from pathlib import Path from typing import Any @@ -764,24 +763,6 @@ def _check_dense_full_sweep(args, *, yeto_policy_sync: bool) -> None: raise ValueError("dense full-parameter mode requires --syncer-profile-sha256") -def _trajectory_evidence_kind(args) -> str: - """The evidence kind of a dense full-parameter island, chosen by its reward. - - The SecRLEnv reward binds the grader's signed verdict as evidence. Any - other reward function (for example a GSM8K answer check) binds its own - output, ``sample.reward``; a sample that still carries a grader verdict is - rejected there. Terminal-Bench 2.1 runs the SAO sync, whose runtime - contract names its own evidence kind. - """ - - from yeto.rl import SECRLENV_REWARD - from yeto.rl.trajectory_evidence import REWARD_FUNCTION_EVIDENCE_KIND - - if args.reward_function == SECRLENV_REWARD: - return "secrlenv" - return REWARD_FUNCTION_EVIDENCE_KIND - - def _private_trajectory_evidence_dir(args) -> Path: """A fresh owner-only directory for this island's trajectory evidence.""" @@ -1017,93 +998,13 @@ def _dense_full_training_contract_sha256( "yeto_source_sha256": args.source_sha256, "miles_backend_identity_sha256": _backend_identity_sha256(args), "reward_sha256": args.reward_sha256, - "codex_harness_contract": getattr( - args, - "codex_harness_contract", - None, - ), + # The Codex harness contract left with the harness (S26 E0.5); the + # key stays so the v6 contract hash of existing runs is unchanged. + "codex_harness_contract": None, } ) -def _strict_json_file_semantics(path: Path) -> tuple[Any, ...]: - """Parse exact, typed JSON semantics independent of object-key order.""" - - def unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: - value: dict[str, Any] = {} - for name, item in pairs: - if name in value: - raise ValueError(f"duplicate JSON object key: {name!r}") - value[name] = item - return value - - def finite_number(name: str) -> None: - raise ValueError(f"non-finite JSON number: {name}") - - try: - value = json.loads( - path.read_bytes(), - object_pairs_hook=unique_object, - parse_constant=finite_number, - parse_float=Decimal, - parse_int=Decimal, - ) - except (OSError, UnicodeError, ValueError, DecimalException) as exc: - raise ValueError("Codex app-server schema is malformed JSON") from exc - - def typed_semantics(item: Any) -> tuple[Any, ...]: - if item is None: - return ("null",) - if isinstance(item, bool): - return ("boolean", item) - if isinstance(item, Decimal): - return ("number", item) - if isinstance(item, str): - return ("string", item) - if isinstance(item, list): - return ("array", tuple(typed_semantics(value) for value in item)) - if isinstance(item, dict): - return ( - "object", - tuple( - sorted( - (name, typed_semantics(value)) - for name, value in item.items() - ) - ), - ) - raise TypeError(f"unsupported parsed JSON value: {type(item).__name__}") - - return typed_semantics(value) - - -def _canonical_schema_sha256(path: Path) -> str | None: - from yeto.provenance import canonical_json_sha256 - - try: - return canonical_json_sha256(path) - except (OSError, ValueError): - return None - - -def _verify_live_codex_app_server_schema(pinned: Path, generated: Path) -> None: - try: - pinned_semantics = _strict_json_file_semantics(pinned) - generated_semantics = _strict_json_file_semantics(generated) - except ValueError as exc: - raise ValueError("live stock Codex app-server schema drifted") from exc - if generated_semantics != pinned_semantics: - raise ValueError("live stock Codex app-server schema drifted") - - -def _preflight_codex_openenv_adapter(args, profile_name: str) -> None: - """Attest the in-tree OpenEnv adapter (forwarded to harness.codex.preflight).""" - - from yeto.rl.harness.codex.preflight import forward_legacy_openenv_preflight - - forward_legacy_openenv_preflight(args, profile_name) - - def miles_revision_pins(args: Any) -> Any: """The Miles commit the island checks for ``--miles-commit`` (full SHA); None without it (``verify_miles_revision`` then checks the pinned MILES_NEXT_PINS).""" @@ -1118,279 +1019,6 @@ def miles_revision_pins(args: Any) -> Any: return MilesRevisionPins(MILES_NEXT_REPOSITORY, str(commit).lower()) -def _preflight_codex_harness(args) -> None: - """Fail before model allocation if the signed Codex surface has drifted.""" - - from yeto.rl import ( - CODEX_APP_SERVER_PROTOCOL_REVISION, - CODEX_APP_SERVER_SCHEMA_SHA256, - CODEX_BASE_INSTRUCTIONS_SHA256, - CODEX_CLI_VERSION, - CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH, - CODEX_CONTAINER_BINARY_PATH, - CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256, - CODEX_HARNESS_AGENT, - CODEX_HARNESS_AGENT_SHA256, - CODEX_LINUX_BINARY_SHA256, - CODEX_LINUX_BINARY_SIZE_BYTES, - CODEX_LINUX_TARGET, - CODEX_NPM_PACKAGE, - CODEX_NPM_TARBALL_SHA256, - CODEX_OPENENV_AGENT, - CODEX_OPENENV_IDENTITY_ENV, - CODEX_PACKAGE_MANIFEST_SHA256, - CODEX_SUBMIT_TOOL_SCHEMA_SHA256, - CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256, - SIGNED_CODEX_AGENTS, - ) - - contract = getattr(args, "codex_harness_contract", None) - if args.custom_agent_function_path not in SIGNED_CODEX_AGENTS: - if contract is not None: - raise ValueError( - "--codex-harness-contract requires a signed Codex agent" - ) - return - if not isinstance(contract, dict): - raise ValueError("the signed Codex agent requires its harness contract") - required = { - "agent_function_path", - "agent_source_sha256", - "controller_binary_path", - "controller_package_manifest_path", - "controller_app_server_schema_path", - "bundle_binary_path", - "bundle_package_manifest_path", - "bundle_app_server_schema_path", - "container_binary_path", - "container_app_server_schema_path", - "binary_sha256", - "binary_size_bytes", - "cli_version", - "npm_package", - "target", - "npm_tarball_sha256", - "package_manifest_sha256", - "app_server_protocol_revision", - "app_server_schema_sha256", - "base_instructions_sha256", - "terminal_exec_tool_schema_sha256", - "submit_tool_schema_sha256", - "dynamic_tools_schema_sha256", - "reasoning_effort", - "backend", - } - if args.custom_agent_function_path == CODEX_OPENENV_AGENT: - required.add("openenv_identity_env") - if set(contract) != required: - raise ValueError("stock Codex harness contract shape drifted") - pinned = { - "agent_function_path": CODEX_HARNESS_AGENT, - "agent_source_sha256": CODEX_HARNESS_AGENT_SHA256, - "container_binary_path": CODEX_CONTAINER_BINARY_PATH, - "container_app_server_schema_path": ( - CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH - ), - "binary_sha256": CODEX_LINUX_BINARY_SHA256, - "binary_size_bytes": CODEX_LINUX_BINARY_SIZE_BYTES, - "cli_version": CODEX_CLI_VERSION, - "npm_package": CODEX_NPM_PACKAGE, - "target": CODEX_LINUX_TARGET, - "npm_tarball_sha256": CODEX_NPM_TARBALL_SHA256, - "package_manifest_sha256": CODEX_PACKAGE_MANIFEST_SHA256, - "app_server_protocol_revision": CODEX_APP_SERVER_PROTOCOL_REVISION, - "app_server_schema_sha256": CODEX_APP_SERVER_SCHEMA_SHA256, - "base_instructions_sha256": CODEX_BASE_INSTRUCTIONS_SHA256, - "terminal_exec_tool_schema_sha256": ( - CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256 - ), - "submit_tool_schema_sha256": CODEX_SUBMIT_TOOL_SCHEMA_SHA256, - "dynamic_tools_schema_sha256": CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256, - "reasoning_effort": "xhigh", - } - if any(contract.get(name) != expected for name, expected in pinned.items()): - raise ValueError("stock Codex runtime identity drifted") - if ( - args.custom_agent_function_path == CODEX_OPENENV_AGENT - and contract.get("openenv_identity_env") != CODEX_OPENENV_IDENTITY_ENV - ): - raise ValueError("stock Codex OpenEnv surface identity drifted") - from yeto.rl.codex_backend import ( - stock_codex_backend_contract_with_reasoning_policy, - validate_stock_codex_fields, - ) - - profile_name = getattr(args, "codex_backend_profile", None) or args.tito_model - expected_backend = stock_codex_backend_contract_with_reasoning_policy( - profile_name, - args.rollout_max_response_len, - not getattr(args, "codex_exclude_reasoning_from_loss", False), - ) - validate_stock_codex_fields( - tito_model=args.tito_model, - codex_backend_profile=profile_name, - rl_model_recipe=args.rl_model_recipe, - model=args.model, - model_revision=args.model_revision, - rollout_model=getattr(args, "rollout_model", None), - rollout_model_revision=getattr(args, "rollout_model_revision", None), - apply_chat_template_kwargs=args.apply_chat_template_kwargs, - tito_allowed_append_roles=args.tito_allowed_append_roles, - codex_reasoning_effort=args.codex_reasoning_effort, - # The signed Qwen rollout backend profile predates full-parameter - # training and uses ``attention`` only as its legacy tuning-profile - # discriminator. Full mode emits no LoRA runtime flags below. - lora_targets=( - args.lora_targets - if getattr(args, "parameter_mode", "lora") == "lora" - else "attention" - ), - expert_full_count=getattr(args, "expert_full_count", 0), - lora_expert_rank=( - int(getattr(args, "rl_lora_expert_rank", None) or 0) - if getattr(args, "parameter_mode", "lora") == "lora" - else 0 - ), - ) - if contract.get("backend") != expected_backend: - raise ValueError("stock Codex backend/TITO identity drifted") - from yeto.provenance import file_sha256 - - # OpenEnv owns and attests its adapter inside the pinned Miles source. The - # legacy SecRLEnv-backed Codex adapter remains an optional integration, but - # it must not be required for the independent Terminal-Bench path. - if args.custom_agent_function_path != CODEX_OPENENV_AGENT: - try: - from yeto.rl.harness.codex import codex_harness_agent - - live_identity = codex_harness_agent.codex_harness_identity() - except ( - ImportError, - AttributeError, - TypeError, - ValueError, - RuntimeError, - ) as exc: - raise ValueError("cannot attest the Yeto Codex harness adapter") from exc - identity_names = ( - "base_instructions_sha256", - "terminal_exec_tool_schema_sha256", - "submit_tool_schema_sha256", - "dynamic_tools_schema_sha256", - ) - if live_identity != {name: contract.get(name) for name in identity_names}: - raise ValueError("Yeto Codex adapter identity drifted") - - adapter_path = Path(codex_harness_agent.__file__) - if adapter_path.is_symlink() or not adapter_path.is_file(): - raise ValueError("Yeto Codex adapter source identity drifted") - adapter_path = adapter_path.resolve() - if file_sha256(adapter_path) != contract.get("agent_source_sha256"): - raise ValueError("Yeto Codex adapter source identity drifted") - binary = Path(CODEX_CONTAINER_BINARY_PATH) - manifest = Path("/opt/yeto/codex/codex-package.json") - schema = Path(CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH) - if ( - binary.is_symlink() - or not binary.is_file() - or binary.stat().st_size != CODEX_LINUX_BINARY_SIZE_BYTES - or file_sha256(binary) != CODEX_LINUX_BINARY_SHA256 - or manifest.is_symlink() - or not manifest.is_file() - or file_sha256(manifest) != CODEX_PACKAGE_MANIFEST_SHA256 - or schema.is_symlink() - or not schema.is_file() - or _canonical_schema_sha256(schema) != CODEX_APP_SERVER_SCHEMA_SHA256 - ): - raise ValueError("mounted stock Codex artifact does not match its Yeto pin") - try: - version = subprocess.run( - [str(binary), "--version"], - check=True, - capture_output=True, - text=True, - timeout=15, - ).stdout.strip() - except (OSError, subprocess.SubprocessError) as exc: - raise ValueError("cannot execute the pinned stock Codex binary") from exc - if version != CODEX_CLI_VERSION: - raise ValueError("stock Codex executable version drifted") - try: - with tempfile.TemporaryDirectory(prefix="yeto-codex-schema-") as temporary: - root = Path(temporary) - output = root / "schema" - subprocess.run( - [ - str(binary), - "app-server", - "generate-json-schema", - "--experimental", - "--out", - str(output), - ], - check=True, - capture_output=True, - text=True, - timeout=30, - ) - generated_schema = ( - output / "codex_app_server_protocol.v2.schemas.json" - ) - if ( - generated_schema.is_symlink() - or not generated_schema.is_file() - ): - raise ValueError("live stock Codex app-server schema drifted") - _verify_live_codex_app_server_schema(schema, generated_schema) - except (OSError, subprocess.SubprocessError) as exc: - raise ValueError("cannot verify the stock Codex app-server schema") from exc - expected_env = { - "YETO_CODEX_BINARY_PATH": CODEX_CONTAINER_BINARY_PATH, - "YETO_CODEX_BINARY_SHA256": CODEX_LINUX_BINARY_SHA256, - "YETO_CODEX_BINARY_SIZE_BYTES": str(CODEX_LINUX_BINARY_SIZE_BYTES), - "YETO_CODEX_VERSION": CODEX_CLI_VERSION, - "YETO_CODEX_APP_SERVER_PROTOCOL_REVISION": ( - CODEX_APP_SERVER_PROTOCOL_REVISION - ), - "YETO_CODEX_APP_SERVER_SCHEMA_SHA256": CODEX_APP_SERVER_SCHEMA_SHA256, - "YETO_CODEX_BASE_INSTRUCTIONS_SHA256": contract[ - "base_instructions_sha256" - ], - "YETO_CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256": contract[ - "terminal_exec_tool_schema_sha256" - ], - "YETO_CODEX_SUBMIT_TOOL_SCHEMA_SHA256": contract[ - "submit_tool_schema_sha256" - ], - "YETO_CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256": contract[ - "dynamic_tools_schema_sha256" - ], - "YETO_CODEX_REASONING_EFFORT": "xhigh", - "YETO_CODEX_BACKEND_MAX_TOKENS": str(args.rollout_max_response_len), - "YETO_CODEX_BACKEND_REASONING_EFFORT": expected_backend[ - "reasoning_effort" - ], - "YETO_CODEX_BACKEND_THINKING": "enabled", - "YETO_CODEX_CHAT_TEMPLATE": expected_backend["chat_template"], - "YETO_CODEX_CHAT_TEMPLATE_KWARGS": json.dumps( - expected_backend["chat_template_kwargs"], - sort_keys=True, - separators=(",", ":"), - ), - "YETO_CODEX_TITO_ALLOWED_APPEND_ROLES": "tool,user", - "YETO_CODEX_HARNESS_CONTRACT_SHA256": _canonical_sha256(contract), - } - mismatched = [ - name for name, expected in expected_env.items() if os.getenv(name) != expected - ] - if mismatched: - raise ValueError( - "stock Codex container environment drifted: " + ", ".join(mismatched) - ) - if args.custom_agent_function_path == CODEX_OPENENV_AGENT: - _preflight_codex_openenv_adapter(args, profile_name) - - def _miles_callable(spec: str) -> str: module, separator, function = spec.partition(":") if not separator or not module or not function.isidentifier(): @@ -1989,6 +1617,7 @@ def run_miles( miles_args.external_policy_sync_run_until_stop = sync_preset == "decoupled" if dense_full: from yeto.rl.adapters.miles.full_parameter.session import CONFIG_ATTRIBUTE + from yeto.rl.trajectory_evidence import REWARD_FUNCTION_EVIDENCE_KIND miles_args.external_policy_identity_setter_path = ( "yeto.rl.adapters.miles.publish.set_current_published_policy_identity" @@ -1996,7 +1625,7 @@ def run_miles( miles_args.yeto_rl_trajectory_evidence_dir = str( _private_trajectory_evidence_dir(args) ) - miles_args.yeto_rl_trajectory_evidence_kind = _trajectory_evidence_kind(args) + miles_args.yeto_rl_trajectory_evidence_kind = REWARD_FUNCTION_EVIDENCE_KIND miles_args.yeto_rl_trajectory_evidence_schema_version = 1 setattr( miles_args, @@ -2258,7 +1887,7 @@ def _run_test_debug_rollout_replay(args, miles_args) -> bool: from yeto.rl.adapters.miles.entry import connect_island_ray # same job runtime_env as the ports island (PYTHONPATH -> pinned Miles/Megatron in the actors) - connect_island_ray(miles_args=miles_args) + connect_island_ray() print(f"[rl] TEST ONLY: upstream Miles train loop on {miles_args.load_debug_rollout_data}", flush=True) asyncio.run(with_disposer(module.train, miles_args)) print("[rl] TEST ONLY: upstream Miles train loop finished", flush=True) @@ -2383,7 +2012,6 @@ def main(argv=None) -> None: verify_miles_revision(miles_root) # the pinned MILES_NEXT_PINS else: verify_miles_revision(miles_root, expected=pins) - _preflight_codex_harness(args) probe_syncers(args) from miles.utils.misc import load_function @@ -2393,16 +2021,6 @@ def main(argv=None) -> None: load_function(args.custom_generate_function_path) if args.custom_agent_function_path: load_function(args.custom_agent_function_path) - from yeto.rl import CODEX_HARNESS_AGENT, SECRLENV_AGENT - - if args.custom_agent_function_path in { - SECRLENV_AGENT, - CODEX_HARNESS_AGENT, - }: - from yeto.rl.harness.codex.client import require_daemon_ready - - require_daemon_ready() - print("[rl] secrlenv episode daemon ready") from huggingface_hub import snapshot_download from yeto.models import resolve diff --git a/yeto/rl/adapters/miles/policy_age.py b/yeto/rl/adapters/miles/policy_age.py index d52886feb..53bcab818 100644 --- a/yeto/rl/adapters/miles/policy_age.py +++ b/yeto/rl/adapters/miles/policy_age.py @@ -44,9 +44,9 @@ AGENTIC_GENERATE = "miles.rollout.generate_hub.agentic_tool_call.generate" # Agent functions whose module defines the suspend()/resume() hooks Miles calls # under --agentic-suspend-between-turns (model-turn gate, survival limit). -SUSPEND_CAPABLE_AGENTS = frozenset({ - "yeto.rl.harness.codex.codex_openenv_subprocess_agent_function.run", -}) +# Empty since S26 E0.5: the only such agent was the removed self-written Codex +# harness; the Harbor agent function does not implement the hooks. +SUSPEND_CAPABLE_AGENTS: frozenset[str] = frozenset() def policy_age_argv(limit: int, *, agentic: bool = False, overlap: bool = False) -> tuple[str, ...]: diff --git a/yeto/rl/adapters/miles/rewards.py b/yeto/rl/adapters/miles/rewards.py index 060e15b15..32be9ea51 100644 --- a/yeto/rl/adapters/miles/rewards.py +++ b/yeto/rl/adapters/miles/rewards.py @@ -15,7 +15,7 @@ from typing import Any, Callable -from yeto.rl.rewards.registry import RewardRegistryError, load_reward, register_reward +from yeto.rl.rewards.registry import RewardRegistryError, load_reward from yeto.rl.rewards.types import FilterDecision, RewardResult, Trajectory _MODULE = __name__ @@ -122,29 +122,6 @@ def miles_custom_rm_path(ref: str) -> str: return f"{_MODULE}.{attr}" -def rlf_success_reward(trajectory: Trajectory) -> RewardResult: - """Neutral success declaration for an authenticated RLF grader result. - - This is the framework-neutral ``Trajectory -> RewardResult`` function. - Do NOT pass ``yeto.rl.adapters.miles.rewards:rlf_success_reward`` - directly as Miles ``--custom-rm-path``: Miles invokes reward callables - as ``fn(args, samples)``, so the unwrapped neutral signature would not - match. Use ``miles_custom_rm_path("rlf_success")`` instead; it resolves - to the ``__getattr__``-wrapped attribute below, which adapts the neutral - function to the Miles calling convention via :func:`miles_reward`. - """ - from yeto.rl.rlf.preflight import validate_rollout - validate_rollout({"metadata": trajectory.metadata}) - value = float(trajectory.metadata.get("reward", 0.0)) - return RewardResult(value=value, metadata={"success": bool(trajectory.metadata["success"])}) - - -# Registered so Miles reaches the wrapped callable as -# ``yeto.rl.adapters.miles.rewards.rlf_success`` (module ``__getattr__`` below -# wraps registered neutral rewards with :func:`miles_reward`). -register_reward("rlf_success", rlf_success_reward) - - def __getattr__(attr: str): if attr.startswith("__"): raise AttributeError(attr) diff --git a/yeto/rl/adapters/miles/rollout_meta_hook.py b/yeto/rl/adapters/miles/rollout_meta_hook.py index d34e761d9..b10d7e902 100644 --- a/yeto/rl/adapters/miles/rollout_meta_hook.py +++ b/yeto/rl/adapters/miles/rollout_meta_hook.py @@ -154,7 +154,7 @@ def record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: mask_placeholder_logprobs(data, tokenizer_loader=_load_miles_tokenizer, args=args)) # CompactionRL: the masks above may drop optimised tokens; the per-segment # counts (tokens_after, gae_length) must follow the final masks. - from yeto.rl.harness.codex.codex_openenv_generate import refresh_segment_counts + from yeto.rl.compaction import refresh_segment_counts refresh_segment_counts([s for group in data for s in _flat(group)]) setattr(args, _TRAINED_ATTR, {_group_key(group) for group in data}) diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 22da4c7d9..da8645bc4 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -310,7 +310,7 @@ def emit(event: str, **fields: Any) -> None: incarnation, resume, emit) progress = {"rollout_id": None, "rounds": 0} # the round in flight, rounds done here - connect_island_ray(miles_args=miles_args) + connect_island_ray() # The rollout process's metadata hook pushes each round's group records to this # named actor (on a trainer island the driver creates it and consumes them). metadata = RayMetadataSink() diff --git a/yeto/rl/adapters/miles/secrlenv_miles.py b/yeto/rl/adapters/miles/secrlenv_miles.py deleted file mode 100644 index 70e5add6e..000000000 --- a/yeto/rl/adapters/miles/secrlenv_miles.py +++ /dev/null @@ -1,74 +0,0 @@ -"""Miles classes for the SecRLEnv same-task retry (P0-1.7b.4-Q1). - -Loaded by Miles in the rollout process (``--data-source-path`` and -``--rollout-function-path``, set by :func:`.secrlenv_retry.rollout_argv`); -the contract and its logic live in :mod:`.secrlenv_retry`. -""" - -from __future__ import annotations - -import asyncio -from typing import Any - -from miles.rollout.data_source import RolloutDataSourceWithBuffer -from miles.rollout.inference_rollout.inference_rollout_common import InferenceRolloutFn -from miles.rollout.inference_rollout.inference_rollout_train import abort - -from .secrlenv_retry import ( - RETRY_STATS_ATTR, - TERMINAL_ERRORS, - SecRLEnvTaskRetry, - first_terminal_error, - raise_after_terminal_cleanup, - retry_metrics, -) - - -class SecRLEnvRetryDataSource(SecRLEnvTaskRetry, RolloutDataSourceWithBuffer): - """Miles' buffered data source with the reserved same-task retry slot.""" - - -class SecRLEnvRolloutFn(InferenceRolloutFn): - """Miles' rollout function; on a terminal SecRLEnv error it drains first. - - The training round also reports the retry counters as rollout metrics. - """ - - async def _call_train(self, input: Any) -> Any: - try: - output = await super()._call_train(input) - except TERMINAL_ERRORS as error: - await raise_after_terminal_cleanup( - first_terminal_error(self.state.args, error), - # Miles' end-of-round abort: stop every engine request and - # call the agent's abort hook (tears down local episodes). - abort=lambda: abort(self.state, set(), input.rollout_id), - in_flight=_in_flight_groups, - require_clean=_require_no_episode_residue, - ) - raise # unreachable: the cleanup always raises - stats = getattr(self.state.args, RETRY_STATS_ATTR, None) - metrics = getattr(output, "metrics", None) - if isinstance(metrics, dict): - metrics.update(retry_metrics(stats)) - return output - - -# Miles keeps a round's group tasks in a local of generate_rollout_async; when -# a terminal error leaves that loop they are only reachable on the event loop. -_GROUP_TASK = "generate_and_rm_group" - - -def _in_flight_groups() -> list[asyncio.Task[Any]]: - current = asyncio.current_task() - return [ - task for task in asyncio.all_tasks() - if task is not current and not task.done() - and getattr(task.get_coro(), "__name__", None) == _GROUP_TASK - ] - - -def _require_no_episode_residue() -> None: - from yeto.rl.harness.codex.agent import require_no_episode_residue - - require_no_episode_residue() diff --git a/yeto/rl/adapters/miles/secrlenv_retry.py b/yeto/rl/adapters/miles/secrlenv_retry.py deleted file mode 100644 index 504b1de76..000000000 --- a/yeto/rl/adapters/miles/secrlenv_retry.py +++ /dev/null @@ -1,318 +0,0 @@ -"""SecRLEnv same-task retry on the ports engine (P0-1.7b.4-Q1). - -The SecRLEnv group check (``yeto.rl.harness.codex.reward.check_group``) -verifies each group's signed episode outcomes. A non-trainable group (an -infrastructure failure, or a sample that did not complete) is rejected and the -check asks for one exact rerun of the same task. This module holds the -Miles-independent half of that contract; :mod:`.secrlenv_miles` binds it to -Miles' data source and rollout function classes. - -Rollout contract (unchanged from the removed legacy rollout): - -* the round draws exactly ``rollout_batch_size`` new task groups; the - ``over_sampling_batch_size = rollout_batch_size + 1`` slot is reserved for - the retry, it is not over-sampling; -* at most ``SECRLENV_INFRASTRUCTURE_REPLACEMENTS`` (1) retry per round; a - refill without a queued retry, a second retry, or a retry whose task or - 503-retry ledger does not match the drawn task ends the run with a - sanitized :class:`SecRLEnvRolloutTerminalError`; -* on a terminal error the in-flight work is aborted and the run checks that no - episode is left behind before it raises. - -Miles marks a rollout boundary for the data source: its rollout function -draws the round's groups with ``get_samples`` and returns the aborted groups -with ``add_samples`` once the round ends. -""" - -from __future__ import annotations - -import asyncio -import copy -from collections.abc import Awaitable, Callable, Iterable, Mapping -from typing import Any - -from yeto.rl import SECRLENV_AGENTS, SECRLENV_INFRASTRUCTURE_REPLACEMENTS -from yeto.rl.engine.algorithm import SECRLENV_GROUP_CHECK_FILTER -from yeto.rl.harness.codex.reward import ( - INFRASTRUCTURE_503_RETRIES_KEY, - OUTCOME_KEY, - SecRLEnvReplacementExhausted, - SecRLEnvRolloutCleanupError, - SecRLEnvRolloutTerminalError, - TERMINAL_ERROR_ATTR, - SecRLEnvUntrustedEvidence, - UntrustedOutcome, - sample_statuses, -) - -DATA_SOURCE_PATH = "yeto.rl.adapters.miles.secrlenv_miles.SecRLEnvRetryDataSource" -ROLLOUT_FUNCTION_PATH = "yeto.rl.adapters.miles.secrlenv_miles.SecRLEnvRolloutFn" -# Read by the group check (reward.check_group) to queue the retry. -RETRY_CALLBACK_ATTR = "_yeto_secrlenv_retry_callback" -RETRY_STATS_ATTR = "_yeto_secrlenv_retry_stats" -# The group check's per-round decision memo; reset at each round start. -FILTER_STATE_ATTR = "_yeto_secrlenv_filter_state" -METRIC_PREFIX = "rl/secrlenv/" -TERMINAL_CLEANUP_TIMEOUT_S = 240.0 -_STAT_NAMES = ( - "replacement_attempts", - "nontrainable_groups", - "infrastructure_groups", - "pending_groups", - "failed_groups", -) - - -def rollout_argv(*, custom_agent: str | None, dynamic_filter: str | None, - rollout_batch_size: int, over_sampling_batch_size: int, - max_policy_age: int) -> tuple[str, ...]: - """Miles flags that bind the SecRLEnv retry; empty for any other agent. - - Raises ``ValueError`` when the agent and the group check are not used - together, when the refill slot is not exactly one group, or when the run - carries rollouts across rounds (the retry slot is counted per round). - """ - - secrlenv = custom_agent in SECRLENV_AGENTS - if secrlenv != (dynamic_filter == SECRLENV_GROUP_CHECK_FILTER): - raise ValueError( - f"the SecRLEnv agents and the {SECRLENV_GROUP_CHECK_FILTER!r} filter go " - f"together (agent {custom_agent!r}, filter {dynamic_filter!r})" - ) - if not secrlenv: - return () - if over_sampling_batch_size != rollout_batch_size + 1: - raise ValueError( - "the SecRLEnv agents reserve exactly one refill slot for the same-task " - f"retry: over-sampling must be {rollout_batch_size + 1} " - f"(got {over_sampling_batch_size})" - ) - if max_policy_age: - raise ValueError( - "the SecRLEnv same-task retry is counted per round; it does not support " - f"a policy-age limit (got {max_policy_age})" - ) - return ( - "--data-source-path", DATA_SOURCE_PATH, - "--rollout-function-path", ROLLOUT_FUNCTION_PATH, - ) - - -def _task_indices(group: object) -> tuple[int, ...] | None: - """The sample indexes of a group (nested sample lists use their first item).""" - - if not isinstance(group, list) or not group: - return None - values = [] - for item in group: - while isinstance(item, list): - if not item: - return None - item = item[0] - try: - values.append(int(item.index)) - except (AttributeError, TypeError, ValueError): - return None - return tuple(values) - - -def _strict_int(value: Any) -> bool: - return type(value) is int - - -class SecRLEnvTaskRetry: - """Data-source mixin: reserve the spare refill slot for one same-task retry. - - Put it before a Miles data source class (``RolloutDataSourceWithBuffer``): - it overrides ``get_samples``/``add_samples`` and defers the rest. - """ - - def __init__(self, args: Any) -> None: - super().__init__(args) - target = getattr(args, "rollout_batch_size", None) - capacity = getattr(args, "over_sampling_batch_size", None) - limit = getattr(args, "yeto_rl_secrlenv_max_infrastructure_replacements", None) - if ( - not _strict_int(target) - or target <= 0 - or capacity != target + 1 - or not _strict_int(limit) - or limit != SECRLENV_INFRASTRUCTURE_REPLACEMENTS - ): - raise SecRLEnvReplacementExhausted( - "SecRLEnv rollout replacement contract is invalid" - ) - self._retry_args = args - self._retry_target = target - self._retry_capacity = capacity - self._retry_limit = limit - self._round_open = False - self._reset_round() - setattr(args, RETRY_CALLBACK_ATTR, self.request_retry) - - def _reset_round(self) -> None: - self._templates: dict[tuple[int, ...], list[Any]] = {} - self._queued: list[list[Any]] = [] - self._seen_evidence: set[object] = set() - self._retries_used = 0 - self._retry_stats = dict.fromkeys(_STAT_NAMES, 0) - setattr(self._retry_args, RETRY_STATS_ATTR, self._retry_stats) - setattr(self._retry_args, FILTER_STATE_ATTR, None) - setattr(self._retry_args, TERMINAL_ERROR_ATTR, None) - - def get_samples(self, num_samples: int) -> list[list[Any]]: - if not _strict_int(num_samples) or num_samples != self._retry_capacity: - raise SecRLEnvReplacementExhausted("SecRLEnv rollout refill contract is invalid") - if not self._round_open: - return self._open_round() - if not self._queued: - raise SecRLEnvReplacementExhausted( - "SecRLEnv rollout refill had no authenticated task retry" - ) - groups, self._queued = self._queued, [] - return groups - - def _open_round(self) -> list[list[Any]]: - self._reset_round() - groups = super().get_samples(self._retry_target) - if not isinstance(groups, list) or len(groups) != self._retry_target: - raise SecRLEnvReplacementExhausted( - "SecRLEnv rollout did not receive its exact task batch" - ) - for group in groups: - key = _task_indices(group) - if key is None: - raise SecRLEnvReplacementExhausted("SecRLEnv rollout task identity is unavailable") - self._templates[key] = copy.deepcopy(group) - self._round_open = True - return groups - - def add_samples(self, samples: list[list[Any]]) -> None: - """Miles returns the round's aborted groups here once the round ends.""" - - super().add_samples(samples) - self._round_open = False - - def retry_stats(self) -> dict[str, int]: - return dict(self._retry_stats) - - def request_retry( - self, - group: list[Any], - *, - evidence_key: object, - reason: str, - infrastructure_503_retries: Mapping[int, int], - ) -> None: - """Queue one exact rerun of ``group``'s task (called by the group check).""" - - if evidence_key in self._seen_evidence: - return - self._seen_evidence.add(evidence_key) - if self._retries_used >= self._retry_limit or self._queued: - raise SecRLEnvReplacementExhausted("SecRLEnv rollout replacement budget was exhausted") - key = _task_indices(group) - template = self._templates.get(key) if key is not None else None - if template is None: - raise SecRLEnvReplacementExhausted( - "SecRLEnv rollout replacement task identity did not match" - ) - if ( - not isinstance(infrastructure_503_retries, Mapping) - or set(infrastructure_503_retries) != set(key) - or any( - not _strict_int(index) or not _strict_int(value) or value not in {0, 1} - for index, value in infrastructure_503_retries.items() - ) - ): - raise SecRLEnvUntrustedEvidence("SecRLEnv rollout retry ledger authentication failed") - retry = copy.deepcopy(template) - for sample in retry: - reset = getattr(sample, "reset_for_retry", None) - if not callable(reset): - raise SecRLEnvReplacementExhausted("SecRLEnv rollout sample cannot be reset safely") - reset() - metadata = getattr(sample, "metadata", None) - if isinstance(metadata, dict): - metadata.pop(OUTCOME_KEY, None) - metadata[INFRASTRUCTURE_503_RETRIES_KEY] = ( - infrastructure_503_retries[int(sample.index)] - ) - statuses = sample_statuses(group) - stats = self._retry_stats - stats["nontrainable_groups"] += 1 - stats["pending_groups"] += int("pending" in statuses) - stats["failed_groups"] += int("failed" in statuses) - stats["infrastructure_groups"] += int(reason == "secrlenv_infrastructure_failure") - self._retries_used += 1 - stats["replacement_attempts"] = self._retries_used - self._queued.append(retry) - - -def retry_metrics(stats: Mapping[str, int] | None) -> dict[str, float]: - """The round's retry counters as rollout metrics (``rl/secrlenv/``).""" - - if not isinstance(stats, Mapping): - return {} - return {f"{METRIC_PREFIX}{name}": float(stats.get(name, 0)) for name in _STAT_NAMES} - - -TERMINAL_ERRORS = (SecRLEnvRolloutTerminalError, UntrustedOutcome) - - -def first_terminal_error(args: Any, error: BaseException) -> BaseException: - """The round's first terminal error, else ``error``. - - Miles swallows an exception raised inside a generate task (it only counts - the group as failed); the round then stops later on a different error, - e.g. a refill without a queued retry. The SecRLEnv code records the first - one on ``args`` (``reward.remember_terminal_error``). - """ - - recorded = getattr(args, TERMINAL_ERROR_ATTR, None) - return recorded if isinstance(recorded, TERMINAL_ERRORS) else error - - -async def cancel_and_wait(tasks: Iterable[asyncio.Task[Any]]) -> None: - """Cancel ``tasks`` and wait until every one has finished.""" - - pending = [task for task in tasks if not task.done()] - for task in pending: - task.cancel() - if pending: - await asyncio.gather(*pending, return_exceptions=True) - - -async def raise_after_terminal_cleanup( - error: BaseException, - *, - abort: Callable[[], Awaitable[Any]], - in_flight: Callable[[], Iterable[asyncio.Task[Any]]], - require_clean: Callable[[], None], - timeout_s: float = TERMINAL_CLEANUP_TIMEOUT_S, -) -> None: - """Stop the in-flight work after a terminal error, then raise it sanitized. - - Order: ``abort`` stops the engines and the agent's episodes; the group - tasks ``in_flight`` returns are cancelled and awaited (the backend loop - that owned them has already exited); ``require_clean`` raises if an - episode is left behind. A cleanup failure or timeout (one bound for the - whole sequence) raises :class:`SecRLEnvRolloutCleanupError` instead. - Error text never carries episode payloads (``from None``). - """ - - if isinstance(error, SecRLEnvRolloutTerminalError): - terminal: SecRLEnvRolloutTerminalError = error - else: - terminal = SecRLEnvUntrustedEvidence("SecRLEnv rollout evidence authentication failed") - - async def stop() -> None: - await abort() - await cancel_and_wait(in_flight()) - - try: - await asyncio.wait_for(stop(), timeout=timeout_s) - require_clean() - except Exception: - raise SecRLEnvRolloutCleanupError("SecRLEnv terminal rollout cleanup failed") from None - raise terminal from None diff --git a/yeto/rl/algos/compactionrl.py b/yeto/rl/algos/compactionrl.py index 31fc46315..bb35a8c8e 100644 --- a/yeto/rl/algos/compactionrl.py +++ b/yeto/rl/algos/compactionrl.py @@ -92,8 +92,9 @@ def compactionrl_spec(**overrides: Any) -> AlgorithmSpec: # --- Rollout-side switch (progress.md "S13 Codex 压缩接线") ------------------------ # The trainer-side gae variant decides whether rollouts must be compacted: # ``cross_segment_per_sample`` (this spec) needs one sample per compaction -# segment, which only the Codex bridge (``harness.codex.compaction_bridge``, -# switched on by ``YETO_CODEX_COMPACTIONRL``) produces. The control arm +# segment. The Codex bridge that produced them (switched on by +# ``YETO_CODEX_COMPACTIONRL``) was removed in S26 E0.5; only the TEST synthetic +# reward (``yeto.rl.synthetic_segments``) labels segments now. The control arm # ``cross_segment_whole_rollout`` needs the whole rollout as ONE sample with # per-token segment ids; the bridge runs every segment in its own session-server # session with a rebuilt context, so its segments cannot be concatenated into @@ -102,9 +103,6 @@ def compactionrl_spec(**overrides: Any) -> AlgorithmSpec: COMPACTION_SWITCH_ENV = "YETO_CODEX_COMPACTIONRL" COMPACTION_T_COMP_ENV = "YETO_CODEX_COMPACTIONRL_T_COMP" COMPACTED_GAE_VARIANT = "cross_segment_per_sample" -# yeto's Codex generate wrapper (segment collection); the launcher checks it -# without importing the adapter layer (import boundary, yeto/rl/engine/README.md). -CODEX_OPENENV_GENERATE = "yeto.rl.harness.codex.codex_openenv_generate.generate" WHOLE_ROLLOUT_CONTROL_GAE_VARIANT = "cross_segment_whole_rollout" _SWITCH_FALSE = frozenset({"", "0", "false", "no", "off"}) _SWITCH_TRUE = frozenset({"1", "true", "yes", "on"}) @@ -138,7 +136,7 @@ def check_rollout_compaction( """Fail closed unless spec, harness and switch agree. ``gae_variant``: the spec's ``advantage.gae_variant`` (= Miles ``--gae-variant``); - ``harness_compacts``: the run's agent can compact (the Codex OpenEnv agent). + ``harness_compacts``: the run produces compaction segments. """ if gae_variant == WHOLE_ROLLOUT_CONTROL_GAE_VARIANT and harness_compacts: raise ValueError( @@ -151,7 +149,7 @@ def check_rollout_compaction( if not harness_compacts: raise ValueError( f"gae_variant {COMPACTED_GAE_VARIANT} (CompactionRL) needs compacted rollouts; " - "only the Codex OpenEnv harness can compact" + "no agent harness compacts (only the TEST synthetic segments reward)" ) if not switch_on: raise ValueError( diff --git a/yeto/rl/compaction.py b/yeto/rl/compaction.py index c23e4b36a..dd4730e03 100644 --- a/yeto/rl/compaction.py +++ b/yeto/rl/compaction.py @@ -336,3 +336,40 @@ def validate_compaction_segments(segments: Sequence[Any]) -> None: for value in identity ): raise RuntimeError("Miles compaction segments changed trajectory identity") + + +def _optimized_tokens(sample: Any) -> int: + mask = getattr(sample, "loss_mask", None) or [] + return int(sum(1 for value in mask if value)) + + +def refresh_segment_counts(samples: Any) -> int: + """Recompute ``segment_tokens``/``tokens_after``/``gae_length`` from the final loss masks. + + Segment samples (``metadata.segment_index``) count optimised tokens when + they are assembled; later hooks still set loss_mask 1 -> 0 (S19 #8 + placeholder-logprob stop tokens, S17 WP6 reasoning tokens) before the + backend converts the samples. N_{>s} and l are optimised-token counts, so + they are refreshed after those hooks (s19-compaction-g1-20261010g: + tokens_after was 1-17 tokens high). Returns the number of segment samples + whose counts changed. + """ + rollouts: dict[Any, list[Any]] = {} + for sample in samples: + meta = getattr(sample, "metadata", None) + if not isinstance(meta, dict) or not isinstance(meta.get("segment_index"), int): + continue + key = (getattr(sample, "rollout_id", None), getattr(sample, "group_index", None), + getattr(sample, "index", None)) + rollouts.setdefault(key, []).append(sample) + changed = 0 + for segments in rollouts.values(): + segments.sort(key=lambda s: s.metadata["segment_index"]) + tokens = [_optimized_tokens(s) for s in segments] + for index, sample in enumerate(segments): + new = {"segment_tokens": tokens[index], "tokens_after": sum(tokens[index + 1:]), + "gae_length": sum(tokens)} + if any(sample.metadata.get(k) != v for k, v in new.items()): + changed += 1 + sample.metadata.update(new) + return changed diff --git a/yeto/rl/engine/algorithm.py b/yeto/rl/engine/algorithm.py index 95e67a837..67bc5283e 100644 --- a/yeto/rl/engine/algorithm.py +++ b/yeto/rl/engine/algorithm.py @@ -52,20 +52,10 @@ # *new* ``AlgorithmSpec.sha256()`` (hash-migration.md, 阶段 4 记录). BOUNDED_NONZERO_STD_FILTER = "nonzero_reward_std_bounded" STOCK_NONZERO_STD_FILTER = "nonzero_reward_std" -# SecRLEnv group check (P0-1.7b.4-Q1): verifies each group's signed episode -# outcomes, rejects non-trainable groups and asks the backend to rerun the -# same task once; zero-variance groups follow the bounded rule. -SECRLENV_GROUP_CHECK_FILTER = "secrlenv_group_check" -# Terminal-Bench group check (S24 T-TB2): validates the signed TB2 outcomes and -# drops a group with any infrastructure (aborted) member or zero reward -# variance, so an environment failure never reaches the advantage. -TBENCH_GROUP_CHECK_FILTER = "tbench_group_check" LEGACY_FILTER_NAMES = { "yeto.rl.filters.bounded_nonzero_reward_std": BOUNDED_NONZERO_STD_FILTER, "miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std": STOCK_NONZERO_STD_FILTER, - "yeto.rl.harness.codex.reward.check_group": SECRLENV_GROUP_CHECK_FILTER, - "yeto.rl.harness.codex.tbench_reward.check_group": TBENCH_GROUP_CHECK_FILTER, } @@ -75,8 +65,7 @@ def normalize_filter_name(name: str | None) -> str | None: # v1 (R0) vocabulary: what the flat v1 fields can express. SUPPORTED_ADVANTAGE_ESTIMATORS = frozenset({"grpo", "ppo"}) # ppo: rl-algo-critic-family 2.2 SUPPORTED_LOSSES = frozenset({"policy_loss"}) -SUPPORTED_FILTERS = frozenset({BOUNDED_NONZERO_STD_FILTER, SECRLENV_GROUP_CHECK_FILTER, - TBENCH_GROUP_CHECK_FILTER}) +SUPPORTED_FILTERS = frozenset({BOUNDED_NONZERO_STD_FILTER}) # v2 vocabulary: what the structured spec can *express* (upstream Miles # ``arguments.py`` choices at MILES_NEXT_COMMIT). Expressible is not @@ -136,10 +125,7 @@ def normalize_filter_name(name: str | None) -> str | None: ) # Pre-D1 names of the bounds, still accepted on input (constructor and JSON). CORRECTION_FIELD_ALIASES = {"tis_clip": "upper", "tis_clip_low": "lower"} -DYNAMIC_SAMPLING_FILTERS = { - BOUNDED_NONZERO_STD_FILTER, STOCK_NONZERO_STD_FILTER, SECRLENV_GROUP_CHECK_FILTER, - TBENCH_GROUP_CHECK_FILTER, -} +DYNAMIC_SAMPLING_FILTERS = {BOUNDED_NONZERO_STD_FILTER, STOCK_NONZERO_STD_FILTER} # Plugins may only live in yeto or in Miles' own (built-in) namespace (D7). PLUGIN_NAMESPACES = ["yeto.", "miles."] diff --git a/yeto/rl/engine/backends.py b/yeto/rl/engine/backends.py index 243e07c98..ad12e92f2 100644 --- a/yeto/rl/engine/backends.py +++ b/yeto/rl/engine/backends.py @@ -49,7 +49,6 @@ class BackendEntry: # publish publisher (test-only fault injection env names) # algorithm_flags AlgorithmSpec -> backend argv (critic warmup) # selection_argv facts read from pass-through backend argv (--rl-engine matrix) - # harness_glue codex OpenEnv generate: upstream generate, aborted status, session collect # launch_flags island-entry flag spellings of the neutral layout (launcher) # run_config_rules backend-only checks of resolve_rl_run_config (ref-load format, divisibility) # binding neutral name -> backend binding check (4.4a) @@ -74,7 +73,6 @@ class BackendEntry: "publish": f"{_MILES}.publish", "algorithm_flags": f"{_MILES}.algorithm_flags", "selection_argv": f"{_MILES}.selection_argv", - "harness_glue": f"{_MILES}.harness_glue.codex_openenv", "launch_flags": f"{_MILES}.launch_flags", "run_flags": f"{_MILES}.run_flags", "run_config_rules": f"{_MILES}.run_config_rules", diff --git a/yeto/rl/eval/codex_infer.py b/yeto/rl/eval/codex_infer.py deleted file mode 100644 index 3adc1206e..000000000 --- a/yeto/rl/eval/codex_infer.py +++ /dev/null @@ -1,420 +0,0 @@ -"""Eval-island inference agent: SGLang + Miles session server + the training Codex driver -(rl-eval-difficulty-buckets 5.3/5.9). - -What runs inside the eval-island container, in this order: - -1. :func:`write_peft_dir` -- rebuild the canonical LoRA state from the stored - ``policy.safetensors`` / ``policy.json`` (the store already checked every - file's sha256), require its ``policy_tensor_hash`` to equal the manifest's, - and write a standard PEFT adapter directory; -2. :class:`SglangSessionLoader` (the island's ``PolicyLoader``) -- start one - SGLang server for the base model with LoRA enabled, load the adapter under - Miles' served adapter name (``miles_lora``), and start Miles' own standalone - session server (``miles.rollout.session.server``) in front of it with the - training TITO settings; the next version only swaps the adapter. ``load`` - returns the manifest's ``rl/policy_token`` only after the tensor-hash check - and a successful adapter load (fail closed otherwise); -3. :class:`CodexEvalAgent` (the ``tb2_attempt`` agent port) -- one evaluation - session per attempt, then exactly the training rollout path: the task - statement from ``task_prompt``, the Codex worker subprocess - (``codex_openenv_subprocess_agent_function._drive_worker``, the same - ``drive_untrusted`` + stock Responses bridge), the lease's worker env and - deadline. Judging stays on the trusted side (``tb2_attempt``). - -Nothing here re-implements the harness; the only new pieces are process -supervision and the PEFT writer. Factories for ``island_main``: -``yeto.rl.eval.codex_infer:loader_factory`` / ``:attempt_factory``. -""" - -from __future__ import annotations - -import asyncio -import json -import os -import subprocess -import sys -import time -import urllib.error -import urllib.request -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Callable, Mapping - -SERVED_LORA_NAME = "miles_lora" # miles.utils.lora.utils.LORA_ADAPTER_NAME (the session server sends it) -STATUS_TO_END_REASON = {"completed": "completed", "timeout": "timed_out", - "max_turns": "max_turns", "max_seq_len": "max_seq_len"} - - -# --- 1. stored policy -> PEFT adapter dir ------------------------------------------------ - - -def peft_adapter_config(*, rank: int, targets: list[str], base_model: str, revision: str | None) -> dict[str, Any]: - """The fields ``yeto.rl.initial_adapter.load_initial_adapter`` accepts (alpha = rank, no dropout).""" - return {"peft_type": "LORA", "task_type": "CAUSAL_LM", "r": int(rank), "lora_alpha": int(rank), - "target_modules": sorted(targets), "lora_dropout": 0.0, "bias": "none", "use_rslora": False, - "use_dora": False, "fan_in_fan_out": False, "rank_pattern": {}, "alpha_pattern": {}, - "inference_mode": True, "base_model_name_or_path": base_model, "revision": revision} - - -def lora_rank_and_targets(shapes: Mapping[str, tuple[int, ...]]) -> tuple[int, list[str]]: - """Rank from the ``lora_A`` rows (must be one value); targets = module leaf names.""" - ranks = {int(shape[0]) for name, shape in shapes.items() if ".lora_A." in name} - if len(ranks) != 1: - raise ValueError(f"stored adapter has no single LoRA rank: {sorted(ranks)}") - targets = sorted({name.rsplit(".lora_", 1)[0].rsplit(".", 1)[-1] for name in shapes}) - return ranks.pop(), targets - - -def write_peft_dir(files_dir: str | Path, manifest: Mapping[str, Any], out_dir: str | Path, *, - base_model: str) -> dict[str, Any]: - """Recompute ``policy_tensor_hash`` from the stored bytes, then write the PEFT dir. Needs torch.""" - from safetensors.torch import load_file, save_file - - from .export import POLICY_IDENTITY, POLICY_TENSORS, verify_policy_tensor_hash - - t0 = time.monotonic() - verify_policy_tensor_hash(files_dir, manifest) # raises EvalIntegrityError on mismatch - root = Path(files_dir) - identity = json.loads((root / POLICY_IDENTITY).read_text()) - tensors = load_file(str(root / POLICY_TENSORS)) - rank, targets = lora_rank_and_targets({k: tuple(v.shape) for k, v in tensors.items()}) - out = Path(out_dir) - out.mkdir(parents=True, exist_ok=True) - save_file(tensors, str(out / "adapter_model.safetensors")) - config = peft_adapter_config(rank=rank, targets=targets, base_model=base_model, - revision=identity.get("base_model_revision")) - (out / "adapter_config.json").write_text(json.dumps(config, indent=1, sort_keys=True)) - return {"rank": rank, "targets": targets, "tensors": len(tensors), "path": str(out), - "seconds": round(time.monotonic() - t0, 3)} - - -# --- 2. SGLang + session server ---------------------------------------------------------- - - -def _http(method: str, url: str, body: Any = None, *, timeout: float = 60.0) -> tuple[int, Any]: - data = None if body is None else json.dumps(body).encode() - req = urllib.request.Request(url, data=data, method=method, headers={"Content-Type": "application/json"}) - try: - with urllib.request.urlopen(req, timeout=timeout) as resp: - raw = resp.read() - status = resp.status - except urllib.error.HTTPError as exc: - raw, status = exc.read(), exc.code - try: - return status, json.loads(raw) if raw else None - except ValueError: - return status, raw.decode("utf-8", "replace") - - -@dataclass -class InferConfig: - """Serving settings; the defaults are the S17 M1 codex shape (qwen35_08b, LoRA r16 attention).""" - - base_model: str # HF id or local snapshot dir - revision: str | None = None - tito_model: str = "qwen35" - chat_template_kwargs: Mapping[str, Any] = field(default_factory=lambda: {"clear_thinking": False}) - context_length: int = 8192 # --agent-max-seq-len (the bridge's max_seq_len) - max_tokens: int = 4096 # --rollout-max-response-len (YETO_CODEX_BACKEND_MAX_TOKENS) - temperature: float = 1.0 # Miles rollout defaults - top_p: float = 1.0 - top_k: int | None = None - mem_fraction_static: float = 0.8 - host: str = "127.0.0.1" - sglang_port: int = 30000 - session_port: int = 30100 - work_dir: str = "/tmp/yeto-eval-infer" - sglang_extra_args: tuple[str, ...] = () - start_timeout_s: float = 900.0 - log_dir: str | None = None - - @classmethod - def from_mapping(cls, raw: Mapping[str, Any]) -> "InferConfig": - known = {k: v for k, v in raw.items() if k in cls.__dataclass_fields__} - if "sglang_extra_args" in known: - known["sglang_extra_args"] = tuple(known["sglang_extra_args"]) - return cls(**known) - - @property - def sglang_url(self) -> str: - return f"http://{self.host}:{self.sglang_port}" - - @property - def session_url(self) -> str: - return f"http://{self.host}:{self.session_port}" - - def request_kwargs(self) -> dict[str, Any]: - out: dict[str, Any] = {"temperature": self.temperature, "top_p": self.top_p, "max_tokens": self.max_tokens} - if self.top_k is not None: - out["top_k"] = self.top_k - return out - - -def resolve_serving(tito_model: str) -> dict[str, Any]: - """What the training island derives from ``--tito-model`` (Miles' own resolvers): - the SGLang reasoning / tool-call parsers (``yeto.rl.adapters.miles.config`` adds - them as ``--sglang-*-parser``) and the family's fixed chat template + kwargs - (``miles.utils.arguments``: ``chat_template_path`` for the session server, - ``--sglang-chat-template`` for SGLang).""" - from miles.utils.chat_template_utils import resolve_fixed_chat_template, resolve_reasoning_and_tool_call_parser - - reasoning, tool_call = resolve_reasoning_and_tool_call_parser(tito_model) - template, kwargs = resolve_fixed_chat_template(tito_model) - return {"reasoning_parser": reasoning, "tool_call_parser": tool_call, - "chat_template_path": template, "chat_template_kwargs": dict(kwargs)} - - -def sglang_argv(cfg: InferConfig, *, rank: int, targets: list[str], python: str = sys.executable, - serving: Mapping[str, Any] | None = None) -> list[str]: - """``rank == 0`` serves the base model alone (no LoRA flags): the baseline shape.""" - argv = [python, "-m", "sglang.launch_server", "--model-path", cfg.base_model, - "--host", cfg.host, "--port", str(cfg.sglang_port), "--trust-remote-code", - "--context-length", str(cfg.context_length), "--mem-fraction-static", str(cfg.mem_fraction_static)] - if rank: - argv += ["--enable-lora", "--max-lora-rank", str(rank), "--lora-target-modules", *targets, - "--max-loaded-loras", "2", "--max-loras-per-batch", "1"] - if cfg.revision and not Path(cfg.base_model).exists(): - argv += ["--revision", cfg.revision] - serving = serving or {} - if serving.get("reasoning_parser") is not None: - argv += ["--reasoning-parser", serving["reasoning_parser"]] - if serving.get("tool_call_parser") is not None: - argv += ["--tool-call-parser", serving["tool_call_parser"]] - if serving.get("chat_template_path") is not None: - argv += ["--chat-template", serving["chat_template_path"]] - return argv + list(cfg.sglang_extra_args) - - -def session_server_config(cfg: InferConfig, *, rank: int, serving: Mapping[str, Any] | None = None) -> dict[str, Any]: - """``miles.rollout.session.config.SessionServerConfig`` fields, as the training island sets them - (``compute_session_server_config``) for a LoRA codex run without replay / speculative decoding. - ``rank == 0`` with no adapter path is Miles' only LoRA-free shape (``is_lora_enabled``).""" - serving = serving or {} - kwargs = dict(cfg.chat_template_kwargs) - for key, value in (serving.get("chat_template_kwargs") or {}).items(): - if key in kwargs and kwargs[key] != value: # miles.utils.arguments: conflicting values are invalid - raise ValueError(f"chat template kwarg {key}={kwargs[key]!r} conflicts with the {cfg.tito_model} " - f"family's fixed {value!r}") - kwargs[key] = value - return { - "host": cfg.host, "port": cfg.session_port, "instance_id": "eval-island", - "backend_url": cfg.sglang_url, "timeout": 1800.0, "hf_checkpoint": cfg.base_model, - "chat_template_path": serving.get("chat_template_path"), "tito_model": cfg.tito_model, - "apply_chat_template_kwargs": kwargs, - "use_rollout_routing_replay": False, "use_rollout_indexer_replay": False, - "use_sampling_support_replay": False, "sglang_speculative_algorithm": None, - "num_layers": None, "moe_router_topk": None, "save_debug_trajectory_data": None, - "lora_rank": int(rank), "lora_adapter_path": None, "lora_train_only": False, - "use_session_server": True, "session_message_matcher": "strict", "pause_generation_mode": None, - "session_sample_picker_path": None, "session_sample_postprocessor_path": None, - } - - -SESSION_SERVER_BOOT = ( - "import json, sys\n" - "from miles.rollout.session.config import SessionServerConfig\n" - "from miles.rollout.session.server import run_session_server\n" - "run_session_server(SessionServerConfig(**json.loads(sys.argv[1])))\n" -) - - -class SglangSessionLoader: - """``PolicyLoader`` for ``EvalIsland``: one SGLang + one session server per island.""" - - def __init__(self, cfg: InferConfig, *, popen: Callable[..., Any] = subprocess.Popen, - http: Callable[..., tuple[int, Any]] = _http, sleep: Callable[[float], None] = time.sleep, - clock: Callable[[], float] = time.monotonic, emit: Callable[..., None] | None = None, - serving: Callable[[str], Mapping[str, Any]] = resolve_serving) -> None: - self.cfg = cfg - self.serving = serving - self.popen = popen - self.http = http - self.sleep = sleep - self.clock = clock - self.emit = emit - self.procs: dict[str, Any] = {} - self.served: dict[str, Any] | None = None - self.loads: list[dict[str, Any]] = [] - - def _spawn(self, name: str, argv: list[str]) -> None: - log = None - if self.cfg.log_dir: - Path(self.cfg.log_dir).mkdir(parents=True, exist_ok=True) - log = open(Path(self.cfg.log_dir) / f"{name}.log", "ab") - self.procs[name] = self.popen(argv, stdout=log, stderr=subprocess.STDOUT if log else None, - start_new_session=True) - - def _wait(self, name: str, url: str) -> float: - t0 = self.clock() - while True: - proc = self.procs.get(name) - if proc is not None and proc.poll() is not None: - raise RuntimeError(f"{name} exited with {proc.returncode} before it was ready") - try: - status, _ = self.http("GET", url, timeout=10.0) - except OSError: - status = None - if status == 200: - return round(self.clock() - t0, 3) - if self.clock() - t0 > self.cfg.start_timeout_s: - raise TimeoutError(f"{name} not ready after {self.cfg.start_timeout_s} s ({url})") - self.sleep(2.0) - - def _start(self, rank: int, targets: list[str]) -> dict[str, float]: - serving = dict(self.serving(self.cfg.tito_model)) - self._spawn("sglang", sglang_argv(self.cfg, rank=rank, targets=targets, serving=serving)) - sglang_s = self._wait("sglang", f"{self.cfg.sglang_url}/health_generate") - self._spawn("session-server", [sys.executable, "-c", SESSION_SERVER_BOOT, - json.dumps(session_server_config(self.cfg, rank=rank, serving=serving))]) - session_s = self._wait("session-server", f"{self.cfg.session_url}/health") - self.served = {"rank": rank, "targets": list(targets)} - return {"sglang_start_s": sglang_s, "session_server_start_s": session_s, "serving": serving} - - def start_base(self) -> dict[str, Any]: - """Serve the base model without LoRA (baseline runs: no stored policy to load).""" - if self.served is not None: - raise RuntimeError(f"serving already started: {self.served}") - return self._start(0, []) - - def load(self, manifest: Mapping[str, Any], files_dir: Path) -> str: - version = int(manifest["policy_version"]) - t0 = self.clock() - adapter = write_peft_dir(files_dir, manifest, Path(self.cfg.work_dir) / f"v{version:06d}", - base_model=str(manifest.get("base_model") or self.cfg.base_model)) - record: dict[str, Any] = {"policy_version": version, "verify_and_write_s": adapter["seconds"], - "rank": adapter["rank"], "targets": adapter["targets"]} - if self.served is None: - record.update(self._start(adapter["rank"], adapter["targets"])) - elif [adapter["rank"], adapter["targets"]] != [self.served["rank"], self.served["targets"]]: - raise RuntimeError(f"v{version}: LoRA shape {adapter['rank']}/{adapter['targets']} differs from the " - f"served {self.served}; restart the island") - else: - self.http("POST", f"{self.cfg.sglang_url}/unload_lora_adapter", {"lora_name": SERVED_LORA_NAME}) - t1 = self.clock() - status, body = self.http("POST", f"{self.cfg.sglang_url}/load_lora_adapter", - {"lora_name": SERVED_LORA_NAME, "lora_path": adapter["path"]}, timeout=600.0) - if status != 200 or (isinstance(body, dict) and body.get("success") is False): - raise RuntimeError(f"v{version}: SGLang refused the adapter: HTTP {status} {str(body)[:300]}") - record["adapter_load_s"] = round(self.clock() - t1, 3) - record["load_total_s"] = round(self.clock() - t0, 3) - self.loads.append(record) - if self.emit is not None: - self.emit("rl_eval_policy_load", **record) - return str(manifest["rl/policy_token"]) - - def close(self) -> None: - for proc in self.procs.values(): - if proc.poll() is None: - proc.terminate() - for proc in self.procs.values(): - try: - proc.wait(timeout=30) - except Exception: # noqa: BLE001 - proc.kill() - - -# --- 3. the agent port -------------------------------------------------------------------- - -DETAIL_KEYS = ("end_reason", "last_completion", "parse_failures", "tool_calls", "turns", - "max_model_total_tokens", "timed_out", "max_seq_len_hit", "tito_session_mismatch") - - -def agent_detail(metrics: Mapping[str, Any] | None) -> dict[str, Any] | None: - """Why the episode ended, as the stock driver recorded it (bounded; raw data for the run record).""" - if not metrics: - return None - out = {k: metrics[k] for k in DETAIL_KEYS if k in metrics} - if isinstance(out.get("last_completion"), (str, dict, list)): - out["last_completion"] = json.dumps(out["last_completion"], default=str)[:1500] - return out - - -def _default_drive() -> Callable[..., Any]: - from yeto.rl.harness.codex.codex_openenv_subprocess_agent_function import _drive_worker - - return _drive_worker - - -class CodexEvalAgent: - """``agent(lease, task, trial, *, policy_token)`` for ``tb2_attempt``.""" - - def __init__(self, cfg: InferConfig, *, drive: Callable[..., Any] | None = None, - http: Callable[..., tuple[int, Any]] = _http, clock: Callable[[], float] = time.monotonic) -> None: - self.cfg = cfg - self.drive = drive - self.http = http - self.clock = clock - - async def _session(self) -> str: - body = {"evaluation": True, "temperature": float(self.cfg.temperature), "top_p": float(self.cfg.top_p)} - if self.cfg.top_k is not None: - body["top_k"] = int(self.cfg.top_k) - status, out = await asyncio.to_thread(self.http, "POST", f"{self.cfg.session_url}/sessions", body) - if status != 200 or not isinstance(out, dict) or not out.get("session_id"): - raise RuntimeError(f"session server refused a session: HTTP {status} {str(out)[:200]}") - return str(out["session_id"]) - - async def __call__(self, lease: Any, task: Any, trial: int, *, policy_token: str) -> dict[str, Any]: - from yeto.rl.harness.codex import codex_openenv_agent_function as adapter - from yeto.rl.harness.codex.codex_openenv_subprocess_agent_function import task_prompt - - drive = self.drive or _default_drive() - session_id = await self._session() - episode_id = adapter.new_episode_id() - job = { - "base_url": f"{self.cfg.session_url}/sessions/{session_id}", - "prompt": task_prompt(dict(getattr(task, "row", {}).get("metadata") or {}), None, lease), - "request_kwargs": self.cfg.request_kwargs(), - "episode_id": episode_id, - "max_seq_len": self.cfg.context_length, - "env_url": lease.env_url, - "env_token": lease.env_token, - "driver": "stock", - } - trajectory = f"eval-{task.task_id}-t{trial}-{episode_id}" - t0 = self.clock() - try: - try: - untrusted = await asyncio.wait_for( - drive(job, trajectory, None, getattr(lease, "worker_env", None)), - timeout=float(lease.deadline_seconds)) - except asyncio.TimeoutError: - untrusted = {"status": "timeout", "metrics": {"timed_out": 1}, "episode_id": episode_id} - except adapter.harness.CodexHarnessError as exc: # training: infrastructure, not a reward - return {"episode_id": episode_id, "end_reason": "infra_error", "error": str(exc)[:500], - "metrics": getattr(exc, "metrics", None), "agent_s": round(self.clock() - t0, 3), - "agent_detail": agent_detail(getattr(exc, "metrics", None))} - finally: - await asyncio.to_thread(self.http, "DELETE", f"{self.cfg.session_url}/sessions/{session_id}") - metrics = dict(untrusted.get("metrics") or {}) - status = str(untrusted.get("status")) - if status not in STATUS_TO_END_REASON: - return {"episode_id": episode_id, "end_reason": "infra_error", "error": f"unknown status {status!r}"} - return {"episode_id": untrusted.get("episode_id", episode_id), "end_reason": STATUS_TO_END_REASON[status], - "turns": metrics.get("turns"), "tokens": metrics.get("max_model_total_tokens"), - "metrics": metrics, "agent_s": round(self.clock() - t0, 3), "agent_detail": agent_detail(metrics)} - - -# --- island_main factories ---------------------------------------------------------------- - -_LOADERS: dict[int, SglangSessionLoader] = {} - - -def loader_factory(config: Mapping[str, Any]) -> SglangSessionLoader: - cfg = InferConfig.from_mapping(config.get("infer") or {}) - loader = SglangSessionLoader(cfg) - _LOADERS[id(config)] = loader - return loader - - -def attempt_factory(config: Mapping[str, Any]) -> Callable[..., Mapping[str, Any]]: - """``tb2_attempt`` over the Modal reward-env provider (prebaked images) and :class:`CodexEvalAgent`.""" - from .tb2_attempt import tb2_attempt - - provider_spec = config.get("provider") or os.environ.get("YETO_HARNESS_ENVIRONMENT_PROVIDER") \ - or "yeto.cloud.modal_reward_env:modal_provider" - module, _, name = provider_spec.partition(":") - import importlib - - provider = getattr(importlib.import_module(module), name)(None) - return tb2_attempt(provider, CodexEvalAgent(InferConfig.from_mapping(config.get("infer") or {}))) diff --git a/yeto/rl/eval/rlf_baseline.py b/yeto/rl/eval/rlf_baseline.py deleted file mode 100644 index bcc3505f9..000000000 --- a/yeto/rl/eval/rlf_baseline.py +++ /dev/null @@ -1,714 +0,0 @@ -"""RLF-shape Qwen3-30B-A3B baseline runner (rlf-codex-rl task 5.6). - -The RLF counterpart of ``codex_infer.py`` (whose docstring covers the TB2 -attempt shape). One run, in this order: - -1. admit the CTF task packages into the RLF catalog and render the prompt - rows (``yeto.rl.rlf.catalog.prompt_rows``); -2. bring up serving on the local GPU: one SGLang ``launch_server`` for the - bf16 base model (no LoRA) plus Miles' standalone session server — the - same wire the training rollout uses (Modal deployment shape: - 1xA100-80GB, ``--context-length 32768 --mem-fraction-static 0.9``, model - from the ``yeto-fn-models`` Volume ``hf/Qwen3-30B-A3B/ad44e777/``); -3. per prompt row: acquire one RLF CTF environment lease - (``RlfModalProvider``: nested-Docker target + trusted flag verifier), - create one session, and drive the training Codex subprocess worker - (``_drive_worker``) with the ``qwen3_30b_a3b`` profile budgets - (48 turns, 1024 tokens/turn, 8192-byte tool output) under a HARD - 30-minute per-task deadline; -4. finish on the trusted side (``finish_trusted`` -> RLF grade) and write - one trace record per episode plus a run summary. - -Flag hygiene: the trusted flags reach the target containers only as -``-e FLAG`` (W4); here they additionally scrub every persisted trace string -(``yeto.rl.rlf.modal.scrub``). Submitted flag values are never recorded — -only booleans (``has_flag`` / ``flag_correct``) and a ``flag_seen`` marker -per tool output. Serving that is not ready within ``start_timeout_s`` -(15 min) aborts the whole run as ``infrastructure``; a failed episode -(acquire/drive/grade rejection) is recorded as its own trace with -``classification`` != "success" and never stops the remaining tasks. - -CPU tests inject the provider, the drive function and the HTTP client; the -real ones (Modal sandboxes, SGLang, the Miles session server) only exist on -the GPU container. Modal entry: ``yeto.cloud.modal_rlf_baseline``. -""" - -from __future__ import annotations - -import argparse -import asyncio -import json -import subprocess -import sys -import time -import uuid -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Callable, Mapping - -from yeto.rl.codex_backend import ( - QWEN3_30B_A3B_MODEL, - QWEN3_30B_A3B_REVISION, - stock_codex_backend_profile, -) -from yeto.rl.rlf.catalog import TaskCatalog, import_catalog, prompt_rows -from yeto.rl.rlf.modal import RlfCtfConfig, RlfModalProvider, scrub - -from .codex_infer import SESSION_SERVER_BOOT, STATUS_TO_END_REASON, _http, resolve_serving - -TRACE_SCHEMA = "yeto.rlf.baseline.trace/v1" -SUMMARY_SCHEMA = "yeto.rlf.baseline.summary/v1" -DEFAULT_PROFILE = "qwen3_30b_a3b" -HARD_EPISODE_SECONDS = 1800.0 # 30 min per task (HARD) -SERVING_START_TIMEOUT_S = 900.0 # 15 min: not ready -> infrastructure, abort - - -def volume_snapshot_path(model: str = QWEN3_30B_A3B_MODEL, revision: str = QWEN3_30B_A3B_REVISION, - mount: str = "/mnt/yeto-models") -> str: - """The Modal model Volume layout (``hf///``, dl_hf.py).""" - return f"{mount}/hf/{model.split('/')[-1]}/{revision[:8]}" - - -@dataclass -class BaselineConfig: - """One baseline run. Defaults come from the ``qwen3_30b_a3b`` profile.""" - - tasks_dir: str # admitted CTF packages (task_id -> dir) - flags: Mapping[str, str] # task_id -> trusted flag (never logged) - profile: str = DEFAULT_PROFILE - model_path: str = "" # filled by __post_init__ (Volume snapshot) - revision: str | None = QWEN3_30B_A3B_REVISION - model_identifier: str = QWEN3_30B_A3B_MODEL - tito_model: str = "qwen3" - context_length: int = 32_768 - max_tokens: int = 8_192 # default matches the qwen3_30b_a3b profile (widened 2026-10-11) - max_turns: int = 48 - temperature: float = 1.0 - top_p: float = 1.0 - top_k: int | None = None - mem_fraction_static: float = 0.9 - host: str = "127.0.0.1" - sglang_port: int = 30000 - session_port: int = 30100 - start_timeout_s: float = SERVING_START_TIMEOUT_S - episode_seconds: float = HARD_EPISODE_SECONDS - repeat: int = 1 - task_ids: tuple[str, ...] | None = None - sglang_extra_args: tuple[str, ...] = () - log_dir: str | None = None - preview_bytes: int = 2048 # per-tool-output trace preview cap - - def __post_init__(self) -> None: - if not self.model_path and self.revision: - self.model_path = volume_snapshot_path(self.model_identifier, str(self.revision)) - - @classmethod - def for_profile(cls, profile: str = DEFAULT_PROFILE, **kwargs: Any) -> "BaselineConfig": - """Derive every budget from the registered profile (one source of truth).""" - p = stock_codex_backend_profile(profile) - # explicit kwargs beat the profile (e.g. context_length=24576 for the - # qwen3_cyber compaction budget, which the bridge requires == max_seq_len) - derived = {"revision": p.get("model_revision"), - "model_identifier": str(p.get("model_identifier") or ""), - "tito_model": str(p.get("tito_model") or p["model"]), - "context_length": int(p["max_seq_len"]), - "max_tokens": int(p["max_tokens_per_turn"]), - "max_turns": int(p["max_turns"])} - derived.update({k: v for k, v in kwargs.items() if v is not None or k not in derived}) - return cls(profile=profile, **derived) - - @property - def policy_id(self) -> str: - """The baseline policy is the pinned base model itself.""" - return f"{self.model_identifier}@{self.revision}" - - @property - def session_url(self) -> str: - return f"http://{self.host}:{self.session_port}" - - @property - def sglang_url(self) -> str: - return f"http://{self.host}:{self.sglang_port}" - - def request_kwargs(self) -> dict[str, Any]: - out: dict[str, Any] = {"temperature": self.temperature, "top_p": self.top_p, - "max_tokens": self.max_tokens} - if self.top_k is not None: - out["top_k"] = self.top_k - return out - - def worker_budget_env(self) -> dict[str, str]: - """Per-episode worker env: the profile selection + the launch-time budget knobs - (docs/RLF_CODEX_RUN.md "Qwen3 cyber baseline profile"). The 8192-byte tool - output budget needs no env: the worker reads it from the selected profile.""" - return {"YETO_CODEX_CHAT_TEMPLATE": self.profile, - "SECRLENV_MAX_TURNS": str(self.max_turns), - "YETO_CODEX_BACKEND_MAX_TOKENS": str(self.max_tokens)} - - -class InfrastructureError(RuntimeError): - """Serving never came up (or died) — the run is aborted as infrastructure.""" - - -# --- serving (SGLang + Miles session server, no LoRA) -------------------------------------- - - -def sglang_argv(cfg: BaselineConfig, *, python: str = sys.executable, - serving: Mapping[str, Any] | None = None) -> list[str]: - """Baseline serving: same ``launch_server`` shape as the eval island, minus LoRA.""" - argv = [python, "-m", "sglang.launch_server", "--model-path", cfg.model_path, - "--host", cfg.host, "--port", str(cfg.sglang_port), "--trust-remote-code", - "--context-length", str(cfg.context_length), - "--mem-fraction-static", str(cfg.mem_fraction_static)] - if cfg.revision and not Path(cfg.model_path).exists(): - argv += ["--revision", cfg.revision] - serving = serving or {} - if serving.get("reasoning_parser") is not None: - argv += ["--reasoning-parser", serving["reasoning_parser"]] - if serving.get("tool_call_parser") is not None: - argv += ["--tool-call-parser", serving["tool_call_parser"]] - if serving.get("chat_template_path") is not None: - argv += ["--chat-template", serving["chat_template_path"]] - return argv + list(cfg.sglang_extra_args) - - -def session_server_config(cfg: BaselineConfig, *, serving: Mapping[str, Any] | None = None) -> dict[str, Any]: - """``SessionServerConfig`` fields, as the training island sets them — but LoRA-free. - - ``miles.utils.lora.utils.is_lora_enabled`` = ``lora_rank > 0 or - lora_adapter_path is not None``; pydantic rejects ``lora_rank=None``. So the - only genuinely LoRA-free shape is ``lora_rank=0`` + ``lora_adapter_path=None`` - (second A100 attempt 2026-10-11: rank=32 made the session server request the - ``miles_lora`` adapter from a LoRA-less SGLang -> HTTP 400).""" - serving = serving or {} - return { - "host": cfg.host, "port": cfg.session_port, "instance_id": "rlf-baseline", - "backend_url": cfg.sglang_url, "timeout": 1800.0, "hf_checkpoint": cfg.model_path, - "chat_template_path": serving.get("chat_template_path"), "tito_model": cfg.tito_model, - "apply_chat_template_kwargs": dict(serving.get("chat_template_kwargs") or {}), - "use_rollout_routing_replay": False, "use_rollout_indexer_replay": False, - "use_sampling_support_replay": False, "sglang_speculative_algorithm": None, - "num_layers": None, "moe_router_topk": None, "save_debug_trajectory_data": None, - "lora_rank": 0, "lora_adapter_path": None, "lora_train_only": False, - "use_session_server": True, "session_message_matcher": "strict", - "pause_generation_mode": None, - "session_sample_picker_path": None, "session_sample_postprocessor_path": None, - } - - -class ServingStack: - """One SGLang server + one session server on the local GPU; fail closed on slow starts.""" - - def __init__(self, cfg: BaselineConfig, *, popen: Callable[..., Any] = subprocess.Popen, - http: Callable[..., tuple[int, Any]] = _http, - sleep: Callable[[float], None] = time.sleep, - clock: Callable[[], float] = time.monotonic, - serving: Callable[[str], Mapping[str, Any]] = resolve_serving) -> None: - self.cfg = cfg - self.serving = serving - self.popen = popen - self.http = http - self.sleep = sleep - self.clock = clock - self.procs: dict[str, Any] = {} - - def _spawn(self, name: str, argv: list[str]) -> None: - log = None - if self.cfg.log_dir: - Path(self.cfg.log_dir).mkdir(parents=True, exist_ok=True) - log = open(Path(self.cfg.log_dir) / f"{name}.log", "ab") - self.procs[name] = self.popen(argv, stdout=log, stderr=subprocess.STDOUT if log else None, - start_new_session=True) - - def _wait(self, name: str, url: str) -> float: - t0 = self.clock() - while True: - proc = self.procs.get(name) - if proc is not None and proc.poll() is not None: - raise InfrastructureError(f"{name} exited with {proc.returncode} before it was ready") - try: - status, _ = self.http("GET", url, timeout=10.0) - except OSError: - status = None - if status == 200: - return round(self.clock() - t0, 3) - if self.clock() - t0 > self.cfg.start_timeout_s: - raise InfrastructureError( - f"{name} not ready after {self.cfg.start_timeout_s} s ({url}): infrastructure") - self.sleep(2.0) - - def start(self) -> dict[str, Any]: - serving = dict(self.serving(self.cfg.tito_model)) - self._spawn("sglang", sglang_argv(self.cfg, serving=serving)) - sglang_s = self._wait("sglang", f"{self.cfg.sglang_url}/health_generate") - self._spawn("session-server", [sys.executable, "-c", SESSION_SERVER_BOOT, - json.dumps(session_server_config(self.cfg, serving=serving))]) - session_s = self._wait("session-server", f"{self.cfg.session_url}/health") - return {"sglang_start_s": sglang_s, "session_server_start_s": session_s, - "serving": {k: v for k, v in serving.items()}} - - def close(self) -> None: - for proc in self.procs.values(): - if proc.poll() is None: - proc.terminate() - for proc in self.procs.values(): - try: - proc.wait(timeout=30) - except Exception: # noqa: BLE001 - proc.kill() - - -# --- per-episode trace recording ------------------------------------------------------------- - - -def _scrubbed(value: Any, secrets: tuple[str, ...]) -> Any: - """JSON-round-trip any structure and scrub the secrets out of every string.""" - if not secrets: - return json.loads(json.dumps(value, default=str)) - return json.loads(scrub(json.dumps(value, default=str), *secrets)) - - -class TraceRecorder: - """Wraps the lease's relay environment: per-tool trace events without the flag. - - Records the scrubbed command, a bounded scrubbed output preview and boolean - markers (``flag_seen`` in the output, ``flag_correct`` on submission). The - submitted flag value itself is never stored. - """ - - def __init__(self, relay: Any, *, secrets: tuple[str, ...] = (), preview_bytes: int = 2048) -> None: - self._relay = relay - self._secrets = tuple(s for s in secrets if s) - self._preview_bytes = int(preview_bytes) - self.tool_events: list[dict[str, Any]] = [] - self.submissions: list[dict[str, Any]] = [] - - def attach(self) -> "TraceRecorder": - relay = self._relay - original_execute = relay.execute - original_submit = relay.submit - - async def execute(episode_id: str, command: str, timeout_seconds: float, output_bytes: int): - result = await original_execute(episode_id, command, timeout_seconds, output_bytes) - output = str((result or {}).get("output") or "") - self.tool_events.append({ - "seq": len(self.tool_events), - "command": scrub(command, *self._secrets), - "exit_code": (result or {}).get("exit_code"), - "timed_out": bool((result or {}).get("timed_out")), - "truncated": bool((result or {}).get("truncated")), - "output_bytes": len(output.encode("utf-8")), - "flag_seen": any(s in output for s in self._secrets), - "output_preview": scrub(output[: self._preview_bytes], *self._secrets), - }) - return result - - async def submit(episode_id: str, submission: dict[str, Any]): - result = await original_submit(episode_id, submission) - candidate = submission.get("flag") if isinstance(submission, Mapping) else None - import hmac as _hmac - self.submissions.append({ - "after_tool_calls": len(self.tool_events), - "has_flag": isinstance(candidate, str) and bool(candidate), - "flag_correct": any(isinstance(candidate, str) - and _hmac.compare_digest(candidate, s) for s in self._secrets), - }) - return result - - relay.execute = execute - relay.submit = submit - return self - - -def _turn_records(metrics: Mapping[str, Any]) -> list[dict[str, Any]]: - """Per-turn token/byte rows from the harness's capped per-turn metric lists.""" - completion = list(metrics.get("turn_completion_tokens") or []) - context = list(metrics.get("turn_context_tokens") or []) - tool_bytes = list(metrics.get("turn_tool_output_bytes") or []) - n = max(len(completion), len(context), len(tool_bytes)) - return [{"turn": i, - "completion_tokens": completion[i] if i < len(completion) else None, - "context_tokens": context[i] if i < len(context) else None, - "tool_output_bytes": tool_bytes[i] if i < len(tool_bytes) else None} - for i in range(n)] - - -def _default_drive() -> Callable[..., Any]: - from yeto.rl.harness.codex.codex_openenv_subprocess_agent_function import _drive_worker - - return _drive_worker - - -async def run_episode(cfg: BaselineConfig, provider: Any, row: Mapping[str, Any], *, - session_url: str, drive: Callable[..., Any] | None = None, - http: Callable[..., tuple[int, Any]] = _http, - clock: Callable[[], float] = time.monotonic) -> dict[str, Any]: - """One baseline episode: lease -> session -> Codex worker -> trusted finish -> trace. - - Never raises for episode-level failures: they become trace records with - ``classification`` "infrastructure" (acquire/drive/serving) or "rejected" - (the trusted side refused the grade). - """ - from yeto.rl import tbench_outcome - from yeto.rl.harness.codex import codex_openenv_agent_function as adapter - from yeto.rl.harness.codex.codex_openenv_subprocess_agent_function import task_prompt - from yeto.rl.rlf.preflight import RolloutRejected - - drive = drive or _default_drive() - metadata = dict(row.get("metadata") or {}) - task_id = str(metadata.get("task_id") or "") - trajectory_id = f"rlf-baseline-{task_id}-{uuid.uuid4().hex[:12]}" - secrets = (str(cfg.flags[task_id]),) if cfg.flags.get(task_id) else () - started = clock() - trace: dict[str, Any] = { - "schema": TRACE_SCHEMA, "task_id": task_id, "rlf_task_id": task_id, - "rlf_trajectory_id": trajectory_id, "policy_id": cfg.policy_id, - "prompt": scrub(str(metadata.get("prompt") or row.get("prompt") or ""), *secrets), - "classification": "infrastructure", "success": False, "reward": None, - } - - def infra(reason: str) -> dict[str, Any]: - return {**trace, "error": scrub(reason, *secrets), "wall_s": round(clock() - started, 3)} - - lease = None - session_id = None - segment_ids: list[str] = [] - try: - lease = await provider.acquire(task_id, trajectory_id) - except Exception as exc: # noqa: BLE001 - a failed acquire is infrastructure, never a 0 reward - return infra(f"acquire: {type(exc).__name__}: {exc}") - trace["rlf_episode_id"] = getattr(lease, "episode_id", None) - trace["rlf_trajectory_id"] = getattr(lease, "trajectory_id", None) or trajectory_id - recorder = None - relay = getattr(lease, "environment", None) - if relay is not None: - recorder = TraceRecorder(relay, secrets=secrets, preview_bytes=cfg.preview_bytes).attach() - try: - try: - prompt = task_prompt(metadata, row.get("prompt"), lease) - except Exception as exc: # noqa: BLE001 - a missing statement is a packaging error - return infra(f"prompt: {type(exc).__name__}: {exc}") - body: dict[str, Any] = {"evaluation": True, "temperature": float(cfg.temperature), - "top_p": float(cfg.top_p)} - if cfg.top_k is not None: - body["top_k"] = int(cfg.top_k) - status, out = await asyncio.to_thread(http, "POST", f"{session_url}/sessions", body) - if status != 200 or not isinstance(out, dict) or not out.get("session_id"): - return infra(f"session server refused a session: HTTP {status} " - f"{scrub(str(out)[:200], *secrets)}") - session_id = str(out["session_id"]) - episode_id = str(getattr(lease, "episode_id", None) or adapter.new_episode_id()) - job = { - "base_url": f"{session_url}/sessions/{session_id}", - "prompt": prompt, - "request_kwargs": cfg.request_kwargs(), - "episode_id": episode_id, - "max_seq_len": cfg.context_length, - "env_url": lease.env_url, - "env_token": lease.env_token, - "driver": "stock", - } - # CompactionRL only (no-op when disabled): pre-create segment sessions in - # the trusted layer, same as the training agent function. - try: - segments_meta = await adapter.prepare_segment_sessions(job) - except Exception as exc: # noqa: BLE001 - session-server failures are infrastructure - return infra(f"segment sessions: {type(exc).__name__}: {exc}") - segment_ids = list(segments_meta.get("codex_compaction_sessions") or []) - worker_env = {**dict(getattr(lease, "worker_env", None) or {}), **cfg.worker_budget_env()} - deadline = min(float(getattr(lease, "deadline_seconds", 0) or cfg.episode_seconds), - float(cfg.episode_seconds)) - agent_t0 = clock() - try: - untrusted = await asyncio.wait_for( - drive(job, trajectory_id, None, worker_env), timeout=deadline) - except asyncio.TimeoutError: - untrusted = {"status": "timeout", "metrics": {"timed_out": 1}, "episode_id": episode_id} - except adapter.harness.CodexHarnessError as exc: - return {**infra(f"drive: {type(exc).__name__}: {exc}"), - "agent_metrics": _scrubbed(getattr(exc, "metrics", None) or {}, secrets)} - agent_s = round(clock() - agent_t0, 3) - status = str(untrusted.get("status")) - metrics = dict(untrusted.get("metrics") or {}) - if status not in adapter.POLICY_STATUSES: - return {**infra(f"drive: unknown status {status!r}"), - "agent_metrics": _scrubbed(metrics, secrets)} - try: - trusted = await adapter.finish_trusted(untrusted, lease.verifier, task_id=task_id, - sample_id=trajectory_id) - except RolloutRejected as exc: - return {**infra(""), "classification": "rejected", - "error": scrub(f"grade rejected: {exc}", secrets), - "agent_metrics": _scrubbed(metrics, secrets)} - reward = trusted.get("reward") - if reward is None: - # the timeout path carries the 0 reward only in the outcome envelope - reward = (trusted.get(tbench_outcome.OUTCOME_KEY) or {}).get("reward") - record: dict[str, Any] = { - **trace, - "status": status, - "end_reason": STATUS_TO_END_REASON.get(status, status), - "end_kind": metrics.get("end_kind"), - "classification": "success" if trusted.get("success") else "failed", - "success": bool(trusted.get("success")), - "reward": reward, - "turns": _turn_records(metrics), - "turn_count": metrics.get("turns"), - "tool_calls": metrics.get("tool_calls"), - "submit_calls": metrics.get("submit_calls"), - "tool_events": recorder.tool_events if recorder else [], - "submissions": recorder.submissions if recorder else [], - "max_model_total_tokens": metrics.get("max_model_total_tokens"), - "max_seq_len_hit": metrics.get("max_seq_len_hit"), - "compactions": int((metrics.get("codex_compaction") or {}).get("compactions", 0)), - "segment_sessions": len(segment_ids), - "last_completion": scrub(json.dumps(metrics.get("last_completion"), default=str)[:1500], - *secrets) if metrics.get("last_completion") is not None else None, - "grader_evidence": _scrubbed(trusted.get("grader_evidence") or {}, secrets), - "verifier_log": scrub(str(trusted.get("verifier_log") or "")[-2000:], *secrets) or None, - "agent_metrics": _scrubbed(metrics, secrets), - "agent_s": agent_s, - "wall_s": round(clock() - started, 3), - } - return record - finally: - for sid in segment_ids: - try: - await asyncio.to_thread(http, "DELETE", f"{session_url}/sessions/{sid}") - except Exception: # noqa: BLE001 - teardown never masks the episode outcome - pass - if session_id is not None: - try: - await asyncio.to_thread(http, "DELETE", f"{session_url}/sessions/{session_id}") - except Exception: # noqa: BLE001 - teardown never masks the episode outcome - pass - if lease is not None: - try: - await lease.destroy() - except Exception: # noqa: BLE001 - pass - - -# --- run loop + summary ------------------------------------------------------------------------ - - -async def run_baseline(cfg: BaselineConfig, provider: Any, rows: list[Mapping[str, Any]], *, - session_url: str, drive: Callable[..., Any] | None = None, - http: Callable[..., tuple[int, Any]] = _http, - clock: Callable[[], float] = time.monotonic, - progress: Callable[[str], None] | None = print) -> list[dict[str, Any]]: - """Sequentially run every prompt row (one 32K window serves one episode at a time).""" - traces = [] - for i, row in enumerate(rows): - trace = await run_episode(cfg, provider, row, session_url=session_url, drive=drive, - http=http, clock=clock) - traces.append(trace) - if progress is not None: - progress(f"[rlf-baseline] {i + 1}/{len(rows)} {trace['task_id']}: " - f"{trace['classification']} end={trace.get('end_reason')} " - f"reward={trace.get('reward')}") - return traces - - -def _mean(values: list[float]) -> float | None: - return round(sum(values) / len(values), 4) if values else None - - -def _min_mean_max(values: list[Any]) -> dict[str, Any]: - nums = [float(v) for v in values if isinstance(v, (int, float)) and not isinstance(v, bool)] - if not nums: - return {"mean": None, "max": None} - return {"mean": round(sum(nums) / len(nums), 4), "max": max(nums)} - - -def summarize(traces: list[Mapping[str, Any]], *, policy_id: str | None = None, - config: Mapping[str, Any] | None = None) -> dict[str, Any]: - """Baseline analysis: success rate, reward distribution, budgets usage per episode.""" - infra = [t for t in traces if t.get("classification") == "infrastructure"] - rejected = [t for t in traces if t.get("classification") == "rejected"] - evaluated = [t for t in traces if t.get("classification") in ("success", "failed")] - successes = [t for t in evaluated if t.get("success")] - rewards = [float(t["reward"]) for t in evaluated if t.get("reward") is not None] - end_reasons: dict[str, int] = {} - for t in evaluated: - key = str(t.get("end_reason")) - end_reasons[key] = end_reasons.get(key, 0) + 1 - per_task: dict[str, dict[str, Any]] = {} - for t in traces: - bucket = per_task.setdefault(str(t.get("task_id")), - {"episodes": 0, "successes": 0, "rewards": [], - "infrastructure": 0, "rejected": 0}) - bucket["episodes"] += 1 - if t.get("classification") == "infrastructure": - bucket["infrastructure"] += 1 - elif t.get("classification") == "rejected": - bucket["rejected"] += 1 - if t.get("reward") is not None: - bucket["rewards"].append(float(t["reward"])) - if t.get("success"): - bucket["successes"] += 1 - for bucket in per_task.values(): - bucket["success_rate"] = _mean([1.0 if s else 0.0 for s in - [r == 1.0 for r in bucket.pop("rewards")]]) - n_seq = sum(1 for t in evaluated if t.get("max_seq_len_hit")) - return { - "schema": SUMMARY_SCHEMA, - "policy_id": policy_id, - "config": dict(config or {}), - "episodes": len(traces), - "evaluated": len(evaluated), - "infrastructure": len(infra), - "rejected": len(rejected), - "success_count": len(successes), - "success_rate": _mean([1.0 if t.get("success") else 0.0 for t in evaluated]), - "reward_distribution": {"mean": _mean(rewards), - "counts": {str(r): rewards.count(r) for r in sorted(set(rewards))}}, - "end_reasons": end_reasons, - "max_seq_len_hit": {"count": n_seq, - "rate": _mean([1.0] * n_seq + [0.0] * (len(evaluated) - n_seq))}, - "compactions_total": sum(int(t.get("compactions") or 0) for t in evaluated), - "turns": _min_mean_max([t.get("turn_count") for t in evaluated]), - "tool_calls": _min_mean_max([t.get("tool_calls") for t in evaluated]), - "max_model_total_tokens": _min_mean_max([t.get("max_model_total_tokens") for t in evaluated]), - "per_task": per_task, - } - - -def write_outputs(out_dir: str | Path, traces: list[Mapping[str, Any]], - summary: Mapping[str, Any]) -> dict[str, str]: - out = Path(out_dir) - out.mkdir(parents=True, exist_ok=True) - traces_path = out / "traces.jsonl" - traces_path.write_text("".join(json.dumps(t, ensure_ascii=False, default=str) + "\n" - for t in traces)) - summary_path = out / "summary.json" - summary_path.write_text(json.dumps(summary, indent=1, ensure_ascii=False, default=str)) - return {"traces": str(traces_path), "summary": str(summary_path)} - - -# --- wiring: catalog, provider, serving, CLI ---------------------------------------------------- - - -def load_catalog(tasks_dir: str | Path, *, evidence: Mapping[str, Any] | None = None) -> TaskCatalog: - """Admit every task package (a subdir with task.yaml) under ``tasks_dir``; fail closed.""" - root = Path(tasks_dir) - packages = sorted(p for p in root.iterdir() if p.is_dir() and (p / "task.yaml").is_file()) - if not packages: - from yeto.rl.rlf.contracts import AdmissionError - - raise AdmissionError(f"{tasks_dir}: no RLF task packages (subdir with task.yaml)") - return import_catalog(packages, evidence=evidence) - - -def run_from_config(config: Mapping[str, Any], *, provider: Any = None, - drive: Callable[..., Any] | None = None, - http: Callable[..., tuple[int, Any]] | None = None, - progress: Callable[[str], None] | None = print) -> dict[str, Any]: - """The whole run (in-container body and CLI share this). - - ``config`` keys: ``tasks_dir``, ``flags`` (trusted; never logged), - ``evidence`` (catalog admission), ``out_dir``, plus any - :class:`BaselineConfig` field override under ``baseline``. With - ``serve`` = "start" (default) the local SGLang + session server are - brought up first; "external" expects ``session_url`` (CPU smoke). - """ - base_raw = dict(config.get("baseline") or {}) - if "sglang_extra_args" in base_raw: - base_raw["sglang_extra_args"] = tuple(base_raw["sglang_extra_args"]) - if "task_ids" in base_raw and base_raw["task_ids"] is not None: - base_raw["task_ids"] = tuple(base_raw["task_ids"]) - cfg = BaselineConfig.for_profile(str(config.get("profile") or DEFAULT_PROFILE), - tasks_dir=str(config["tasks_dir"]), - flags=dict(config.get("flags") or {}), **base_raw) - out_dir = Path(str(config.get("out_dir") or "rlf-baseline-out")) - catalog = load_catalog(cfg.tasks_dir, evidence=config.get("evidence")) - rows = prompt_rows(catalog, list(cfg.task_ids) if cfg.task_ids else None, repeat=cfg.repeat) - provider = provider or RlfModalProvider( - ctf=RlfCtfConfig(tasks_dir=cfg.tasks_dir, flags=cfg.flags, - lease_seconds=cfg.episode_seconds)) - stack = None - session_url = str(config.get("session_url") or "") - serving_info: dict[str, Any] = {} - t0 = time.monotonic() - try: - if str(config.get("serve") or "start") == "start": - stack = ServingStack(cfg) - serving_info = stack.start() - session_url = cfg.session_url - if not session_url: - raise InfrastructureError("serve=external needs config['session_url']") - traces = asyncio.run(run_baseline(cfg, provider, rows, session_url=session_url, - drive=drive, http=http or _http, - progress=progress)) - except InfrastructureError as exc: - traces = [] - summary = summarize(traces, policy_id=cfg.policy_id) - summary["aborted"] = {"classification": "infrastructure", "stage": "serving", - "reason": str(exc)[:500]} - paths = write_outputs(out_dir, traces, summary) - return {"ok": False, "summary": summary, "paths": paths} - finally: - if stack is not None: - stack.close() - summary = summarize(traces, policy_id=cfg.policy_id, - config={"profile": cfg.profile, "context_length": cfg.context_length, - "max_tokens": cfg.max_tokens, "max_turns": cfg.max_turns, - "model_path": cfg.model_path, "revision": cfg.revision, - "repeat": cfg.repeat, **serving_info}) - summary["wall_s"] = round(time.monotonic() - t0, 3) - paths = write_outputs(out_dir, traces, summary) - return {"ok": True, "summary": summary, "paths": paths} - - -def main(argv: list[str] | None = None) -> int: - """``python -m yeto.rl.eval.rlf_baseline --tasks-dir DIR --flags-json JSON --out DIR``.""" - parser = argparse.ArgumentParser(description="RLF Qwen3-30B-A3B Codex baseline") - parser.add_argument("--tasks-dir", default=None, - help="admitted CTF packages dir (default $YETO_RLF_CTF_TASKS_DIR)") - parser.add_argument("--flags-json", default=None, - help='{"task_id": "flag{...}"} or @file (default $YETO_RLF_CTF_FLAGS_JSON)') - parser.add_argument("--evidence", default=None, help="JSON catalog admission evidence " - "(must include runtime_quality_gate), or @file") - parser.add_argument("--out", required=True) - parser.add_argument("--repeat", type=int, default=1) - parser.add_argument("--task", action="append", default=None, help="task id (repeatable)") - parser.add_argument("--profile", default=DEFAULT_PROFILE) - parser.add_argument("--model-path", default=None) - parser.add_argument("--serve", choices=["start", "external"], default="start") - parser.add_argument("--session-url", default=None) - parser.add_argument("--episode-seconds", type=float, default=HARD_EPISODE_SECONDS) - parser.add_argument("--start-timeout-s", type=float, default=SERVING_START_TIMEOUT_S) - args = parser.parse_args(argv) - - import os - - def _json_arg(value: str | None, env: str) -> Any: - raw = value if value is not None else os.environ.get(env) - if not raw: - return None - if raw.startswith("@"): - raw = Path(raw[1:]).read_text() - return json.loads(raw) - - tasks_dir = args.tasks_dir or os.environ.get("YETO_RLF_CTF_TASKS_DIR") - if not tasks_dir: - parser.error("--tasks-dir or YETO_RLF_CTF_TASKS_DIR is required") - flags = _json_arg(args.flags_json, "YETO_RLF_CTF_FLAGS_JSON") - if not isinstance(flags, dict) or not flags: - parser.error("--flags-json or YETO_RLF_CTF_FLAGS_JSON must be a JSON object of task_id -> flag") - baseline: dict[str, Any] = {"repeat": args.repeat, "task_ids": args.task, - "episode_seconds": args.episode_seconds, - "start_timeout_s": args.start_timeout_s} - if args.model_path: - baseline["model_path"] = args.model_path - config = {"tasks_dir": tasks_dir, "flags": flags, - "evidence": _json_arg(args.evidence, "YETO_RLF_EVIDENCE_JSON"), - "out_dir": args.out, "profile": args.profile, "serve": args.serve, - "session_url": args.session_url, "baseline": baseline} - result = run_from_config(config) - print(json.dumps(result["summary"], indent=1, ensure_ascii=False, default=str)) - return 0 if result["ok"] else 2 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/yeto/rl/eval/rlf_rl_gate.py b/yeto/rl/eval/rlf_rl_gate.py deleted file mode 100644 index 21e602965..000000000 --- a/yeto/rl/eval/rlf_rl_gate.py +++ /dev/null @@ -1,500 +0,0 @@ -"""Real-GPU RLF RL gate runner and evidence chain (rlf-codex-rl task 5.8). - -One gate run, in this order: - -1. admit the CTF task packages (``fixture-ctf-ezshell``: solvable in ~2 - turns, provides success; ``fixture-ctf-ret2win``: provides reward - variance) into the RLF catalog and render Miles ``--prompt-data`` rows - (``yeto.rl.rlf.catalog.write_prompt_rows`` via :func:`write_prompt_data`); -2. launch one Miles GRPO island on the small ``qwen35_08b`` profile (0.8B, - the S17-verified harness stack; LoRA, 1xGPU) with - ``--custom-rm-path yeto.rl.adapters.miles.rewards.rlf_success`` and - ``--advantage-estimator grpo`` (both already fixed by - the Miles run translation) and the W4 RLF environment provider - (``YETO_HARNESS_ENVIRONMENT_PROVIDER=yeto.cloud.modal_reward_env:rlf_modal_provider`` - + ``YETO_RLF_CTF_TASKS_DIR`` / ``YETO_RLF_CTF_FLAGS_JSON``). The 30B - capability surface is covered separately by the 5.6 baseline; 30B RL - training is a follow-up item (design decision 2026-10-11); -3. after training, collect the gate evidence from real learner/rollout - artifacts — never hand-filled: - - rewards/success: the island's completed-groups checkpoint - (``--completed-groups-path``), i.e. the serialized Miles samples the - learner actually trained on, each carrying the ``rlf_success`` - reward value and ``metadata.success``; - - advantages: recomputed per GRPO group with - ``yeto.rl.rlf.contracts.group_advantages`` (the computation of record - designated by W3 for paths where the learner does not export - per-sample advantages); - - policy update: ``contracts.policy_update_evidence`` over a parameter - summary of the initial adapter vs the final trained adapter - (see "代拍板" note below); -4. judge with ``contracts.rl_gate`` (hard criteria: non-empty reward, - reward_std>0, success_count>0, advantage_abs_sum>0, policy_update) and - write the evidence JSON (reward detail, success, advantage_abs_sum, - policy-update digest before/after). All-zero or constant reward MUST - fail the gate. - -代拍板 (policy-update collection, pending main-agent sign-off): the Miles -fork exports no per-step parameter snapshot hook, so the minimal viable -collection is a parameter summary (per-tensor ``float(sum)``) of the -initial LoRA adapter (``--initial-adapter`` snapshot, or the base snapshot -at t0) versus the final saved adapter, fed to -``contracts.policy_update_evidence``. Rationale: it reuses W3's digest -contract unchanged, proves the optimizer step changed policy parameters, -and needs no fork patch; it cannot attribute the change to a specific -optimizer step, which the gate does not require. - -代拍板 v2 (2026-10-11): after gate attempt 1 lost every island-side file -on teardown, the collection of record is the island event tape's -``rl_publication`` payload hashes (initial vs final policy version) — see -:func:`collect_policy_update_from_tape`. The adapter-file path remains as -``--adapter-before/--adapter-after`` for runs that arrange adapter -snapshots themselves. - -Flag hygiene: trusted flags live only in ``YETO_RLF_CTF_FLAGS_JSON`` (the -Modal Secret). Every persisted evidence string is scrubbed through -``yeto.rl.rlf.modal.scrub`` with the flag values; flags are never logged. - -The verdict is announced only for evidence collected from the real island -artifacts (``collector="miles_island"``). CPU tests inject fake learners -and only exercise the orchestration: an injected collection can never -produce a passing verdict, so the gate cannot be declared green off GPU. -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import time -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Callable, Mapping, Sequence - -from yeto.rl.rlf.catalog import import_catalog, write_prompt_data -from yeto.rl.rlf.contracts import ( - AdmissionError, - PolicyUpdateEvidence, - group_advantages, - policy_update_evidence, - rl_gate, -) -from yeto.rl.rlf.modal import scrub - -EVIDENCE_SCHEMA = "yeto.rlf.rl_gate.evidence/v1" -DEFAULT_PROFILE = "qwen35_08b" -DEFAULT_TASK_IDS = ("fixture-ctf-ezshell", "fixture-ctf-ret2win") -GATE_TASKS_DIR_ENV = "YETO_RLF_CTF_TASKS_DIR" -GATE_FLAGS_ENV = "YETO_RLF_CTF_FLAGS_JSON" - - -class GateEvidenceError(RuntimeError): - """Collected artifacts cannot support a gate judgment (fail closed).""" - - -@dataclass -class GateConfig: - """One 5.8 gate run. Flags never leave the trusted side (never logged).""" - - tasks_dir: str - flags: Mapping[str, str] # task_id -> trusted flag - profile: str = DEFAULT_PROFILE - task_ids: tuple[str, ...] = DEFAULT_TASK_IDS - repeat: int = 2 # prompt rows per task - rounds: int = 2 # Miles --num-rollout - groups_per_round: int = 4 # --rollout-batch-size - samples_per_group: int = 4 # --n-samples-per-prompt - lora_rank: int = 16 - evidence: Mapping[str, Any] = field( - default_factory=lambda: {"runtime_quality_gate": {"healthy": True}}) - - def __post_init__(self) -> None: - missing = [t for t in self.task_ids if t not in self.flags] - if missing: - raise GateEvidenceError(f"no trusted flag for tasks: {missing}") - - def digest(self) -> str: - """Config digest WITHOUT the flags (safe to persist).""" - blob = json.dumps( - {"profile": self.profile, "task_ids": sorted(self.task_ids), - "repeat": self.repeat, "rounds": self.rounds, - "groups_per_round": self.groups_per_round, - "samples_per_group": self.samples_per_group, - "lora_rank": self.lora_rank, - "flag_sha256": {t: hashlib.sha256(self.flags[t].encode()).hexdigest() - for t in sorted(self.flags)}}, - sort_keys=True, separators=(",", ":")) - return hashlib.sha256(blob.encode()).hexdigest() - - -# ---------------------------------------------------------------- prompt data - -def prepare_prompt_data(config: GateConfig, out_path: str | Path) -> Path: - """Admit the gate task packages and write Miles ``--prompt-data`` JSONL.""" - roots = [Path(config.tasks_dir) / task_id for task_id in config.task_ids] - catalog = import_catalog(roots, evidence=dict(config.evidence)) - return write_prompt_data(catalog, out_path, repeat=config.repeat) - - -# ---------------------------------------------------- rollout sample collection - -def _load_island_checkpoint(path: str | Path) -> Mapping[str, Any]: - """Real loader: the island's completed-groups checkpoint (lazy torch).""" - import torch # GPU/learner side only; CPU tests inject ``loader`` - - payload = torch.load(Path(path).expanduser(), map_location="cpu", weights_only=True) - if not isinstance(payload, Mapping) or "completed_groups" not in payload: - raise GateEvidenceError(f"island checkpoint {path} lacks completed_groups") - return payload - - -def collect_rollout_samples( - checkpoint_path: str | Path, - *, - loader: Callable[[str | Path], Mapping[str, Any]] | None = None, -) -> list[dict[str, Any]]: - """Real reward/success collection from the learner's completed groups. - - Returns one record per sample: ``group_index``/``index`` (position in the - checkpoint), ``task_id``, ``reward`` and ``success``. A sample without a - graded ``reward`` or ``metadata.success`` aborts the gate (fail closed) — - silent zeros are exactly what this gate exists to reject. - """ - payload = (loader or _load_island_checkpoint)(checkpoint_path) - groups = payload.get("completed_groups") - if not isinstance(groups, list) or not groups: - raise GateEvidenceError("island checkpoint has no completed groups") - records: list[dict[str, Any]] = [] - for group_index, group in enumerate(groups): - if not isinstance(group, list): - raise GateEvidenceError(f"completed group {group_index} is not a sample list") - for index, sample in enumerate(group): - if isinstance(sample, list): # compaction segments share one outcome - if not sample: - raise GateEvidenceError("empty segment list in completed group") - sample = sample[0] - if not isinstance(sample, Mapping): - raise GateEvidenceError(f"sample {group_index}/{index} is not a mapping") - reward = sample.get("reward") - metadata = sample.get("metadata") or {} - if not isinstance(reward, (int, float)) or isinstance(reward, bool): - raise GateEvidenceError(f"sample {group_index}/{index} lacks a numeric reward") - if "success" not in metadata: - raise GateEvidenceError( - f"sample {group_index}/{index} lacks metadata.success " - "(rlf_success reward must write it)") - records.append({"group_index": group_index, "index": index, - "task_id": metadata.get("rlf_task_id") or metadata.get("task_id"), - "reward": float(reward), - "success": bool(metadata["success"])}) - return records - - -def collect_rollout_samples_from_tape( - tape_path: str | Path, -) -> list[dict[str, Any]]: - """Real reward/success collection from ``rl_sample_outcome`` tape events. - - rlf-codex-rl 5.8 代拍板: on the ports engine the island checkpoint only - ever holds an empty completed-group queue, so sample-level evidence rides - the event tape (mirrored to the Modal volume, survives the island). The - same fail-closed rules as the checkpoint path: any gap in reward/success - aborts the gate. - - Two event shapes are accepted (rlf-codex-rl 5.8 attempt 8 lesson: the - ports launcher binds ``--rollout-all-samples-process-path`` to - ``rollout_meta_hook.extract_rollout_metadata``, so the legacy - ``queue_completed_groups`` hook never runs there; the meta hook already - reports per-sample records which the driver tapes as - ``rl_trajectory_reward`` behind ``--rl-observe-timeline``): - - - ``rl_sample_outcome`` (tapes of the removed legacy engine, 054c20d5); - - ``rl_trajectory_reward`` (rl-fn-codex-rollout 1.0 observe-only pipeline; - ``sample_index`` fills ``index``). - """ - outcomes: list[dict[str, Any]] = [] - for line in Path(tape_path).expanduser().read_text().splitlines(): - line = line.strip() - if not line: - continue - try: - record = json.loads(line) - except json.JSONDecodeError: - continue - if not isinstance(record, dict): - continue - if record.get("event") not in ("rl_sample_outcome", "rl_trajectory_reward"): - continue - outcomes.append(record) - if not outcomes: - raise GateEvidenceError( - "event tape has no rl_sample_outcome/rl_trajectory_reward records" - ) - records: list[dict[str, Any]] = [] - for order, outcome in enumerate(outcomes): - reward = outcome.get("reward") - if not isinstance(reward, (int, float)) or isinstance(reward, bool): - raise GateEvidenceError(f"tape sample {order} lacks a numeric reward") - if outcome.get("success") is None: - raise GateEvidenceError( - f"tape sample {order} lacks success " - "(rlf_success reward must write it)" - ) - records.append({"group_index": int(outcome.get("group_index", order)), - "index": int(outcome.get("index", outcome.get("sample_index", 0))), - "task_id": outcome.get("task_id"), - "reward": float(reward), - "success": bool(outcome["success"])}) - return records - - -# ------------------------------------------------------- policy update evidence - -def _summarize_params_real(path: str | Path) -> dict[str, float]: - """Real loader: per-tensor ``float(sum)`` over a safetensors/torch file or dir.""" - tensors: dict[str, Any] = {} - root = Path(path).expanduser() - files = sorted(root.rglob("*.safetensors")) if root.is_dir() else [root] - if files and all(f.suffix == ".safetensors" for f in files): - from safetensors.torch import load_file # lazy, learner side only - - for file in files: - tensors.update(load_file(str(file))) - else: - import torch # lazy, learner side only - - payload = torch.load(root, map_location="cpu", weights_only=True) - state = payload.get("state_dict", payload) if isinstance(payload, Mapping) else payload - if not isinstance(state, Mapping): - raise GateEvidenceError(f"cannot read parameter tensors from {path}") - tensors = dict(state) - summary: dict[str, float] = {} - for name, tensor in tensors.items(): - try: - summary[str(name)] = float(tensor.float().sum()) - except (TypeError, AttributeError, RuntimeError) as error: - raise GateEvidenceError(f"non-tensor parameter {name!r} in {path}") from error - if not summary: - raise GateEvidenceError(f"no parameter tensors found in {path}") - return summary - - -def collect_policy_update( - before_path: str | Path, - after_path: str | Path, - *, - step: int, - loader: Callable[[str | Path], Mapping[str, float]] | None = None, -) -> PolicyUpdateEvidence: - """Digest the initial vs final adapter parameters (see module 代拍板 note).""" - load = loader or _summarize_params_real - return policy_update_evidence(dict(load(before_path)), dict(load(after_path)), step=step) - - -def collect_policy_update_from_tape( - tape_path: str | Path, - *, - step: int, -) -> PolicyUpdateEvidence: - """Policy-update proof from the island event tape's ``rl_publication`` records. - - 代拍板 v2 (2026-10-11, gate attempt 2): the first gate run lost every - island-side artifact when the island was torn down after the trainer - crash. The event tape is the one channel that is mirrored to a Modal - Volume (``--modal-tape-volume``) and pulled even on failure, and the - driver publishes one ``rl_publication`` event per policy version with a - learner-computed payload hash. Comparing the lowest- and - highest-version payload hashes proves an optimizer step changed policy - parameters without any island-side file surviving — the same - step-attribution limitation as the initial-vs-final adapter digest - (代拍板 v1), with strictly better survivability. - - Fail closed: fewer than two distinct published policy versions, or a - publication without a payload hash, aborts the gate. - """ - publications: dict[int, str] = {} - with Path(tape_path).expanduser().open(encoding="utf-8") as handle: - for line in handle: - line = line.strip() - if not line: - continue - try: - event = json.loads(line) - except json.JSONDecodeError: - continue # prefixed/garbled lines are not gate evidence - if not isinstance(event, Mapping) or event.get("event") != "rl_publication": - continue - version = event.get("policy_version") - digest = event.get("sync/publication_payload_hash") - if not isinstance(version, int) or not isinstance(digest, str) or not digest: - raise GateEvidenceError( - f"rl_publication event lacks policy_version/payload hash: {event!r}" - ) - publications.setdefault(version, digest) - if len(publications) < 2: - raise GateEvidenceError( - f"event tape has {len(publications)} published policy version(s); " - "the gate needs the initial and at least one post-training publication" - ) - first, last = min(publications), max(publications) - before_digest, after_digest = publications[first], publications[last] - updated = before_digest != after_digest - return PolicyUpdateEvidence( - int(step), updated, before_digest, after_digest, - ("policy_payload",) if updated else (), - ) - - -# ------------------------------------------------------------------ gate rows - -def build_gate_rows( - samples: Sequence[Mapping[str, Any]], - update: PolicyUpdateEvidence, -) -> list[dict[str, Any]]: - """Gate rows: real rewards/success, GRPO group advantages, real update flag. - - Advantages come from ``contracts.group_advantages`` per rollout group - (reward minus group mean) — the designated computation of record; the - ``policy_update`` flag comes from the collected PolicyUpdateEvidence, - never from a literal. - """ - by_group: dict[Any, list[Mapping[str, Any]]] = {} - for record in samples: - by_group.setdefault(record["group_index"], []).append(record) - rows: list[dict[str, Any]] = [] - for group_index in sorted(by_group): - group = by_group[group_index] - advantages = group_advantages([r["reward"] for r in group]) - for record, advantage in zip(group, advantages): - rows.append({"reward": float(record["reward"]), - "advantage": float(advantage), - "metadata": {"success": bool(record["success"]), - "task_id": record.get("task_id"), - "group_index": group_index}, - "policy_update": update.updated}) - return rows - - -# ---------------------------------------------------------------------- run - -def run_gate( - rows: Sequence[Mapping[str, Any]], - *, - collector: str, - config: GateConfig | None = None, - update: PolicyUpdateEvidence | None = None, -) -> dict[str, Any]: - """Judge the rows with ``rl_gate`` and assemble the evidence record. - - ``collector`` names the evidence source; only ``"miles_island"`` (the real - learner artifacts) may produce ``verdict=True``. Anything else (CPU fake - learners in tests, manual drills) exercises the orchestration but the - verdict stays False — the gate is never announced off the real run. - """ - try: - result = rl_gate(list(rows)) - gate: dict[str, Any] = {"passed": True, "result": result} - except AdmissionError as error: - gate = {"passed": False, "error": str(error)} - real = collector == "miles_island" - verdict = bool(gate["passed"]) and real - evidence: dict[str, Any] = { - "schema": EVIDENCE_SCHEMA, - "collected_at": time.time(), - "collector": collector, - "verdict": verdict, - "verdict_note": ( - "gate criteria evaluated on real Miles island artifacts" - if real else - "orchestration-only collection (injected/fake learner): " - "the 5.8 gate is never announced off the real GPU run"), - "gate": gate, - "rows": [dict(row) for row in rows], - } - if update is not None: - evidence["policy_update_evidence"] = { - "step": update.step, "updated": update.updated, - "before_digest": update.before_digest, "after_digest": update.after_digest, - "changed": list(update.changed)} - if config is not None: - evidence["config_digest"] = config.digest() - evidence["profile"] = config.profile - evidence["task_ids"] = sorted(config.task_ids) - return evidence - - -def write_evidence(evidence: Mapping[str, Any], out_path: str | Path, - *, secrets: Sequence[str] = ()) -> Path: - """Persist the evidence JSON with every trusted flag scrubbed.""" - text = json.dumps(evidence, ensure_ascii=False, indent=2, sort_keys=True, default=str) - path = Path(out_path) - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(scrub(text, *secrets) + "\n") - return path - - -def main(argv: list[str] | None = None) -> None: - """``python -m yeto.rl.eval.rlf_rl_gate {prompt-data|gate} ...`` (island side).""" - parser = argparse.ArgumentParser(description=__doc__) - sub = parser.add_subparsers(dest="command", required=True) - pd = sub.add_parser("prompt-data", help="admit packages, write Miles --prompt-data") - pd.add_argument("--tasks-dir", required=True) - pd.add_argument("--out", required=True) - pd.add_argument("--repeat", type=int, default=2) - pd.add_argument("--flags-json", default=None, - help="path to the trusted flags JSON (never logged)") - gate = sub.add_parser("gate", help="collect island artifacts, judge, write evidence") - gate.add_argument("--completed-groups", default=None, - help="island checkpoint with completed groups; optional when " - "--tape-events carries rl_sample_outcome/rl_trajectory_reward " - "records (ports engine)") - gate.add_argument("--adapter-before", default=None) - gate.add_argument("--adapter-after", default=None) - gate.add_argument("--tape-events", default=None, - help="island event tape (rl-island-*.jsonl): policy-update proof " - "from rl_publication payload hashes (代拍板 v2); mutually " - "exclusive with --adapter-before/--adapter-after") - gate.add_argument("--step", type=int, required=True) - gate.add_argument("--out", required=True) - gate.add_argument("--flags-json", default=None) - args = parser.parse_args(argv) - - secrets: list[str] = [] - flags_path = getattr(args, "flags_json", None) - if flags_path: - flags = json.loads(Path(flags_path).read_text()) - secrets = [str(v) for v in flags.values()] - if args.command == "prompt-data": - if not flags_path: - raise SystemExit("--flags-json is required (fail closed)") - config = GateConfig(tasks_dir=args.tasks_dir, flags=flags, repeat=args.repeat) - prepare_prompt_data(config, args.out) - return - if args.completed_groups: - samples = collect_rollout_samples(args.completed_groups) - elif args.tape_events: - samples = collect_rollout_samples_from_tape(args.tape_events) - else: - raise SystemExit("rollout samples need --completed-groups or --tape-events") - adapter_pair = (args.adapter_before, args.adapter_after) - if args.tape_events and any(adapter_pair): - raise SystemExit("--tape-events is mutually exclusive with --adapter-before/--adapter-after") - if args.tape_events: - update = collect_policy_update_from_tape(args.tape_events, step=args.step) - update_source = "event_tape" - elif all(adapter_pair): - update = collect_policy_update(args.adapter_before, args.adapter_after, step=args.step) - update_source = "adapter_files" - else: - raise SystemExit("policy-update proof needs --tape-events or --adapter-before + --adapter-after") - rows = build_gate_rows(samples, update) - evidence = run_gate(rows, collector="miles_island", update=update) - evidence["policy_update_source"] = update_source - evidence["sample_source"] = "completed_groups" if args.completed_groups else "event_tape" - write_evidence(evidence, args.out, secrets=secrets) - if not evidence["verdict"]: - raise SystemExit(f"RL gate failed: {evidence['gate']}") - - -if __name__ == "__main__": - main() diff --git a/yeto/rl/eval/tb2_attempt.py b/yeto/rl/eval/tb2_attempt.py deleted file mode 100644 index da329af80..000000000 --- a/yeto/rl/eval/tb2_attempt.py +++ /dev/null @@ -1,63 +0,0 @@ -"""TB2 eval attempt: sandbox lease + agent + trusted verifier (rl-eval-difficulty-buckets 3.1). - -Composes the reward environment's provider (``tb2_provider.Tb2EnvironmentProvider``; -on Modal ``yeto.cloud.modal_reward_env:modal_provider``) with an agent port: - - agent(lease, task, trial, *, policy_token) -> awaitable mapping - {"episode_id", "end_reason", "turns", "tokens"} - -The verifier runs on the trusted side (``lease.verifier.evaluate``) after the -agent stops; the sandbox is always destroyed. A sandbox that cannot be created -or a verifier error propagates and the island records ``infra_error``; an agent -that ran out of turns / context / time keeps its ``end_reason`` and is still -judged (D6.d: only infrastructure errors are excluded from the pass rate). -""" - -from __future__ import annotations - -import asyncio -import time -from typing import Any, Awaitable, Callable, Mapping - -Agent = Callable[..., Awaitable[Mapping[str, Any]]] - - -def trajectory_id(policy_version: int, task_id: str, trial: int) -> str: - return f"eval-v{int(policy_version)}-{task_id}-t{int(trial)}" - - -def tb2_attempt(provider: Any, agent: Agent, *, run: Callable[[Awaitable[Any]], Any] = asyncio.run, - clock: Callable[[], float] = time.monotonic): - """``timing`` (seconds) in the result: ``sandbox_s`` (acquire), ``agent_s``, ``judge_s``, ``destroy_s``.""" - - def attempt(task: Any, trial: int, *, policy_version: int, policy_token: str) -> dict[str, Any]: - tid = trajectory_id(policy_version, task.task_id, trial) - - async def go() -> dict[str, Any]: - timing: dict[str, float] = {} - t = clock() - lease = await provider.acquire(task.task_id, tid) - timing["sandbox_s"] = round(clock() - t, 3) - try: - t = clock() - out = dict(await agent(lease, task, trial, policy_token=policy_token)) - timing["agent_s"] = round(clock() - t, 3) - end = str(out.get("end_reason") or "completed") - if end == "infra_error": - return {**out, "trajectory_id": tid, "reward": 0.0, "success": False, "timing": timing} - t = clock() - verdict = await lease.verifier.evaluate(out.get("episode_id", tid)) - timing["judge_s"] = round(clock() - t, 3) - finally: - t = clock() - await lease.destroy() - timing["destroy_s"] = round(clock() - t, 3) - passed = bool(verdict.get("passed")) - return {"trajectory_id": tid, "reward": 1.0 if passed else 0.0, "success": passed, - "end_reason": end, "turns": out.get("turns"), "tokens": out.get("tokens"), - "verifier_timed_out": bool(verdict.get("timed_out")), "timing": timing, - **({"agent_detail": out["agent_detail"]} if out.get("agent_detail") else {})} - - return run(go()) - - return attempt diff --git a/yeto/rl/eval/tb2_baseline.py b/yeto/rl/eval/tb2_baseline.py deleted file mode 100644 index 059f4b42c..000000000 --- a/yeto/rl/eval/tb2_baseline.py +++ /dev/null @@ -1,169 +0,0 @@ -"""TB2 base-model baseline: rollout only, no training, no stored policy (S24 T-TB2). - -Measures how often a base model passes the TB2 training tasks, to decide -whether TB2 gives a usable RL signal (pass rate neither ~0 nor ~1). One run: - -1. serve the base model on the local GPU without LoRA - (:meth:`codex_infer.SglangSessionLoader.start_base`: SGLang + Miles' - standalone session server with the training TITO settings); -2. run every ``(task, trial)`` pair through exactly the eval-island attempt - (``tb2_attempt`` over the Modal reward-env provider and - :class:`codex_infer.CodexEvalAgent`), ``concurrency`` attempts at a time in - one event loop; -3. append one trace per attempt to ``traces.jsonl`` as it finishes and write - ``summary.json`` at the end. - -Pass-rate rule = the eval island's (D6.d): ``infra_error`` attempts are -excluded from the denominator; an agent that ran out of turns / context / -time is still judged. Entry for Modal: ``yeto.cloud.modal_tb2_baseline``. -""" - -from __future__ import annotations - -import asyncio -import json -import time -from collections import Counter, defaultdict -from pathlib import Path -from typing import Any, Callable, Iterable, Mapping - -from .codex_infer import CodexEvalAgent, InferConfig, SglangSessionLoader -from .island import EvalTask -from .tb2_attempt import tb2_attempt - -SUMMARY_SCHEMA = "yeto.tb2.baseline.summary/v1" -BASE_POLICY_VERSION = 0 -TRACES = "traces.jsonl" -SUMMARY = "summary.json" - - -def load_tasks(split: Mapping[str, Any], rows: Iterable[Mapping[str, Any]], - task_ids: Iterable[str] | None = None) -> list[EvalTask]: - """Tasks of a TB2 split file (``tb2-train.json`` items) with their jsonl rows (prompt metadata).""" - by_id = {str(r["metadata"]["task_id"]): r for r in rows} - keep = None if task_ids is None else set(task_ids) - tasks = [EvalTask(task_id=i["task_id"], eval_bucket=i.get("eval_bucket", "unknown"), - difficulty=i.get("difficulty", "unknown"), benchmark=str(split.get("benchmark", "tb2")), - row=by_id.get(i["task_id"], {})) - for i in split["items"] if keep is None or i["task_id"] in keep] - missing = [t.task_id for t in tasks if not t.row] - if missing: - raise ValueError(f"tasks without a jsonl row: {missing}") - if keep is not None and len(tasks) != len(keep): - raise ValueError(f"tasks not in the split: {sorted(keep - {t.task_id for t in tasks})}") - return tasks - - -async def run_attempts(tasks: list[EvalTask], trials: int, attempt: Callable[..., Any], *, - concurrency: int, on_trace: Callable[[dict[str, Any]], None], - clock: Callable[[], float] = time.time) -> list[dict[str, Any]]: - """``attempt(task, trial, ...)`` must return an awaitable (``tb2_attempt(..., run=lambda c: c)``). - - Trials are the outer loop so a run cut short still covers every task.""" - gate = asyncio.Semaphore(max(1, int(concurrency))) - traces: list[dict[str, Any]] = [] - - async def one(task: EvalTask, trial: int) -> None: - async with gate: - base = {"task_id": task.task_id, "difficulty": task.difficulty, "trial": trial, "start_ts": clock()} - try: - out = dict(await attempt(task, trial, policy_version=BASE_POLICY_VERSION, policy_token="base")) - except Exception as exc: # noqa: BLE001 - sandbox/verifier failure = infrastructure - out = {"end_reason": "infra_error", "reward": 0.0, "success": False, - "error": f"{type(exc).__name__}: {str(exc)[:500]}"} - trace = {**base, **out, "end_ts": clock()} - traces.append(trace) - on_trace(trace) - - await asyncio.gather(*(one(task, trial) for trial in range(trials) for task in tasks)) - return traces - - -def _stats(values: list[float]) -> dict[str, Any] | None: - if not values: - return None - ordered = sorted(values) - return {"n": len(ordered), "mean": round(sum(ordered) / len(ordered), 3), "min": ordered[0], - "p50": ordered[len(ordered) // 2], "max": ordered[-1]} - - -def summarize(traces: list[Mapping[str, Any]], *, trials: int) -> dict[str, Any]: - judged = [t for t in traces if t.get("end_reason") != "infra_error"] - passed = sum(1 for t in judged if t.get("success")) - per_task: dict[str, dict[str, Any]] = defaultdict(lambda: {"passed": 0, "judged": 0, "infra": 0}) - by_diff: dict[str, list[int]] = defaultdict(lambda: [0, 0]) - for t in traces: - row = per_task[str(t["task_id"])] - row["difficulty"] = t.get("difficulty") - if t.get("end_reason") == "infra_error": - row["infra"] += 1 - continue - row["judged"] += 1 - row["passed"] += int(bool(t.get("success"))) - by_diff[str(t.get("difficulty"))][0] += int(bool(t.get("success"))) - by_diff[str(t.get("difficulty"))][1] += 1 - # GRPO-style signal: a task gives a non-zero advantage only if its group mixes pass and fail - mixed = sorted(k for k, v in per_task.items() if 0 < v["passed"] < v["judged"]) - timing = [t.get("timing") or {} for t in judged] - return { - "schema": SUMMARY_SCHEMA, "trials": trials, "attempts": len(traces), "judged": len(judged), - "infra_errors": len(traces) - len(judged), "passed": passed, - "pass_rate": round(passed / len(judged), 4) if judged else None, - "by_difficulty": {k: {"passed": p, "judged": n, "pass_rate": round(p / n, 4) if n else None} - for k, (p, n) in sorted(by_diff.items())}, - "tasks": len(per_task), - "tasks_any_pass": sum(1 for v in per_task.values() if v["passed"] > 0), - "tasks_all_pass": sum(1 for v in per_task.values() if v["judged"] and v["passed"] == v["judged"]), - "tasks_mixed": len(mixed), "mixed_task_ids": mixed, - "per_task": dict(sorted(per_task.items())), - "end_reasons": dict(Counter(str(t.get("end_reason")) for t in traces)), - # the driver's own stop reason (e.g. a CodexModelFailure message), finer than end_reason - "driver_end_reasons": dict(Counter(str((t.get("agent_detail") or {}).get("end_reason")) for t in traces)), - "turns": _stats([float(t["turns"]) for t in judged if t.get("turns") is not None]), - "tokens": _stats([float(t["tokens"]) for t in judged if t.get("tokens") is not None]), - "agent_s": _stats([float(x["agent_s"]) for x in timing if "agent_s" in x]), - "sandbox_s": _stats([float(x["sandbox_s"]) for x in timing if "sandbox_s" in x]), - "judge_s": _stats([float(x["judge_s"]) for x in timing if "judge_s" in x]), - } - - -def _provider(spec: str) -> Any: - import importlib - - module, _, name = spec.partition(":") - return getattr(importlib.import_module(module), name)(None) - - -def run_from_config(config: Mapping[str, Any], *, provider: Any = None, loader: Any = None, - agent: Any = None) -> dict[str, Any]: - """``config``: ``split`` / ``rows`` (paths), ``out_dir``, ``trials``, ``concurrency``, - optional ``task_ids``, ``provider`` (``module:factory``) and ``infer`` (:class:`InferConfig`).""" - out = Path(config["out_dir"]) - out.mkdir(parents=True, exist_ok=True) - split = json.loads(Path(config["split"]).read_text()) - rows = [json.loads(line) for line in Path(config["rows"]).read_text().splitlines() if line.strip()] - tasks = load_tasks(split, rows, config.get("task_ids")) - trials = int(config.get("trials", 4)) - infer = InferConfig.from_mapping(config.get("infer") or {}) - loader = loader or SglangSessionLoader(infer) - t0 = time.time() - try: - serving = loader.start_base() - (out / "serving.json").write_text(json.dumps(serving, indent=1, sort_keys=True, default=str)) - provider = provider or _provider(str(config.get("provider") or "yeto.cloud.modal_reward_env:modal_provider")) - attempt = tb2_attempt(provider, agent or CodexEvalAgent(infer), run=lambda coro: coro) - - def on_trace(trace: dict[str, Any]) -> None: - with open(out / TRACES, "a", encoding="utf-8") as fh: - fh.write(json.dumps(trace, sort_keys=True, default=str) + "\n") - - t1 = time.time() - traces = asyncio.run(run_attempts(tasks, trials, attempt, concurrency=int(config.get("concurrency", 8)), - on_trace=on_trace)) - finally: - loader.close() - summary = {**summarize(traces, trials=trials), "serving": serving, - "serving_start_s": round(t1 - t0, 3), "attempts_wall_s": round(time.time() - t1, 3), - "concurrency": int(config.get("concurrency", 8)), "infer": config.get("infer")} - (out / SUMMARY).write_text(json.dumps(summary, indent=1, sort_keys=True, default=str)) - return summary diff --git a/yeto/rl/harness/README.md b/yeto/rl/harness/README.md index b0a8b4ceb..f7a5110d2 100644 --- a/yeto/rl/harness/README.md +++ b/yeto/rl/harness/README.md @@ -1,41 +1,14 @@ -# yeto.rl.harness — agentic harness integrations (ports path) - -Change: `openspec/changes/rl-codex-harness-rollout`. Packages: `codex/` (Codex -Terminal-Bench harness, moved from `yeto_miles_secrlenv@5bfc011` + rewritten -OpenEnv adapter) and `gateway/` (protocol gateway core: Responses / Chat / -Messages → TITO chat, prefix hash chains). - -## Trust boundary (task 6.3, design D7) - -| Item | Trusted layer | Untrusted layer | -|---|---|---| -| Processes | training driver / `agentic_tool_call` worker; `codex_openenv_subprocess_agent_function.run`; verifier runner | `codex_openenv_agent_worker` subprocess (own session, `start_new_session=True`); Codex 0.145.0 binary (isolated `CODEX_HOME`); task container | -| Verifier assets (`tests/test.sh`, solution) | injected by the sandbox backend only at `verify` time (`TB2_WITHHOLD_TESTS` semantics); `evaluate` is not routed on the worker wire (`environment.serve_environment`) | never readable before verification | -| Network | gateway / session-server endpoint; environment control endpoint (bearer token per lease) | task container: default deny egress except the gateway session endpoint (sandbox backend policy, D9); worker reaches only `env_url` and `base_url` | -| Validation points | `tbench_reward.reward_func`, parent `tbench_outcome.verified_outcome`, `trajectory_evidence._verified_outcome` — all three check outcome structure; any failure rejects the sample. Rewards are not signed (openspec `rl-drop-reward-tamper-hash`) | never runs the verifier | -| Failure classes | infrastructure (acquire/lease/worker protocol): no outcome, sample ABORTED, no 0 reward; policy boundary (timeout / max_turns / max_seq_len) and wrong answer: outcome with reward 0 | — | -| Cleanup | lease `destroy` + `describe()=="gone"` before return; worker SIGTERM→SIGKILL on process group; tool-wait scope balanced on cancel | — | - -CPU evidence: `tests/test_harness_codex_openenv.py` (`*_infrastructure_failures_*`, `*_cancellation_*`). - -## Gateway (tasks 4.1/4.2) - -`gateway.Gateway.handle(entry, body, trajectory_id)`: entry ∈ {responses, chat, -messages} → canonical chat → `ChainRegistry.locate` (continue / fork -`retry_fork` / break `history_rewrite|template_drops_reasoning|compaction_window`) -→ `SessionBackend.chat` on the chain's session → entry response. Signed sampling -fields cannot be overridden; unsupported shapes fail closed. `max_chains=1` is -the first-batch Codex assertion. `ContextProvider` (identity by default) is the -CompactionRL hook (task 7.2). - -## Island wiring (stage 4: IR-1..IR-4 as implemented by INFRA, `integ-decl` 8347351) - -| Concern | Interface (INFRA) | Harness side | -|---|---|---| -| Preflight before any allocation (2.3) | `entry.preflight_stage(miles_args, launch, ...)` runs `HarnessPreflight = Callable[[miles_args, launch], None]`, resolved from `miles_args.yeto_harness_preflight` / `YETO_HARNESS_PREFLIGHT`, before `connect_island_ray` | `codex.preflight.harness_preflight` (`HARNESS_PREFLIGHT_SPEC`): `check_harness_reward_scope`, agent-function check, `preflight_codex_openenv`, then installs the `EnvironmentProvider` (`miles_args.yeto_harness_environment_provider` / `YETO_HARNESS_ENVIRONMENT_PROVIDER`) and the island boards into `codex_openenv_subprocess_agent_function.configure` | -| Drain probe (3.2) | `tool_wait.drain_blockers(router, tool_wait, harness: HarnessSnapshot)`; pool reads `HarnessBoard` via `entry.harness_source` (named actor `yeto-rl-harness-`); `MilesRolloutPool.drain` closes admission first | agent entry: `allow_new_session(member)` → `enter_session`/`exit_session`, `lease_acquired(id, deadline=)`/`lease_released` (idle sandboxes count; hard deadline bounds the wait); gateway: same admission + sessions, `AdmissionClosed` on a draining member | -| Target policy version (4.4) | `RolloutPool.generate(rollout_id, *, expected_policy_version)` → `MilesRolloutPool` publishes the token via the metadata sink; `rollout_meta_hook.expected_policy_version(sample)` reads it (metadata first) | `codex_openenv_generate.generate` and the agent entry refuse (`PolicyVersionMissing`) without a token; drift → `policy_age_violation=1` → `harness_counters` → driver `PolicyIdentityError` | -| 1.7 metrics (4.5) | `timeline.LOAD_SAMPLE_SCHEMA`: `harness_in_flight`, `env_live` (gauge), `tito_session_mismatch`, `tito_chain_breaks{reason}`, `policy_age_violation` (counter; no `_total` suffix), labels `profile_hash`/`epoch` | gateway mirrors `record_session_mismatch` / `record_chain_break(reason)` onto the `HarnessBoard`; generate wrapper writes `policy_age_violation` into sample metadata | -| `reward_scope` (7.1) | `config.check_harness_reward_scope(miles_args.yeto_harness_reward_scope)` in `validate_parsed_args` | re-checked in `harness_preflight`; the harness itself only emits `reward_scope=trajectory` | - -CPU evidence: `tests/test_harness_codex_openenv.py` (`*_harness_board_*`, `*_admission_*`, `*_policy_token_*`, `test_entry_preflight_stage_*`, `test_ir4_schema_names_*`), `tests/test_harness_gateway.py::test_gateway_mirrors_counters_*`, `tests/test_rl_ir_harness.py` (INFRA side). +# yeto.rl.harness — agentic harness integrations + +| Package | Role | +|---|---| +| `harbor/` | Harbor as the agent harness: task loading, sandbox (boat provider), agent, verifier. See `harbor/README.md`. | +| `gateway/` | Records proxy (Responses / Chat -> token-recorded chat) in front of the sampler, prefix hash chains. | +| `reward_env/` | Benchmark-neutral hold-out helpers (`benchmark`, `swebench_verified`). | +| `chat_templates.py`, `reasoning_loss_mask.py`, `placeholder_logprob_mask.py`, `reward_scope.py` | Shared token / mask / reward-scope helpers. | + +The self-written Codex harness (`codex/`: Terminal-Bench task loading, sandbox +broker, tool execution, grading, Codex bridge) was removed in S26 E0.5. The last +commit that has it is tagged `s26-pre-harness-removal`. `session_protocol.md` +documents the session-service protocol that harness used; it is kept as the +contract of `tests/session_service_double.py`. diff --git a/yeto/rl/harness/codex/__init__.py b/yeto/rl/harness/codex/__init__.py deleted file mode 100644 index 2d2e5d9e4..000000000 --- a/yeto/rl/harness/codex/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Codex harness for the ports engine path (moved from yeto_miles_secrlenv@5bfc011).""" diff --git a/yeto/rl/harness/codex/agent.py b/yeto/rl/harness/codex/agent.py deleted file mode 100644 index 445fdb281..000000000 --- a/yeto/rl/harness/codex/agent.py +++ /dev/null @@ -1,1070 +0,0 @@ -"""Miles session-server agent for isolated secrlenv security episodes.""" - -from __future__ import annotations - -import asyncio -import hashlib -import hmac -import json -import logging -import math -import os -import random -import re -import secrets -import time -from collections.abc import Awaitable -from dataclasses import dataclass -from pathlib import Path -from threading import Lock -from typing import Any -from urllib.parse import urlparse, urlunparse - -from .pins import AGENT_SHA256 as SECRLENV_AGENT_SHA256 - -from .client import ( - EpisodeAPIError, - EpisodeClient, - EpisodeClientError, -) -from .generate import capture_agent_metadata -from .reward import ( - ADMISSION_ERROR_CODE, - ADMISSION_ERROR_CODE_KEY, - ADMISSION_NONCE_KEY, - ADMISSION_PHASE, - ADMISSION_PHASE_KEY, - ADMISSION_SCHEMA, - CLEANUP_ERROR_STATUS, - INFRASTRUCTURE_503_RETRIES_KEY, - INFRASTRUCTURE_STATUS, - OUTCOME_KEY, -) - - -class _AdmissionInfrastructureExhausted(EpisodeAPIError): - """An exact current pre-create infrastructure response exhausted its ledger.""" - - def __init__(self) -> None: - super().__init__( - 503, - ADMISSION_ERROR_CODE, - "admission infrastructure exhausted", - ) - - -LOGGER = logging.getLogger(__name__) -_TASK_ID = re.compile(r"CVE-\d{4}-\d{4,}") -_EPISODE_PHASES: dict[str, str] = {} -_EPISODE_POLICY_TASKS: dict[str, asyncio.Task[str]] = {} -_ACTIVE_LOCK = Lock() - - -def _register_episode( - episode_id: str, policy_task: asyncio.Task[str] | None = None -) -> None: - with _ACTIVE_LOCK: - _EPISODE_PHASES[episode_id] = "driving" - if policy_task is not None: - _EPISODE_POLICY_TASKS[episode_id] = policy_task - - -def _claim_episode_finalization(episode_id: str) -> bool: - """Transfer one episode from policy driving to evaluation ownership.""" - - with _ACTIVE_LOCK: - phase = _EPISODE_PHASES.get(episode_id) - if phase == "finalizing": - return True - if phase != "driving": - return False - _EPISODE_PHASES[episode_id] = "finalizing" - return True - - -def _claim_episode_cleanup(episode_id: str) -> bool: - """Claim cancelled-policy cleanup only while policy still owns the episode.""" - - with _ACTIVE_LOCK: - phase = _EPISODE_PHASES.get(episode_id) - if phase != "driving": - return False - _EPISODE_PHASES[episode_id] = "aborting" - return True - - -def _claim_cancelled_policy_cleanup( - episode_id: str, policy_task: asyncio.Task[str] | None -) -> bool: - """Transfer this run's drained finalization boundary to cleanup ownership.""" - - if policy_task is None or not policy_task.done(): - return False - with _ACTIVE_LOCK: - if ( - _EPISODE_POLICY_TASKS.get(episode_id) is not policy_task - or _EPISODE_PHASES.get(episode_id) != "finalizing" - ): - return False - _EPISODE_PHASES[episode_id] = "aborting" - return True - - -async def _recover_failed_normal_close( - client: EpisodeClient, - episode_id: str, - policy_task: asyncio.Task[str] | None, -) -> bool: - """Own and bound a close retry before a rollout may be replaced.""" - - if not _claim_cancelled_policy_cleanup(episode_id, policy_task): - return False - closed = False - try: - closed = await _close_aborted_episode( - client, - episode_id, - retry_timeout_seconds=min( - _positive_env("SECRLENV_ABORT_CLOSE_RETRY_SECONDS", 180.0), - 30.0, - ), - ) - return closed - finally: - if not closed: - _mark_episode_cleanup_pending(episode_id) - - -def _claim_driving_episodes_for_abort() -> list[str]: - """Atomically claim only episodes that have not entered evaluation.""" - - with _ACTIVE_LOCK: - episode_ids = [ - episode_id - for episode_id, phase in _EPISODE_PHASES.items() - if phase == "driving" - ] - for episode_id in episode_ids: - _EPISODE_PHASES[episode_id] = "aborting" - return episode_ids - - -def _claim_episodes_and_tasks_for_abort() -> list[tuple[str, asyncio.Task[str] | None]]: - """Atomically transfer policy ownership and snapshot its in-flight tasks.""" - - with _ACTIVE_LOCK: - claimed = [ - (episode_id, _EPISODE_POLICY_TASKS.get(episode_id)) - for episode_id, phase in _EPISODE_PHASES.items() - if phase in {"driving", "cleanup_pending"} - ] - for episode_id, _task in claimed: - _EPISODE_PHASES[episode_id] = "aborting" - return claimed - - -def _mark_episode_cleanup_pending(episode_id: str) -> None: - """Retain abort ownership while making a failed close retryable.""" - - with _ACTIVE_LOCK: - if _EPISODE_PHASES.get(episode_id) == "aborting": - _EPISODE_PHASES[episode_id] = "cleanup_pending" - - -def _release_episode(episode_id: str) -> None: - with _ACTIVE_LOCK: - _EPISODE_PHASES.pop(episode_id, None) - _EPISODE_POLICY_TASKS.pop(episode_id, None) - - -def require_no_episode_residue() -> None: - """Fail without exposing episode identity when terminal cleanup is incomplete.""" - - with _ACTIVE_LOCK: - if _EPISODE_PHASES or _EPISODE_POLICY_TASKS: - raise RuntimeError("SecRLEnv episode cleanup left active residue") - - -async def _await_policy_and_claim_finalization( - episode_id: str, policy_result: Awaitable[str] -) -> str: - """Claim evaluation ownership before a completed policy await can yield.""" - - try: - return await policy_result - finally: - _claim_episode_finalization(episode_id) - - -async def _drain_aborted_policy_task( - task: asyncio.Task[str] | None, *, timeout_seconds: float -) -> bool: - """Cancel and await policy work before the abort hook closes its episode.""" - - if task is None: - return True - if not task.done(): - task.cancel() - done, _pending = await asyncio.wait({task}, timeout=timeout_seconds) - if not done: - LOGGER.warning( - "aborted secrlenv policy task did not drain within %.1fs", - timeout_seconds, - ) - return False - try: - await task - except asyncio.CancelledError: - pass - except Exception: - LOGGER.warning( - "aborted secrlenv policy task failed while draining", exc_info=True - ) - return True - - -async def _close_aborted_episode( - client: EpisodeClient, - episode_id: str, - *, - retry_timeout_seconds: float, -) -> bool: - """Close one abort-owned episode, retrying only transient lifecycle conflicts.""" - - deadline = time.monotonic() + retry_timeout_seconds - retry_delay = 0.25 - while True: - try: - await client.close(episode_id) - except EpisodeAPIError as exc: - if exc.status == 404: - _release_episode(episode_id) - return True - retryable = exc.status == 409 and exc.code == "episode_conflict" - remaining = deadline - time.monotonic() - if not retryable or remaining <= 0: - LOGGER.warning("abort teardown failed for %s: %s", episode_id, exc) - return False - await asyncio.sleep(min(retry_delay, remaining)) - retry_delay = min(retry_delay * 2.0, 2.0) - except Exception: - LOGGER.warning("abort teardown failed for %s", episode_id, exc_info=True) - return False - else: - _release_episode(episode_id) - return True - - -def _verify_pinned_source() -> None: - actual = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() - if not hmac.compare_digest(actual, SECRLENV_AGENT_SHA256): - raise RuntimeError("secrlenv agent source does not match its Yeto pin") - - -_verify_pinned_source() - -SYSTEM_PROMPT = """You are an autonomous security-testing agent inside an isolated, -authorized benchmark. Solve the supplied target task using only the two provided -tools. `terminal.exec` runs a shell command in your persistent operator workspace; -it has TARGET_URL, YETO_NONCE and YETO_SUBMISSION. The scored target is available -only over the network: you have no target shell, source checkout, container API, -or auxiliary copy. Do not merely describe commands: call terminal.exec. When you -have concrete evidence, call submit exactly once with the final JSON submission. -Do not invent tool results or claim success without evidence.""" - -TOOLS = [ - { - "type": "function", - "function": { - "name": "terminal.exec", - "description": ( - "Execute one shell command in the persistent isolated operator. " - "Use it to investigate and attack the scored TARGET_URL over the network." - ), - "parameters": { - "type": "object", - "properties": { - "command": { - "type": "string", - "description": "The exact bash command to execute.", - }, - "timeout_seconds": { - "type": "number", - "minimum": 1, - "maximum": 120, - }, - }, - "required": ["command"], - "additionalProperties": False, - }, - }, - }, - { - "type": "function", - "function": { - "name": "submit", - "description": ( - "Make the episode's one final submission. Call only after obtaining " - "real evidence from the scored target. The full arguments object is " - "written as submission.json and immediately finalized." - ), - "parameters": { - "type": "object", - "properties": { - "evidence": { - "type": "string", - "description": "Concrete evidence required by the task grader.", - }, - "flag": {"type": "string"}, - "class": {"type": "string"}, - "requests": { - "type": "array", - "items": {"type": "string"}, - }, - }, - "required": ["evidence"], - "additionalProperties": True, - }, - }, - }, -] - - -@dataclass -class AgentMetrics: - turns: int = 0 - tool_calls: int = 0 - terminal_calls: int = 0 - submit_calls: int = 0 - parse_failures: int = 0 - create_time: float = 0.0 - total_generation_time: float = 0.0 - total_tool_time: float = 0.0 - evaluate_time: float = 0.0 - close_time: float = 0.0 - timed_out: int = 0 - max_model_total_tokens: int = 0 - max_seq_len_hit: int = 0 - usage_missing: int = 0 - - def to_dict(self) -> dict[str, Any]: - return { - "turns": self.turns, - "tool_calls": self.tool_calls, - "terminal_calls": self.terminal_calls, - "submit_calls": self.submit_calls, - "parse_failures": self.parse_failures, - "create_time": self.create_time, - "total_generation_time": self.total_generation_time, - "total_tool_time": self.total_tool_time, - "evaluate_time": self.evaluate_time, - "close_time": self.close_time, - "timed_out": self.timed_out, - "max_model_total_tokens": self.max_model_total_tokens, - "max_seq_len_hit": self.max_seq_len_hit, - "usage_missing": self.usage_missing, - } - - -def _positive_env(name: str, default: float) -> float: - try: - value = float(os.getenv(name, str(default))) - except ValueError as exc: - raise ValueError(f"{name} must be numeric") from exc - if not math.isfinite(value) or value <= 0: - raise ValueError(f"{name} must be finite and positive") - return value - - -def _positive_int_env(name: str, default: int) -> int: - value = int(_positive_env(name, float(default))) - if value <= 0: - raise ValueError(f"{name} must be positive") - return value - - -def _metadata_max_seq_len(metadata: dict[str, Any]) -> int | None: - value = metadata.get("max_seq_len") - if value is None: - return None - if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise ValueError("sample metadata max_seq_len must be a positive integer") - return value - - -def _task_identity(metadata: dict[str, Any]) -> tuple[str, str]: - task_id = metadata.get("task_id") or metadata.get("cve") or metadata.get( - "task_name" - ) - if not isinstance(task_id, str) or not _TASK_ID.fullmatch(task_id): - raise ValueError("sample metadata must contain a canonical CVE task_id") - tier = metadata.get("prompt_tier", "l2") - if tier != "l2": - raise ValueError("secrlenv RL samples must use prompt_tier l2") - return task_id, tier - - -def _session_url(base_url: str) -> str: - value = f"{base_url.rstrip('/')}/v1" - external_host = os.getenv("MILES_ROUTER_EXTERNAL_HOST") - if external_host: - parsed = urlparse(value) - netloc = f"{external_host}:{parsed.port}" if parsed.port else external_host - value = urlunparse(parsed._replace(netloc=netloc)) - return value - - -def _tool_arguments(call: dict[str, Any]) -> tuple[str, dict[str, Any]]: - function = call.get("function") - if not isinstance(function, dict) or not isinstance(function.get("name"), str): - raise TypeError("tool call has no function name") - arguments = function.get("arguments", "{}") - if isinstance(arguments, str): - try: - arguments = json.loads(arguments) - except json.JSONDecodeError as exc: - raise ValueError("tool arguments are not valid JSON") from exc - if not isinstance(arguments, dict): - raise TypeError("tool arguments must be a JSON object") - return function["name"], arguments - - -def _tool_message(call_id: Any, value: dict[str, Any]) -> dict[str, str]: - identifier = call_id if isinstance(call_id, str) and call_id else "invalid_tool_call" - return { - "role": "tool", - "tool_call_id": identifier, - "content": json.dumps(value, ensure_ascii=False, separators=(",", ":")), - } - - -def _completion_total_tokens(completion: Any) -> int | None: - usage = getattr(completion, "usage", None) - if isinstance(usage, dict): - value = usage.get("total_tokens") - else: - value = getattr(usage, "total_tokens", None) - if isinstance(value, bool) or not isinstance(value, int) or value < 0: - return None - return value - - -def _terminal_observation(result: dict[str, Any], char_cap: int) -> dict[str, Any]: - output = result.get("output", "") - if not isinstance(output, str): - output = str(output) - locally_truncated = len(output) > char_cap - if locally_truncated: - output = output[:char_cap] + "\n[observation truncated by Miles adapter]" - return { - "exit_code": result.get("exit_code"), - "output": output or "(no output)", - "timed_out": bool(result.get("timed_out", False)), - "truncated": bool(result.get("truncated", False)) or locally_truncated, - } - - -async def _drive_policy( - policy: Any, - client: EpisodeClient, - episode: dict[str, Any], - request_kwargs: dict[str, Any], - metrics: AgentMetrics, - *, - max_seq_len: int | None = None, -) -> str: - episode_id = episode["episode_id"] - prompt = episode.get("prompt") - if not isinstance(prompt, str) or not prompt: - raise EpisodeClientError("episode daemon returned no task prompt") - messages: list[dict[str, Any]] = [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt}, - ] - max_turns = _positive_int_env("SECRLENV_MAX_TURNS", 40) - max_no_tool = _positive_int_env("SECRLENV_MAX_NO_TOOL_RETRIES", 3) - tool_timeout = _positive_env("SECRLENV_TOOL_TIMEOUT_SECONDS", 120.0) - output_bytes = _positive_int_env("SECRLENV_TOOL_OUTPUT_BYTES", 32768) - observation_chars = _positive_int_env("SECRLENV_OBSERVATION_CHARS", 16000) - model_name = os.getenv("AGENT_MODEL_NAME", "model") - no_tool_retries = 0 - sampling = { - key: value - for key, value in request_kwargs.items() - if key not in {"model", "messages", "tools", "tool_choice", "stream", "stream_options"} - } - - last_total_tokens = 0 - for turn in range(1, max_turns + 1): - if max_seq_len is not None and last_total_tokens >= max_seq_len: - metrics.max_seq_len_hit = 1 - return "max_seq_len" - metrics.turns = turn - turn_sampling = dict(sampling) - if max_seq_len is not None: - remaining = max_seq_len - last_total_tokens - configured_max = turn_sampling.get("max_tokens") - if configured_max is None: - turn_sampling["max_tokens"] = remaining - elif ( - isinstance(configured_max, bool) - or not isinstance(configured_max, int) - or configured_max <= 0 - ): - raise ValueError("model max_tokens must be a positive integer") - else: - turn_sampling["max_tokens"] = min(configured_max, remaining) - started = time.monotonic() - completion = await policy.chat.completions.create( - model=model_name, - messages=messages, - tools=TOOLS, - tool_choice="auto", - extra_body=turn_sampling, - ) - metrics.total_generation_time += time.monotonic() - started - observed_total_tokens = _completion_total_tokens(completion) - usage_missing = max_seq_len is not None and observed_total_tokens is None - if observed_total_tokens is not None: - last_total_tokens = observed_total_tokens - metrics.max_model_total_tokens = max( - metrics.max_model_total_tokens, - observed_total_tokens, - ) - elif usage_missing: - metrics.usage_missing += 1 - LOGGER.warning( - "model response omitted token usage while max_seq_len is enforced" - ) - message = completion.choices[0].message - assistant = message.model_dump(exclude_none=True) - messages.append(assistant) - calls = assistant.get("tool_calls") - if not isinstance(calls, list) or not calls: - no_tool_retries += 1 - metrics.parse_failures += 1 - if max_seq_len is not None and ( - usage_missing or last_total_tokens >= max_seq_len - ): - metrics.max_seq_len_hit = 1 - return "max_seq_len" - if no_tool_retries >= max_no_tool: - return "max_turns" - messages.append( - { - "role": "user", - "content": ( - "You must act through a structured terminal.exec or submit " - "tool call. Do not answer with prose or a bare command." - ), - } - ) - continue - no_tool_retries = 0 - submitted = False - for raw_call in calls: - metrics.tool_calls += 1 - call = raw_call if isinstance(raw_call, dict) else {} - call_id = call.get("id") - try: - name, arguments = _tool_arguments(call) - except ValueError as exc: - metrics.parse_failures += 1 - messages.append(_tool_message(call_id, {"error": str(exc)})) - continue - if name == "terminal.exec": - command = arguments.get("command") - if not isinstance(command, str) or not command.strip(): - messages.append( - _tool_message(call_id, {"error": "command must be non-empty"}) - ) - continue - requested_timeout = arguments.get("timeout_seconds", tool_timeout) - try: - requested_timeout = min(float(requested_timeout), tool_timeout) - except (TypeError, ValueError): - requested_timeout = tool_timeout - started = time.monotonic() - try: - result = await client.execute( - episode_id, - command, - timeout_seconds=max(1.0, requested_timeout), - output_bytes=output_bytes, - ) - except EpisodeAPIError as exc: - if exc.status == 400: - messages.append(_tool_message(call_id, {"error": exc.message})) - continue - raise - finally: - metrics.total_tool_time += time.monotonic() - started - metrics.terminal_calls += 1 - messages.append( - _tool_message( - call_id, _terminal_observation(result, observation_chars) - ) - ) - elif name == "submit": - started = time.monotonic() - try: - result = await client.submit(episode_id, arguments) - except EpisodeAPIError as exc: - if exc.status == 400: - messages.append(_tool_message(call_id, {"error": exc.message})) - continue - raise - finally: - metrics.total_tool_time += time.monotonic() - started - metrics.submit_calls += 1 - messages.append(_tool_message(call_id, result)) - submitted = True - else: - messages.append( - _tool_message(call_id, {"error": f"unknown tool {name!r}"}) - ) - if submitted: - return "completed" - if max_seq_len is not None and ( - usage_missing or last_total_tokens >= max_seq_len - ): - metrics.max_seq_len_hit = 1 - return "max_seq_len" - return "max_turns" - - -def _validated_outcome( - result: dict[str, Any], - task_id: str, - episode_id: str, - status: str, - *, - infrastructure_503_retries: int = 0, -) -> dict[str, Any]: - if result.get("task_id") != task_id or result.get("episode_id") != episode_id: - raise EpisodeClientError("evaluation response identity mismatch") - reward = result.get("reward") - passed = result.get("passed") - if ( - isinstance(reward, bool) - or not isinstance(reward, (int, float)) - or not math.isfinite(float(reward)) - or not 0.0 <= float(reward) <= 1.0 - or not isinstance(passed, bool) - ): - raise EpisodeClientError("evaluation response has an invalid verdict") - grader = result.get("grader") if isinstance(result.get("grader"), dict) else {} - return { - "schema": 1, - "status": status, - "episode_id": episode_id, - "task_id": task_id, - "reward": float(reward), - "passed": passed, - "class": grader.get("class") if isinstance(grader.get("class"), str) else None, - INFRASTRUCTURE_503_RETRIES_KEY: infrastructure_503_retries, - } - - -def _infrastructure_outcome( - task_id: str, - episode_id: str, - infrastructure_503_retries: int = 0, -) -> dict[str, Any]: - """Build an abort signal, never a synthetic grader verdict.""" - - return { - "schema": 1, - "status": INFRASTRUCTURE_STATUS, - "episode_id": episode_id, - "task_id": task_id, - "reward": 0.0, - "passed": False, - "class": None, - INFRASTRUCTURE_503_RETRIES_KEY: infrastructure_503_retries, - } - - -def _admission_infrastructure_outcome( - task_id: str, - infrastructure_503_retries: int, - *, - nonce: str | None = None, -) -> dict[str, Any]: - """Build evidence for one exhausted authenticated pre-create 503.""" - - if infrastructure_503_retries != 1: - raise ValueError("invalid SecRLEnv admission retry ledger") - return { - "schema": ADMISSION_SCHEMA, - "status": INFRASTRUCTURE_STATUS, - "episode_id": None, - "task_id": task_id, - "reward": 0.0, - "passed": False, - "class": None, - INFRASTRUCTURE_503_RETRIES_KEY: infrastructure_503_retries, - ADMISSION_PHASE_KEY: ADMISSION_PHASE, - ADMISSION_ERROR_CODE_KEY: ADMISSION_ERROR_CODE, - ADMISSION_NONCE_KEY: nonce or secrets.token_hex(16), - } - - -def _agent_result( - outcome: dict[str, Any], outcome_status: str, metrics: AgentMetrics -) -> dict[str, Any]: - result = { - OUTCOME_KEY: outcome, - "exit_status": outcome_status, - "agent_metrics": metrics.to_dict(), - } - capture_agent_metadata(result) - return result - - -def _cleanup_error_outcome( - task_id: str, - episode_id: str, - infrastructure_503_retries: int = 0, -) -> dict[str, Any]: - """Build a terminal signal when episode release is unproven.""" - - outcome = _infrastructure_outcome( - task_id, - episode_id, - infrastructure_503_retries, - ) - outcome["status"] = CLEANUP_ERROR_STATUS - return outcome - - -async def _create_with_capacity_retry( - client: EpisodeClient, - task_id: str, - tier: str, - *, - retry_ledger: dict[str, int] | None = None, -) -> dict[str, Any]: - max_wait = _positive_env("SECRLENV_CAPACITY_MAX_WAIT_SECONDS", 14400.0) - started = time.monotonic() - deadline = started + max_wait - attempts = 0 - if retry_ledger is None: - retry_ledger = {INFRASTRUCTURE_503_RETRIES_KEY: 0} - infrastructure_retries = retry_ledger.get(INFRASTRUCTURE_503_RETRIES_KEY) - if type(infrastructure_retries) is not int or infrastructure_retries not in {0, 1}: - raise ValueError("invalid SecRLEnv infrastructure retry ledger") - next_log = started - last_retryable: EpisodeAPIError | None = None - while True: - remaining: float | None = None - if last_retryable is not None: - remaining = deadline - time.monotonic() - if remaining <= 0: - raise last_retryable - attempts += 1 - try: - if remaining is None: - result = await client.create(task_id, tier) - else: - create_task = asyncio.create_task(client.create(task_id, tier)) - try: - result = await asyncio.wait_for(create_task, timeout=remaining) - except asyncio.TimeoutError: - if last_retryable.code == ADMISSION_ERROR_CODE: - raise EpisodeClientError( - "secrlenv episode admission retry timed out" - ) from None - raise last_retryable from None - except EpisodeAPIError as exc: - now = time.monotonic() - if exc.status == 503 and exc.code == "capacity_reached": - reason = "capacity" - elif exc.status == 503 and exc.code == ADMISSION_ERROR_CODE: - if infrastructure_retries >= 1: - raise _AdmissionInfrastructureExhausted() from None - infrastructure_retries += 1 - retry_ledger[INFRASTRUCTURE_503_RETRIES_KEY] = infrastructure_retries - reason = "infrastructure" - else: - raise - if now >= deadline: - raise - last_retryable = exc - else: - # Once create returns an episode it must be tracked and closed. A - # deadline-edge success cannot be discarded as a stale prior 503. - return result - if now >= next_log: - LOGGER.warning( - "waiting for secrlenv episode daemon reason=%s attempts=%d elapsed_s=%.1f max_wait_s=%.1f", - reason, - attempts, - now - started, - max_wait, - ) - next_log = now + 15.0 - await asyncio.sleep(min(random.uniform(0.5, 3.0), max(0.0, deadline - now))) - - -async def run( - base_url: str, - prompt: Any, - request_kwargs: dict[str, Any] | None = None, - metadata: dict[str, Any] | None = None, - **_kwargs: Any, -) -> dict[str, Any] | None: - """Run one model-driven secrlenv episode and return reward metadata.""" - - del prompt # The authenticated daemon supplies the immutable task prompt. - metadata = dict(metadata or {}) - request_kwargs = dict(request_kwargs or {}) - try: - task_id, tier = _task_identity(metadata) - max_seq_len = _metadata_max_seq_len(metadata) - infrastructure_503_retries = metadata.get( - INFRASTRUCTURE_503_RETRIES_KEY, 0 - ) - if ( - type(infrastructure_503_retries) is not int - or infrastructure_503_retries not in {0, 1} - ): - raise ValueError("invalid SecRLEnv infrastructure retry ledger") - except ValueError as exc: - LOGGER.error("invalid secrlenv sample metadata: %s", exc) - return None - - metrics = AgentMetrics() - episode_id: str | None = None - outcome_status = "completed" - infrastructure_failure = False - cleanup_failure = False - evaluation: dict[str, Any] | None = None - policy = None - policy_task: asyncio.Task[str] | None = None - retry_ledger = { - INFRASTRUCTURE_503_RETRIES_KEY: infrastructure_503_retries - } - try: - from openai import AsyncOpenAI - - policy = AsyncOpenAI(base_url=_session_url(base_url), api_key="EMPTY") - async with EpisodeClient() as client: - started = time.monotonic() - episode = await _create_with_capacity_retry( - client, task_id, tier, retry_ledger=retry_ledger - ) - metrics.create_time = time.monotonic() - started - episode_id = episode.get("episode_id") - if not isinstance(episode_id, str): - raise EpisodeClientError("episode daemon returned no episode ID") - - rollout_timeout = _positive_env( - "SECRLENV_MAX_ROLLOUT_TIME_SECONDS", 3600.0 - ) - policy_task = asyncio.create_task( - _await_policy_and_claim_finalization( - episode_id, - _drive_policy( - policy, - client, - episode, - request_kwargs, - metrics, - max_seq_len=max_seq_len, - ), - ) - ) - _register_episode(episode_id, policy_task) - try: - outcome_status = await asyncio.wait_for( - policy_task, - timeout=rollout_timeout, - ) - except asyncio.TimeoutError: - metrics.timed_out = 1 - outcome_status = "timeout" - LOGGER.warning( - "secrlenv episode %s exceeded %.1fs; evaluating partial state", - episode_id, - rollout_timeout, - ) - except asyncio.CancelledError: - raise - except Exception: - infrastructure_failure = True - LOGGER.exception("secrlenv policy/tool loop failed") - - if _claim_episode_finalization(episode_id): - started = time.monotonic() - try: - evaluation = await client.evaluate(episode_id) - except Exception: - infrastructure_failure = True - LOGGER.exception("secrlenv evaluation failed") - finally: - metrics.evaluate_time = time.monotonic() - started - - started = time.monotonic() - try: - await client.close(episode_id) - except asyncio.CancelledError: - raise - except EpisodeAPIError as exc: - if exc.status == 404: - _release_episode(episode_id) - else: - infrastructure_failure = True - LOGGER.exception("secrlenv teardown failed") - cleanup_failure = not await _recover_failed_normal_close( - client, episode_id, policy_task - ) - except Exception: - infrastructure_failure = True - LOGGER.exception("secrlenv teardown failed") - cleanup_failure = not await _recover_failed_normal_close( - client, episode_id, policy_task - ) - else: - _release_episode(episode_id) - finally: - metrics.close_time = time.monotonic() - started - metrics.total_tool_time += ( - metrics.create_time - + metrics.evaluate_time - + metrics.close_time - ) - else: - infrastructure_failure = True - LOGGER.warning( - "secrlenv episode was claimed by the abort hook before evaluation" - ) - except asyncio.CancelledError: - LOGGER.warning("secrlenv rollout cancelled") - cleanup_claimed = episode_id is not None and ( - _claim_episode_cleanup(episode_id) - or _claim_cancelled_policy_cleanup(episode_id, policy_task) - ) - if cleanup_claimed: - try: - async with EpisodeClient(total_timeout_seconds=180.0) as cleanup_client: - retry_timeout = _positive_env( - "SECRLENV_ABORT_CLOSE_RETRY_SECONDS", 180.0 - ) - await _close_aborted_episode( - cleanup_client, - episode_id, - retry_timeout_seconds=retry_timeout, - ) - except Exception: - LOGGER.warning( - "cancelled episode cleanup failed for %s", episode_id, exc_info=True - ) - finally: - _mark_episode_cleanup_pending(episode_id) - raise - except _AdmissionInfrastructureExhausted: - if ( - episode_id is None - and retry_ledger[INFRASTRUCTURE_503_RETRIES_KEY] == 1 - ): - try: - outcome = _admission_infrastructure_outcome( - task_id, - retry_ledger[INFRASTRUCTURE_503_RETRIES_KEY], - ) - return _agent_result(outcome, INFRASTRUCTURE_STATUS, metrics) - except Exception: - LOGGER.exception("refusing untrusted secrlenv admission metadata") - return None - LOGGER.error("invalid secrlenv admission infrastructure state") - return None - except EpisodeAPIError: - LOGGER.exception("secrlenv episode failed before a trustworthy verdict") - return None - except Exception: - LOGGER.exception("secrlenv episode failed before a trustworthy verdict") - return None - finally: - if policy is not None: - try: - await policy.close() - except Exception: - LOGGER.warning("failed to close policy client", exc_info=True) - - if episode_id is None: - return None - try: - if cleanup_failure: - outcome_status = CLEANUP_ERROR_STATUS - outcome = _cleanup_error_outcome( - task_id, - episode_id, - retry_ledger[INFRASTRUCTURE_503_RETRIES_KEY], - ) - elif infrastructure_failure or evaluation is None: - outcome_status = INFRASTRUCTURE_STATUS - outcome = _infrastructure_outcome( - task_id, - episode_id, - retry_ledger[INFRASTRUCTURE_503_RETRIES_KEY], - ) - else: - outcome = _validated_outcome( - evaluation, - task_id, - episode_id, - outcome_status, - infrastructure_503_retries=retry_ledger[ - INFRASTRUCTURE_503_RETRIES_KEY - ], - ) - result = _agent_result(outcome, outcome_status, metrics) - except Exception: - LOGGER.exception("refusing untrusted secrlenv evaluation metadata") - return None - return result - - -async def abort(_args: Any = None) -> None: - """Miles oversampling abort hook: tear down every in-flight local episode.""" - - claimed = _claim_episodes_and_tasks_for_abort() - if not claimed: - return - try: - retry_timeout = _positive_env("SECRLENV_ABORT_CLOSE_RETRY_SECONDS", 180.0) - drain_timeout = _positive_env("SECRLENV_ABORT_DRAIN_TIMEOUT_SECONDS", 30.0) - drained = await asyncio.gather( - *( - _drain_aborted_policy_task( - task, - timeout_seconds=drain_timeout, - ) - for _episode_id, task in claimed - ) - ) - close_ids = [] - for (episode_id, _task), task_drained in zip(claimed, drained, strict=True): - if task_drained: - close_ids.append(episode_id) - if not close_ids: - return - async with EpisodeClient(total_timeout_seconds=180.0) as client: - await asyncio.gather( - *( - _close_aborted_episode( - client, - episode_id, - retry_timeout_seconds=retry_timeout, - ) - for episode_id in close_ids - ) - ) - except asyncio.CancelledError: - raise - except Exception: - LOGGER.warning("secrlenv abort cleanup failed", exc_info=True) - finally: - for episode_id, _task in claimed: - _mark_episode_cleanup_pending(episode_id) diff --git a/yeto/rl/harness/codex/alignment.py b/yeto/rl/harness/codex/alignment.py deleted file mode 100644 index 0432d100c..000000000 --- a/yeto/rl/harness/codex/alignment.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Mask / token / logprob alignment contract for agentic samples (design D5 assertions). - -A sample produced by the session server (``compute_samples_from_openai_records`` -+ ``merge_samples``) must satisfy, fail closed: - -- ``len(tokens) == len(loss_mask)`` (loss_mask covers the whole sequence), or - ``len(loss_mask) == response_length`` (per-turn layout before merge); -- ``rollout_log_probs`` has one entry per response token; -- every mask=1 position lies in the response region and has a finite logprob; -- mask=1 positions are exactly the generated spans when ``weight_versions`` - spans are present (they delimit SGLang output); tool/template tokens are mask 0. - -Violations raise ``AlignmentError``; callers abort the trajectory and never -convert the failure into a 0 reward. -""" - -from __future__ import annotations - -import math -from typing import Any - - -class AlignmentError(ValueError): - pass - - -def _response_offset(sample: Any) -> tuple[int, int]: - tokens = list(getattr(sample, "tokens", None) or []) - response_length = getattr(sample, "response_length", None) - if not isinstance(response_length, int) or response_length < 0 or response_length > len(tokens): - raise AlignmentError("response_length is missing or out of range") - return len(tokens) - response_length, len(tokens) - - -def generated_spans(sample: Any) -> list[tuple[int, int]]: - spans: list[tuple[int, int]] = [] - for call in getattr(sample, "weight_versions", None) or []: - for span in getattr(call, "spans", None) or (call.get("spans", []) if isinstance(call, dict) else []): - start = getattr(span, "abs_start", None) if not isinstance(span, dict) else span.get("abs_start") - end = getattr(span, "abs_end", None) if not isinstance(span, dict) else span.get("abs_end") - if not isinstance(start, int) or not isinstance(end, int) or start < 0 or end < start: - raise AlignmentError("weight version span is malformed") - spans.append((start, end)) - return spans - - -def assert_sample_alignment(sample: Any) -> None: - tokens = list(getattr(sample, "tokens", None) or []) - loss_mask = getattr(sample, "loss_mask", None) - logprobs = getattr(sample, "rollout_log_probs", None) - if loss_mask is None: - raise AlignmentError("loss_mask is missing") - if logprobs is None: - raise AlignmentError("rollout_log_probs is missing") - start, end = _response_offset(sample) - if len(loss_mask) == len(tokens): - mask_offset = 0 - elif len(loss_mask) == end - start: - mask_offset = start - else: - raise AlignmentError(f"len(loss_mask)={len(loss_mask)} matches neither len(tokens)={len(tokens)} nor response_length={end - start}") - if len(logprobs) != end - start: - raise AlignmentError(f"len(rollout_log_probs)={len(logprobs)} != response_length={end - start}") - ones = {mask_offset + i for i, m in enumerate(loss_mask) if int(m) == 1} - for position in sorted(ones): - if position < start: - raise AlignmentError(f"loss_mask=1 at prompt position {position}") - value = logprobs[position - start] - if value is None or not isinstance(value, (int, float)) or not math.isfinite(float(value)): - raise AlignmentError(f"loss_mask=1 at {position} without a finite rollout logprob") - spans = generated_spans(sample) - if spans: - generated = {p for s, e in spans for p in range(s, e)} - if ones - generated: - raise AlignmentError(f"loss_mask=1 outside generated spans at {sorted(ones - generated)[:5]}") diff --git a/yeto/rl/harness/codex/chain_judge.py b/yeto/rl/harness/codex/chain_judge.py deleted file mode 100644 index b9c97bd8a..000000000 --- a/yeto/rl/harness/codex/chain_judge.py +++ /dev/null @@ -1,78 +0,0 @@ -"""R-D5a chain-count judge for Codex GPU runs (rl-codex-harness-rollout design 9.2). - -The GPU 9.2 hard criterion used to be stated only in prose: every trajectory -has ``chains_total == 1``. With CompactionRL rollout compaction on -(``YETO_CODEX_COMPACTIONRL``, progress.md "S13 Codex 压缩接线") a trajectory is -deliberately split into one sample per compaction segment, so the criterion -becomes, per trajectory: - -- ``chains_total == num_segments == compactions + 1 == number of samples``; -- ``chain_index`` / ``segment_index`` are exactly ``0..n-1``; -- ``chain_break_reason`` is ``compaction_window`` on every sample but the first - (None there). - -With compaction off the criterion is unchanged: ``chains_total == 1`` on every -sample (and one sample per trajectory). Aborted (infrastructure) samples are -judged the same way: the generate wrapper keeps an aborted compacted -trajectory as its single segment-0 sample, which satisfies the off rule. -""" - -from __future__ import annotations - -from collections.abc import Iterable, Mapping -from typing import Any - -COMPACTION_BREAK_REASON = "compaction_window" - - -def _key(meta: Mapping[str, Any], position: int) -> str: - trajectory = meta.get("trajectory_id") - return str(trajectory) if trajectory not in (None, "") else f"" - - -def judge_chain_counts( - samples_metadata: Iterable[Mapping[str, Any]], - *, - compaction_enabled: bool, -) -> list[str]: - """Return the violations (empty list = criterion holds).""" - groups: dict[str, list[Mapping[str, Any]]] = {} - for position, meta in enumerate(samples_metadata): - groups.setdefault(_key(meta, position), []).append(meta) - problems: list[str] = [] - for trajectory, metas in groups.items(): - totals = {m.get("chains_total") for m in metas} - if not compaction_enabled: - if totals != {1} or len(metas) != 1: - problems.append( - f"{trajectory}: chains_total {sorted(map(str, totals))} over {len(metas)} sample(s); " - "expected exactly one sample with chains_total==1" - ) - continue - # Aborted compacted trajectories collapse to one sample (no segment keys). - if len(metas) == 1 and metas[0].get("num_segments") is None: - if totals != {1}: - problems.append(f"{trajectory}: single sample with chains_total {sorted(map(str, totals))}") - continue - n = len(metas) - segments = {m.get("num_segments") for m in metas} - compactions = {m.get("compactions") for m in metas} - if totals != {n} or segments != {n}: - problems.append( - f"{trajectory}: {n} sample(s) but chains_total {sorted(map(str, totals))}, " - f"num_segments {sorted(map(str, segments))}" - ) - if compactions != {n - 1}: - problems.append( - f"{trajectory}: compactions {sorted(map(str, compactions))} != num_segments-1 ({n - 1})" - ) - order = sorted(metas, key=lambda m: (m.get("chain_index") is None, m.get("chain_index") or 0)) - if [m.get("chain_index") for m in order] != list(range(n)) or [ - m.get("segment_index") for m in order - ] != list(range(n)): - problems.append(f"{trajectory}: chain_index/segment_index are not 0..{n - 1}") - continue - reasons = [m.get("chain_break_reason") for m in order] - if reasons != [None] + [COMPACTION_BREAK_REASON] * (n - 1): - problems.append(f"{trajectory}: chain_break_reason {reasons}") - return problems diff --git a/yeto/rl/harness/codex/client.py b/yeto/rl/harness/codex/client.py deleted file mode 100644 index 138059fca..000000000 --- a/yeto/rl/harness/codex/client.py +++ /dev/null @@ -1,215 +0,0 @@ -"""Authenticated async client for the loopback secrlenv episode daemon.""" - -from __future__ import annotations - -import json -import math -import os -import stat -from pathlib import Path -from typing import Any, Self -from urllib.error import HTTPError, URLError -from urllib.parse import urlsplit -from urllib.request import urlopen - -import aiohttp - - -class EpisodeClientError(RuntimeError): - pass - - -class EpisodeAPIError(EpisodeClientError): - def __init__(self, status: int, code: str, message: str) -> None: - super().__init__(f"episode daemon returned {status} {code}: {message}") - self.status = status - self.code = code - self.message = message - - -class EpisodeTransportError(EpisodeClientError): - pass - - -def _read_token() -> str: - token_file = os.getenv("SECRLENV_BEARER_TOKEN_FILE") - if token_file: - path = Path(token_file) - try: - info = path.stat() - if not stat.S_ISREG(info.st_mode) or info.st_mode & 0o077: - raise EpisodeClientError( - "SECRLENV_BEARER_TOKEN_FILE must be a private regular file" - ) - token = path.read_text(encoding="utf-8").strip() - except OSError as exc: - raise EpisodeClientError(f"cannot read episode daemon token file: {exc}") from exc - else: - token = os.getenv("SECRLENV_BEARER_TOKEN", "").strip() - if len(token) < 32: - raise EpisodeClientError("episode daemon bearer token is missing or too short") - return token - - -def daemon_url_from_env() -> str: - value = os.getenv("SECRLENV_DAEMON_URL", "http://127.0.0.1:8765").rstrip("/") - parsed = urlsplit(value) - if ( - parsed.scheme != "http" - or parsed.hostname not in {"127.0.0.1", "::1", "localhost"} - or parsed.username is not None - or parsed.password is not None - or parsed.path not in {"", "/"} - or parsed.query - or parsed.fragment - or parsed.port is None - ): - raise EpisodeClientError( - "SECRLENV_DAEMON_URL must be an unauthenticated loopback HTTP origin" - ) - return value - - -def require_daemon_ready(*, timeout_seconds: float = 10.0) -> None: - """Fail before GPU initialization unless the expected daemon is healthy.""" - - if not math.isfinite(timeout_seconds) or timeout_seconds <= 0: - raise ValueError("episode daemon readiness timeout must be positive") - try: - with urlopen( - f"{daemon_url_from_env()}/healthz", timeout=timeout_seconds - ) as response: - if response.status != 200: - raise EpisodeTransportError( - f"secrlenv episode daemon health check returned {response.status}" - ) - value = json.loads(response.read()) - except EpisodeTransportError: - raise - except (HTTPError, URLError, OSError, json.JSONDecodeError) as exc: - raise EpisodeTransportError( - "secrlenv episode daemon is not healthy on its configured loopback origin" - ) from exc - if not isinstance(value, dict) or value.get("ok") is not True: - raise EpisodeTransportError("secrlenv episode daemon health response is invalid") - expected_task_pack = os.getenv("SECRLENV_TASK_PACK_SHA256", "").strip() - if expected_task_pack and value.get("task_pack_sha256") != expected_task_pack: - raise EpisodeTransportError( - "secrlenv episode daemon task-pack identity does not match the run plan" - ) - - -class EpisodeClient: - def __init__( - self, - base_url: str | None = None, - token: str | None = None, - *, - total_timeout_seconds: float = 1200.0, - ) -> None: - self.base_url = (base_url or daemon_url_from_env()).rstrip("/") - # Apply the same validation to explicit constructor values used in tests - # and direct-drive tools, not only to the environment-derived default. - parsed = urlsplit(self.base_url) - if ( - parsed.scheme != "http" - or parsed.hostname not in {"127.0.0.1", "::1", "localhost"} - or parsed.username is not None - or parsed.password is not None - or parsed.path not in {"", "/"} - or parsed.query - or parsed.fragment - or parsed.port is None - ): - raise EpisodeClientError("episode daemon URL must be a loopback HTTP origin") - self.token = token or _read_token() - if len(self.token) < 32: - raise EpisodeClientError("episode daemon bearer token is too short") - self.timeout = aiohttp.ClientTimeout(total=total_timeout_seconds) - self._session: aiohttp.ClientSession | None = None - - async def __aenter__(self) -> Self: - self._session = aiohttp.ClientSession( - timeout=self.timeout, - connector=aiohttp.TCPConnector(limit=4), - headers={"Authorization": f"Bearer {self.token}"}, - ) - return self - - async def __aexit__(self, *_exc: object) -> None: - if self._session is not None: - await self._session.close() - self._session = None - - async def _request( - self, - method: str, - path: str, - body: dict[str, Any] | None = None, - ) -> dict[str, Any]: - if self._session is None: - raise EpisodeClientError("EpisodeClient must be used as an async context manager") - try: - async with self._session.request( - method, - f"{self.base_url}{path}", - json=body, - headers={"Content-Type": "application/json"}, - ) as response: - raw = await response.text() - except (aiohttp.ClientError, TimeoutError) as exc: - raise EpisodeTransportError(f"episode daemon request failed: {exc}") from exc - try: - value = json.loads(raw) - except json.JSONDecodeError as exc: - raise EpisodeTransportError("episode daemon returned invalid JSON") from exc - if not isinstance(value, dict): - raise EpisodeTransportError("episode daemon returned a non-object response") - if not 200 <= response.status < 300: - error = value.get("error") if isinstance(value.get("error"), dict) else {} - raise EpisodeAPIError( - response.status, - str(error.get("code", "unknown_error")), - str(error.get("message", "request failed")), - ) - return value - - async def create(self, task_id: str, prompt_tier: str) -> dict[str, Any]: - return await self._request( - "POST", - "/v1/episodes", - {"task_id": task_id, "prompt_tier": prompt_tier}, - ) - - async def execute( - self, - episode_id: str, - command: str, - *, - timeout_seconds: float, - output_bytes: int, - ) -> dict[str, Any]: - return await self._request( - "POST", - f"/v1/episodes/{episode_id}/exec", - { - "command": command, - "timeout_seconds": timeout_seconds, - "output_bytes": output_bytes, - }, - ) - - async def submit(self, episode_id: str, submission: dict[str, Any]) -> dict[str, Any]: - return await self._request( - "POST", - f"/v1/episodes/{episode_id}/submit", - {"submission": submission}, - ) - - async def evaluate(self, episode_id: str) -> dict[str, Any]: - return await self._request( - "POST", f"/v1/episodes/{episode_id}/evaluate", {} - ) - - async def close(self, episode_id: str) -> dict[str, Any]: - return await self._request("DELETE", f"/v1/episodes/{episode_id}") diff --git a/yeto/rl/harness/codex/codex_harness_agent.py b/yeto/rl/harness/codex/codex_harness_agent.py deleted file mode 100644 index aab59eb2e..000000000 --- a/yeto/rl/harness/codex/codex_harness_agent.py +++ /dev/null @@ -1,2389 +0,0 @@ -"""Pinned stock-Codex harness for isolated Miles/SecRLEnv episodes. - -The stock Codex app-server owns the agent loop. A per-episode loopback bridge -adapts its Responses API calls to the one Miles session-server Chat Completions -endpoint so Miles remains the sole sampler and logprob/TITO owner. -""" - -from __future__ import annotations - -import asyncio -import copy -import errno -import hashlib -import hmac -import json -import logging -import math -import os -import re -import secrets -import shutil -import signal -import stat -import subprocess -import tempfile -import time -from dataclasses import dataclass -from pathlib import Path -from threading import Lock -from typing import Any, Self - -import aiohttp -from aiohttp import web - -from yeto.rl.codex_backend import stock_codex_backend_profile - -from . import agent as legacy -from .client import EpisodeAPIError, EpisodeClient, EpisodeClientError -from .generate import capture_agent_metadata -from .reward import ( - INFRASTRUCTURE_503_RETRIES_KEY, - INFRASTRUCTURE_STATUS, - OUTCOME_KEY, -) - -LOGGER = logging.getLogger(__name__) - -CODEX_APP_SERVER_PROTOCOL_REVISION = "v2" -CODEX_APP_SERVER_SCHEMA_SHA256 = ( - "a88d865c3ca41fc63672baf423e28b3bfdb85b989e1c95e6e268932daaac91a0" -) -CODEX_CLI_VERSION = "codex-cli 0.145.0" -_BACKEND_PROFILE = stock_codex_backend_profile( - os.getenv("YETO_CODEX_CHAT_TEMPLATE", "deepseekv4") -) -BACKEND_MODEL = _BACKEND_PROFILE["model"] -BACKEND_REASONING_EFFORT = _BACKEND_PROFILE["backend_reasoning_effort"] -BACKEND_CHAT_TEMPLATE_KWARGS = _BACKEND_PROFILE["chat_template_kwargs"] -# rlf-codex-rl 5.6: the tool-output budget is a profile field. The -# qwen3_30b_a3b cyber baseline pins 8192 bytes; profiles without a declaration -# keep the signed 32 KiB default. -MAX_TOOL_OUTPUT_BYTES = int(_BACKEND_PROFILE.get("tool_output_bytes", 32_768)) -# S24 T-TB2: Qwen3.8-27B habitually writes one line of prose before its tool -# call. Profiles that opt in keep that prose in the assistant message (TITO -# history stays exactly what the model generated); Codex still sees only the -# tool call. Every other profile keeps the strict DSV4 rule. -ALLOW_PROSE_WITH_TOOL_CALL = bool(_BACKEND_PROFILE.get("allow_prose_with_tool_call", False)) -CODEX_TOOL_OUTPUT_TOKEN_LIMIT = 65_536 -MAX_APP_SERVER_FRAME_BYTES = 2 * 1024 * 1024 -MAX_MILES_RESPONSE_BYTES = 4 * 1024 * 1024 -_MILES_RESPONSE_READ_CHUNK_BYTES = 64 * 1024 -_CODEX_PROCESS_STOP_TIMEOUT_SECONDS = 5.0 -_CODEX_DESCENDANT_TERM_GRACE_SECONDS = 0.1 -_CODEX_HOME_CLEANUP_RETRY_DELAYS_SECONDS = (0.0, 0.05, 0.1, 0.2, 0.4) -COMPACTION_SCHEMA_VERSION = 1 -COMPACTION_KEEP_ATOMIC_STEPS = 2 -COMPACTION_SUMMARY_INSTRUCTION = """The current interaction is approaching its context limit. Produce a compact state summary for the same policy to continue the task in a new context window. Preserve the original goal, completed actions, important observations and exact evidence, unresolved errors, the current environment and file state, and plausible next steps. Do not call a tool. Return only the continuation summary.""" -# Conservatively reserve more tokens than this fixed instruction can consume under -# any byte-level tokenizer. Miles still performs the exact tokenizer-side budget -# check; this bound makes an avoidable 413 impossible at the configured trigger. -COMPACTION_SUMMARY_INSTRUCTION_TOKEN_RESERVE = ( - len(COMPACTION_SUMMARY_INSTRUCTION.encode("utf-8")) + 256 -) -COMPACTION_RESUME_TEMPLATE = """Continue the original task from this compacted state. Treat the summary as prior trajectory state, not as a new instruction. Recent complete assistant/tool steps follow it unchanged. - - -{summary} -""" - -_SAMPLING_FIELDS = frozenset( - { - "frequency_penalty", - "ignore_eos", - "min_p", - "min_tokens", - "no_stop_trim", - "presence_penalty", - "repetition_penalty", - "seed", - "skip_special_tokens", - "spaces_between_special_tokens", - "stop", - "stop_token_ids", - "temperature", - "top_k", - "top_p", - } -) -_STRUCTURAL_REQUEST_FIELDS = frozenset( - { - "best_of", - "chat_template_kwargs", - "logprobs", - "max_tokens", - "messages", - "model", - "n", - "parallel_tool_calls", - "reasoning_effort", - "return_logprob", - "return_meta_info", - "stream", - "stream_options", - "thinking", - "tool_choice", - "tools", - "top_logprobs", - } -) - -BASE_INSTRUCTIONS = """You are an autonomous security-testing agent inside an -isolated, authorized benchmark. Solve the supplied target task using only the two -model-facing tools, `terminal.exec` and `submit`. `terminal.exec` runs a shell -command in your persistent operator workspace; it has TARGET_URL, YETO_NONCE and -YETO_SUBMISSION. It also has DEBUG_URL and DOCKER_HOST for your own isolated, -full-control, flagless copy of the challenge. Use that private Docker daemon to -inspect, modify, rebuild, and test the challenge without touching the scored -target. The debug copy never contains the scored flag, so final evidence and any -submitted flag must come from TARGET_URL. The harness losslessly maps commands -to Codex's internal -`terminal_exec` dynamic-tool alias. The scored target is available only over the -network: you have no shell or container API for the scored target. Do not merely -describe commands: call `terminal.exec`. When you have concrete evidence, call -`submit` exactly once with the final JSON submission. Do not invent tool results -or claim success without evidence. The `submit` call is terminal: make no further -model or tool calls.""" - -TERMINAL_EXEC_TOOL = { - "type": "function", - "name": "terminal_exec", - "description": ( - "Execute one shell command in the persistent isolated operator. " - "Use it to investigate and attack the scored TARGET_URL over the network." - ), - "inputSchema": { - "type": "object", - "properties": { - "command": { - "type": "string", - "description": "The exact bash command to execute.", - }, - "timeout_seconds": { - "type": "number", - "minimum": 1, - "maximum": 120, - }, - }, - "required": ["command"], - "additionalProperties": False, - }, -} - -SUBMIT_TOOL = { - "type": "function", - "name": "submit", - "description": ( - "Make the episode's one final submission. Call only after obtaining " - "real evidence from the scored target. The full arguments object is " - "written as submission.json and immediately finalized." - ), - "inputSchema": { - "type": "object", - "properties": { - "evidence": { - "type": "string", - "description": "Concrete evidence required by the task grader.", - }, - "flag": {"type": "string"}, - "class": {"type": "string"}, - "requests": { - "type": "array", - "items": {"type": "string"}, - }, - }, - "required": ["evidence"], - "additionalProperties": True, - }, -} - -DYNAMIC_TOOLS = [TERMINAL_EXEC_TOOL, SUBMIT_TOOL] - -# 2026-10-07 (S15 r2 follow-up): the legacy tool descriptions above are CTF -# wording ("attack the scored TARGET_URL", "evidence from the scored target"). -# The TB2 family keeps names and input schemas and only rewords descriptions. -TB2_TERMINAL_EXEC_TOOL = { - **TERMINAL_EXEC_TOOL, - "description": ( - "Execute one shell command in the task container (the graded " - "environment), starting in the task workdir. Returns exit code and output." - ), -} -TB2_SUBMIT_TOOL = { - **SUBMIT_TOOL, - "description": ( - "Finish the episode after the task is done and verified. The hidden " - "tests then grade the container state. Call exactly once." - ), -} -TB2_DYNAMIC_TOOLS = [TB2_TERMINAL_EXEC_TOOL, TB2_SUBMIT_TOOL] - - -def _canonical_json(value: Any) -> bytes: - return json.dumps( - value, - ensure_ascii=False, - allow_nan=False, - sort_keys=True, - separators=(",", ":"), - ).encode("utf-8") - - -def _sha256_json(value: Any) -> str: - return hashlib.sha256(_canonical_json(value)).hexdigest() - - -BASE_INSTRUCTIONS_SHA256 = hashlib.sha256(BASE_INSTRUCTIONS.encode()).hexdigest() - -# 2026-10-07 (S15 stage 2): Terminal-Bench 2 tasks get their own system prompt. -# The legacy text above is a security-testing/CTF prompt (TARGET_URL, flags, -# "no shell for the scored target"); in TB2 the shell *is* the scored task -# container. Selected per episode by the TB2 environment provider through -# ``YETO_CODEX_INSTRUCTIONS_FAMILY=tb2`` plus the signed -# ``YETO_CODEX_TB2_BASE_INSTRUCTIONS_SHA256`` (worker env); every other path -# keeps BASE_INSTRUCTIONS and its unchanged hash. -TB2_BASE_INSTRUCTIONS = """You are an autonomous terminal agent solving a -Terminal-Bench task. You have a shell inside the task container: the container -is the environment that will be graded, and the shell starts in the task -workdir. The task statement is the user message. - -Reply format (strict, every reply): think briefly, then make exactly ONE tool -call and write no other visible text. A reply with prose, no tool call or -several tool calls ends the episode with no credit. -- `terminal.exec`: run one shell command (inspect, edit, install, build, test). -- `submit`: finish the episode; arguments {"evidence": ""}. - -Budget: about ten model calls and a small context window; long thinking or -long outputs end the episode early. Act from the first reply, prefer compact -commands (write files with a heredoc, limit output with head/tail), verify the -result, then call `submit`. The grader runs the task's hidden tests against the -container state when the episode ends.""" -def _tb2_surface_sha256() -> str: - """Signed TB2 surface: system prompt + model/Codex-facing tool definitions.""" - return hashlib.sha256( - TB2_BASE_INSTRUCTIONS.encode() + b"\0" + _canonical_json(TB2_DYNAMIC_TOOLS) - ).hexdigest() - - -TB2_BASE_INSTRUCTIONS_SHA256 = _tb2_surface_sha256() -INSTRUCTIONS_FAMILY_ENV = "YETO_CODEX_INSTRUCTIONS_FAMILY" -TB2_INSTRUCTIONS_SHA_ENV = "YETO_CODEX_TB2_BASE_INSTRUCTIONS_SHA256" -TB2_INSTRUCTIONS_FAMILY = "tb2" - - -def base_instructions() -> str: - """System prompt for this process: legacy unless the TB2 family is signed in.""" - family = os.getenv(INSTRUCTIONS_FAMILY_ENV, "").strip() - if not family: - return BASE_INSTRUCTIONS - if family != TB2_INSTRUCTIONS_FAMILY: - raise CodexHarnessError(f"{INSTRUCTIONS_FAMILY_ENV}={family!r} is not a signed instructions family") - if _tb2_surface_sha256() != TB2_BASE_INSTRUCTIONS_SHA256 or not hmac.compare_digest( - os.getenv(TB2_INSTRUCTIONS_SHA_ENV, ""), TB2_BASE_INSTRUCTIONS_SHA256 - ): - raise CodexHarnessError(f"{TB2_INSTRUCTIONS_SHA_ENV} does not match the signed TB2 instructions") - return TB2_BASE_INSTRUCTIONS - - -def _tb2_selected() -> bool: - return base_instructions() is TB2_BASE_INSTRUCTIONS - - -def dynamic_tools() -> list[dict[str, Any]]: - """Codex dynamic tools for this process (TB2 wording when TB2 is signed in).""" - return TB2_DYNAMIC_TOOLS if _tb2_selected() else DYNAMIC_TOOLS - - -def _miles_tools_for(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: - return [ - { - "type": "function", - "function": { - "name": "terminal.exec" if tool["name"] == "terminal_exec" else tool["name"], - "description": tool["description"], - "parameters": tool["inputSchema"], - }, - } - for tool in tools - ] - - -def miles_tools() -> list[dict[str, Any]]: - """Model-facing tools sent to Miles (``terminal.exec`` / ``submit``).""" - return _MILES_TOOLS if not _tb2_selected() else _miles_tools_for(TB2_DYNAMIC_TOOLS) -TERMINAL_EXEC_TOOL_SCHEMA_SHA256 = _sha256_json(TERMINAL_EXEC_TOOL) -SUBMIT_TOOL_SCHEMA_SHA256 = _sha256_json(SUBMIT_TOOL) -DYNAMIC_TOOLS_SCHEMA_SHA256 = _sha256_json(DYNAMIC_TOOLS) - -_MILES_TOOLS = [ - { - "type": "function", - "function": { - "name": "terminal.exec", - "description": TERMINAL_EXEC_TOOL["description"], - "parameters": TERMINAL_EXEC_TOOL["inputSchema"], - }, - }, - { - "type": "function", - "function": { - "name": "submit", - "description": SUBMIT_TOOL["description"], - "parameters": SUBMIT_TOOL["inputSchema"], - }, - }, -] - -_IDENTITY_ENV = { - "YETO_CODEX_APP_SERVER_PROTOCOL_REVISION": CODEX_APP_SERVER_PROTOCOL_REVISION, - "YETO_CODEX_APP_SERVER_SCHEMA_SHA256": CODEX_APP_SERVER_SCHEMA_SHA256, - "YETO_CODEX_BASE_INSTRUCTIONS_SHA256": BASE_INSTRUCTIONS_SHA256, - "YETO_CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256": TERMINAL_EXEC_TOOL_SCHEMA_SHA256, - "YETO_CODEX_SUBMIT_TOOL_SCHEMA_SHA256": SUBMIT_TOOL_SCHEMA_SHA256, - "YETO_CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256": DYNAMIC_TOOLS_SCHEMA_SHA256, - "YETO_CODEX_REASONING_EFFORT": "xhigh", - "YETO_CODEX_BACKEND_REASONING_EFFORT": BACKEND_REASONING_EFFORT, - "YETO_CODEX_BACKEND_THINKING": "enabled", - "YETO_CODEX_CHAT_TEMPLATE": _BACKEND_PROFILE["chat_template"], - "YETO_CODEX_CHAT_TEMPLATE_KWARGS": json.dumps( - BACKEND_CHAT_TEMPLATE_KWARGS, - sort_keys=True, - separators=(",", ":"), - ), - "YETO_CODEX_TITO_ALLOWED_APPEND_ROLES": "tool,user", -} -_HEX64 = re.compile(r"[0-9a-f]{64}") -_RUNTIME_ATTESTATION_LOCK = Lock() -_RUNTIME_ATTESTATION_CACHE: dict[tuple[Any, ...], Path] = {} - - -class CodexHarnessError(RuntimeError): - """A fail-closed identity, transport, or protocol violation.""" - - -class CodexSequenceLimit(CodexHarnessError): - """The next sample would exceed the signed Miles sequence limit.""" - - -class _CompactionContextDoesNotFit(CodexHarnessError): - """Miles proved that a reconstructed context exceeds its declared budget.""" - - -class CodexTurnLimit(CodexHarnessError): - """Codex reached the signed model-call budget after a valid trajectory.""" - - -class CodexSuspendExpired(CodexHarnessError): - """agentic-rollout-utilization 5.3: a trajectory suspended between model turns - was not continued within the survival limit (discarded; environment released).""" - - -# agentic-rollout-utilization 5.1/5.3: model-turn gate of THIS trajectory (a file -# path; the file exists = gate closed). The trusted side (``suspend``/``resume`` -# hooks of the subprocess agent function) creates and removes it; the bridge -# checks it only before a model request, so a running tool call is never -# interrupted and its result is already in the history when the request waits. -SUSPEND_GATE_ENV = "YETO_CODEX_SUSPEND_GATE" -SUSPEND_MAX_SECONDS_ENV = "YETO_CODEX_SUSPEND_MAX_SECONDS" -SUSPEND_MAX_SECONDS_DEFAULT = 600.0 -SUSPEND_POLL_SECONDS = 0.25 -SUSPEND_METRIC_FIELDS = ("suspensions", "suspended_seconds", "suspend_retried_turns", "suspend_expired") - - -def suspend_max_seconds() -> float: - raw = os.environ.get(SUSPEND_MAX_SECONDS_ENV) - if raw is None or raw == "": - return SUSPEND_MAX_SECONDS_DEFAULT - value = float(raw) - if not math.isfinite(value) or value <= 0: - raise ValueError(f"{SUSPEND_MAX_SECONDS_ENV} must be a positive number of seconds") - return value - - -# The worker process installs a sink for gate edges ("enter"/"exit") so the -# trusted side's PhaseClock keeps parked time out of the generation phase. -_suspend_event_sink: Any = None - - -def set_suspend_event_sink(sink: Any) -> None: - global _suspend_event_sink - _suspend_event_sink = sink - - -def _suspend_edge(phase: str) -> None: - if _suspend_event_sink is not None: - _suspend_event_sink(phase) - - -class SuspendGate: - """The bridge side of one trajectory's model-turn gate.""" - - def __init__(self, path: str | None, metrics: Any, *, max_seconds: float | None = None, - poll_seconds: float = SUSPEND_POLL_SECONDS, clock: Any = time.monotonic) -> None: - self.path = path or None - self.metrics = metrics - self.max_seconds = suspend_max_seconds() if max_seconds is None and self.path else (max_seconds or 0.0) - self.poll_seconds = poll_seconds - self.clock = clock - for name in SUSPEND_METRIC_FIELDS: - if not hasattr(metrics, name): - setattr(metrics, name, 0.0 if name == "suspended_seconds" else 0) - - @property - def enabled(self) -> bool: - return self.path is not None - - def closed(self) -> bool: - return self.path is not None and os.path.exists(self.path) - - async def wait_open(self) -> None: - """Return at once when the gate is open; else wait (counted as one - suspension) until it opens, or raise :class:`CodexSuspendExpired` after - the survival limit.""" - if not self.closed(): - return - started = self.clock() - self.metrics.suspensions += 1 - _suspend_edge("enter") - try: - while self.closed(): - if self.clock() - started >= self.max_seconds: - self.metrics.suspend_expired = 1 - raise CodexSuspendExpired( - f"suspended trajectory not continued within {self.max_seconds:g} s") - await asyncio.sleep(self.poll_seconds) - finally: - self.metrics.suspended_seconds += self.clock() - started - _suspend_edge("exit") - - -class CodexModelFailure(CodexHarnessError): - """The sampled model turn could not produce one executable tool call.""" - - -def codex_harness_identity() -> dict[str, str]: - """Return live, recomputed identities consumed by Yeto plan attestation.""" - - live = { - "base_instructions_sha256": hashlib.sha256( - BASE_INSTRUCTIONS.encode() - ).hexdigest(), - "terminal_exec_tool_schema_sha256": _sha256_json(TERMINAL_EXEC_TOOL), - "submit_tool_schema_sha256": _sha256_json(SUBMIT_TOOL), - "dynamic_tools_schema_sha256": _sha256_json(DYNAMIC_TOOLS), - } - declared = { - "base_instructions_sha256": BASE_INSTRUCTIONS_SHA256, - "terminal_exec_tool_schema_sha256": TERMINAL_EXEC_TOOL_SCHEMA_SHA256, - "submit_tool_schema_sha256": SUBMIT_TOOL_SCHEMA_SHA256, - "dynamic_tools_schema_sha256": DYNAMIC_TOOLS_SCHEMA_SHA256, - } - if ( - live != declared - or _MILES_TOOLS != legacy.TOOLS - or any(not _HEX64.fullmatch(value) for value in live.values()) - ): - raise CodexHarnessError("Codex harness identity constants drifted") - return live - - -def _positive_int(name: str) -> int: - raw = os.getenv(name, "").strip() - try: - value = int(raw) - except ValueError as exc: - raise CodexHarnessError(f"{name} must be an integer") from exc - if value <= 0: - raise CodexHarnessError(f"{name} must be positive") - return value - - -def _strict_env_flag(name: str, default: bool = False) -> bool: - raw = os.getenv(name, "1" if default else "0").strip() - if raw not in {"0", "1"}: - raise CodexHarnessError(f"{name} must be exactly 0 or 1") - return raw == "1" - - -def _attest_runtime() -> Path: - codex_harness_identity() - for name, expected in _IDENTITY_ENV.items(): - if not hmac.compare_digest(os.getenv(name, ""), expected): - raise CodexHarnessError(f"{name} does not match the signed harness") - binary = Path(os.getenv("YETO_CODEX_BINARY_PATH", "")) - expected_size = _positive_int("YETO_CODEX_BINARY_SIZE_BYTES") - expected_sha = os.getenv("YETO_CODEX_BINARY_SHA256", "") - if not _HEX64.fullmatch(expected_sha): - raise CodexHarnessError("pinned Codex binary digest is invalid") - if not hmac.compare_digest(os.getenv("YETO_CODEX_VERSION", ""), CODEX_CLI_VERSION): - raise CodexHarnessError("pinned Codex version contract drifted") - if _positive_int("YETO_CODEX_BACKEND_MAX_TOKENS") <= 0: - raise CodexHarnessError("invalid Codex backend token budget") - with _RUNTIME_ATTESTATION_LOCK: - try: - info = binary.stat() - except OSError as exc: - raise CodexHarnessError("pinned Codex binary is unavailable") from exc - if ( - binary.is_symlink() - or not stat.S_ISREG(info.st_mode) - or not info.st_mode & 0o111 - ): - raise CodexHarnessError("pinned Codex binary is not a regular executable") - if info.st_size != expected_size: - raise CodexHarnessError("pinned Codex binary size drifted") - identity = ( - str(binary.resolve()), - info.st_dev, - info.st_ino, - info.st_size, - info.st_mtime_ns, - info.st_ctime_ns, - info.st_mode, - expected_sha, - CODEX_CLI_VERSION, - ) - if identity in _RUNTIME_ATTESTATION_CACHE: - return _RUNTIME_ATTESTATION_CACHE[identity] - digest = hashlib.sha256() - try: - with binary.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - digest.update(chunk) - except OSError as exc: - raise CodexHarnessError("pinned Codex binary could not be read") from exc - if not hmac.compare_digest(digest.hexdigest(), expected_sha): - raise CodexHarnessError("pinned Codex binary digest drifted") - try: - version = subprocess.run( - [str(binary), "--version"], - check=True, - capture_output=True, - text=True, - timeout=15.0, - env={"LANG": "C.UTF-8", "PATH": "/usr/bin:/bin", "TZ": "UTC"}, - ).stdout.strip() - except (OSError, subprocess.SubprocessError) as exc: - raise CodexHarnessError("pinned Codex executable version check failed") from exc - if not hmac.compare_digest(version, CODEX_CLI_VERSION): - raise CodexHarnessError("pinned Codex executable version drifted") - after = binary.stat() - after_identity = ( - str(binary.resolve()), - after.st_dev, - after.st_ino, - after.st_size, - after.st_mtime_ns, - after.st_ctime_ns, - after.st_mode, - expected_sha, - CODEX_CLI_VERSION, - ) - if after_identity != identity: - raise CodexHarnessError("pinned Codex binary changed during attestation") - _RUNTIME_ATTESTATION_CACHE.clear() - _RUNTIME_ATTESTATION_CACHE[identity] = binary - return binary - - -def _canonical_history_item(item: Any) -> dict[str, Any]: - if not isinstance(item, dict) or not isinstance(item.get("type"), str): - raise CodexHarnessError("Codex supplied an invalid history item") - kind = item["type"] - if kind in {"compaction", "compaction_trigger", "context_compaction"}: - raise CodexHarnessError("Codex history compaction is forbidden") - if kind == "message": - role = item.get("role") - content = item.get("content") - if role not in {"developer", "user", "assistant"} or not isinstance(content, list): - raise CodexHarnessError("Codex supplied an invalid message history item") - normalized_content = [] - for part in content: - if not isinstance(part, dict) or part.get("type") not in { - "input_text", - "output_text", - }: - raise CodexHarnessError("Codex supplied non-text message history") - if not isinstance(part.get("text"), str): - raise CodexHarnessError("Codex supplied invalid message text") - normalized_content.append({"type": part["type"], "text": part["text"]}) - return {"type": "message", "role": role, "content": normalized_content} - if kind == "reasoning": - summary = item.get("summary", []) - content = item.get("content") - encrypted = item.get("encrypted_content") - if not isinstance(summary, list) or content is not None and not isinstance(content, list): - raise CodexHarnessError("Codex supplied invalid reasoning history") - normalized: dict[str, Any] = {"type": "reasoning", "summary": copy.deepcopy(summary)} - if content is not None: - normalized["content"] = copy.deepcopy(content) - if encrypted is not None: - if not isinstance(encrypted, str): - raise CodexHarnessError("Codex supplied invalid encrypted reasoning") - normalized["encrypted_content"] = encrypted - return normalized - if kind == "function_call": - name = item.get("name") - arguments = item.get("arguments") - call_id = item.get("call_id") - if ( - name not in {"terminal_exec", "submit"} - or not isinstance(arguments, str) - or not isinstance(call_id, str) - or not call_id - ): - raise CodexHarnessError("Codex supplied invalid function-call history") - return { - "type": "function_call", - "name": name, - "arguments": arguments, - "call_id": call_id, - } - if kind == "function_call_output": - call_id = item.get("call_id") - output = item.get("output") - if not isinstance(call_id, str) or not call_id or not isinstance(output, str): - raise CodexHarnessError("Codex supplied invalid function output history") - return {"type": kind, "call_id": call_id, "output": output} - raise CodexHarnessError(f"Codex supplied forbidden history type {kind!r}") - - -def _response_tool(tool: dict[str, Any]) -> dict[str, Any]: - parameters = copy.deepcopy(tool["inputSchema"]) - if tool["name"] == "terminal_exec": - # Stock 0.145.0's DynamicToolSpec conversion preserves the numeric type - # but intentionally drops JSON-Schema numeric bounds on the Responses wire. - timeout_schema = parameters["properties"]["timeout_seconds"] - timeout_schema.pop("minimum") - timeout_schema.pop("maximum") - return { - "type": "function", - "name": tool["name"], - "description": tool["description"], - "strict": False, - "parameters": parameters, - } - - -_EXPECTED_RESPONSE_TOOLS = [_response_tool(tool) for tool in DYNAMIC_TOOLS] - - -def _validate_codex_request(body: Any, expected_model: str) -> list[dict[str, Any]]: - if not isinstance(body, dict): - raise CodexHarnessError("Codex Responses request is not an object") - if body.get("model") != expected_model or body.get("instructions") != base_instructions(): - raise CodexHarnessError("Codex model or base instructions drifted") - if body.get("stream") is not True or body.get("store") is not False: - raise CodexHarnessError("Codex must use one non-persistent stream") - if body.get("parallel_tool_calls") is not False or body.get("tool_choice") != "auto": - raise CodexHarnessError("Codex sampling/tool-call policy drifted") - reasoning = body.get("reasoning") - if not isinstance(reasoning, dict) or reasoning.get("effort") != "xhigh": - raise CodexHarnessError("Codex did not request xhigh reasoning") - if reasoning.get("summary") not in {None, "none"}: - raise CodexHarnessError("Codex reasoning summary policy drifted") - if body.get("previous_response_id") is not None: - raise CodexHarnessError("server-side Responses history is forbidden") - if body.get("truncation") not in {None, "disabled"}: - raise CodexHarnessError("Codex Responses truncation is forbidden") - tools = body.get("tools") - if not isinstance(tools, list) or [tool.get("name") for tool in tools if isinstance(tool, dict)] != [ - "update_plan", - "terminal_exec", - "submit", - ]: - raise CodexHarnessError("Codex-side tool surface drifted") - if tools[1:] != [_response_tool(tool) for tool in dynamic_tools()]: - raise CodexHarnessError("Codex dynamic tool schemas drifted") - update_plan = tools[0] - if update_plan.get("type") != "function" or not isinstance( - update_plan.get("parameters"), dict - ): - raise CodexHarnessError("stock update_plan schema is invalid") - values = body.get("input") - if not isinstance(values, list) or not values: - raise CodexHarnessError("Codex supplied no explicit history") - return [_canonical_history_item(item) for item in values] - - -def _message_text(item: dict[str, Any]) -> str: - return "\n".join(part["text"] for part in item["content"]) - - -def _usage(completion: dict[str, Any]) -> tuple[dict[str, Any], int | None]: - usage = completion.get("usage") - if not isinstance(usage, dict): - return { - "input_tokens": 0, - "input_tokens_details": {"cached_tokens": 0}, - "output_tokens": 0, - "output_tokens_details": {"reasoning_tokens": 0}, - "total_tokens": 0, - }, None - prompt = usage.get("prompt_tokens", usage.get("input_tokens", 0)) - output = usage.get("completion_tokens", usage.get("output_tokens", 0)) - total = usage.get("total_tokens") - for value in (prompt, output): - if isinstance(value, bool) or not isinstance(value, int) or value < 0: - raise CodexHarnessError("Miles returned invalid token usage") - if total is not None and ( - isinstance(total, bool) or not isinstance(total, int) or total < 0 - ): - raise CodexHarnessError("Miles returned invalid token usage") - return { - "input_tokens": prompt, - "input_tokens_details": {"cached_tokens": 0}, - "output_tokens": output, - "output_tokens_details": {"reasoning_tokens": output}, - "total_tokens": total, - }, total - - -async def _read_bounded_miles_response(response: Any) -> bytes: - """Read one Miles response without ever buffering past the signed cap.""" - - content_length = response.content_length - if content_length is not None and content_length > MAX_MILES_RESPONSE_BYTES: - raise CodexHarnessError("Miles session response is oversized") - chunks: list[bytes] = [] - size = 0 - async for chunk in response.content.iter_chunked(_MILES_RESPONSE_READ_CHUNK_BYTES): - if len(chunk) > MAX_MILES_RESPONSE_BYTES - size: - raise CodexHarnessError("Miles session response is oversized") - chunks.append(chunk) - size += len(chunk) - return b"".join(chunks) - - -def _normalize_miles_stream_chunk(value: Any) -> dict[str, Any]: - """Normalize pinned Miles' one fake-stream chunk to a completion.""" - - if not isinstance(value, dict) or set(value) != { - "id", - "object", - "created", - "model", - "choices", - "usage", - }: - raise CodexHarnessError("Miles session returned an invalid SSE chunk") - identifier = value.get("id") - created = value.get("created") - if ( - not isinstance(identifier, str) - or not identifier - or value.get("object") != "chat.completion.chunk" - or value.get("model") != BACKEND_MODEL - or isinstance(created, bool) - or not isinstance(created, int) - or created < 0 - or not isinstance(value.get("usage"), dict) - ): - raise CodexHarnessError("Miles session returned an invalid SSE chunk") - choices = value.get("choices") - if not isinstance(choices, list) or len(choices) != 1: - raise CodexHarnessError("Miles session returned an invalid SSE chunk") - choice = choices[0] - if not isinstance(choice, dict) or set(choice) != { - "index", - "delta", - "finish_reason", - }: - raise CodexHarnessError("Miles session returned an invalid SSE chunk") - choice_index = choice.get("index") - finish_reason = choice.get("finish_reason") - if ( - isinstance(choice_index, bool) - or not isinstance(choice_index, int) - or choice_index != 0 - or finish_reason - not in { - "stop", - "length", - "tool_calls", - "content_filter", - "function_call", - "abort", - } - ): - raise CodexHarnessError("Miles session returned an invalid SSE chunk") - delta = choice.get("delta") - if ( - not isinstance(delta, dict) - or not {"role", "content"}.issubset(delta) - or not set(delta).issubset( - {"role", "content", "reasoning_content", "tool_calls"} - ) - or delta.get("role") != "assistant" - or not isinstance(delta.get("content"), str) - ): - raise CodexHarnessError("Miles session returned an invalid SSE delta") - reasoning = delta.get("reasoning_content") - if reasoning is not None and not isinstance(reasoning, str): - raise CodexHarnessError("Miles session returned an invalid SSE delta") - raw_calls = delta.get("tool_calls") - calls = None - if raw_calls is not None: - if not isinstance(raw_calls, list): - raise CodexHarnessError("Miles session returned an invalid SSE delta") - calls = [] - for index, raw_call in enumerate(raw_calls): - call_index = raw_call.get("index") if isinstance(raw_call, dict) else None - if ( - not isinstance(raw_call, dict) - or set(raw_call) != {"id", "index", "type", "function"} - or isinstance(call_index, bool) - or not isinstance(call_index, int) - or call_index != index - or raw_call.get("type") != "function" - or not isinstance(raw_call.get("id"), str) - or not raw_call["id"] - ): - raise CodexHarnessError("Miles session returned an invalid SSE delta") - function = raw_call.get("function") - if ( - not isinstance(function, dict) - or set(function) != {"name", "arguments"} - or not isinstance(function.get("name"), str) - or not isinstance(function.get("arguments"), str) - ): - raise CodexHarnessError("Miles session returned an invalid SSE delta") - calls.append( - { - "id": raw_call["id"], - "type": "function", - "function": copy.deepcopy(function), - } - ) - message: dict[str, Any] = { - "role": "assistant", - "content": delta["content"], - } - if reasoning is not None: - message["reasoning_content"] = reasoning - if calls is not None: - message["tool_calls"] = calls - return { - "id": identifier, - "object": "chat.completion", - "created": created, - "model": BACKEND_MODEL, - "choices": [ - { - "index": 0, - "message": message, - "finish_reason": finish_reason, - } - ], - "usage": copy.deepcopy(value["usage"]), - } - - -def _parse_miles_sse(raw: bytes) -> dict[str, Any]: - """Parse exactly one pinned Miles data event followed by ``[DONE]``.""" - - def object_without_duplicates(pairs: list[tuple[str, Any]]) -> dict[str, Any]: - value = {} - for key, item in pairs: - if key in value: - raise ValueError("duplicate JSON object key") - value[key] = item - return value - - def reject_nonfinite_constant(_value: str) -> None: - raise ValueError("non-finite JSON number") - - try: - text = raw.decode("utf-8") - except UnicodeDecodeError as exc: - raise CodexHarnessError("Miles session returned invalid UTF-8") from exc - blocks = text.split("\n\n") - if len(blocks) != 3 or blocks[-1] or any("\n" in block for block in blocks[:-1]): - raise CodexHarnessError("Miles session returned malformed SSE") - if not all(block.startswith("data: ") for block in blocks[:-1]): - raise CodexHarnessError("Miles session returned malformed SSE") - payload, done = (block[len("data: ") :] for block in blocks[:-1]) - if done != "[DONE]": - raise CodexHarnessError("Miles session returned malformed SSE") - try: - value = json.loads( - payload, - object_pairs_hook=object_without_duplicates, - parse_constant=reject_nonfinite_constant, - ) - except (json.JSONDecodeError, ValueError) as exc: - raise CodexHarnessError("Miles session returned invalid SSE JSON") from exc - return _normalize_miles_stream_chunk(value) - - -def _sse_response( - *, - request: dict[str, Any], - output: list[dict[str, Any]], - usage: dict[str, Any], - response_number: int, -) -> bytes: - response_id = f"resp_yeto_{response_number}" - completed = { - "id": response_id, - "object": "response", - "created_at": int(time.time()), - "status": "completed", - "error": None, - "incomplete_details": None, - "instructions": base_instructions(), - "max_output_tokens": None, - "max_tool_calls": None, - "model": request["model"], - "output": output, - "parallel_tool_calls": False, - "previous_response_id": None, - "prompt_cache_key": request.get("prompt_cache_key"), - "reasoning": {"effort": "xhigh", "summary": "none"}, - "safety_identifier": None, - "service_tier": "default", - "store": False, - "temperature": None, - "text": {"format": {"type": "text"}}, - "tool_choice": "auto", - "tools": request["tools"], - "top_logprobs": 0, - "top_p": None, - "truncation": "disabled", - "usage": usage, - "user": None, - "metadata": {}, - } - events: list[dict[str, Any]] = [ - { - "type": "response.created", - "response": {**completed, "status": "in_progress", "output": []}, - } - ] - for index, item in enumerate(output): - added = copy.deepcopy(item) - added["status"] = "in_progress" - if item["type"] == "function_call": - added["arguments"] = "" - events.append( - {"type": "response.output_item.added", "output_index": index, "item": added} - ) - if item["type"] == "function_call": - events.extend( - [ - { - "type": "response.function_call_arguments.delta", - "item_id": item["id"], - "output_index": index, - "delta": item["arguments"], - }, - { - "type": "response.function_call_arguments.done", - "item_id": item["id"], - "output_index": index, - "arguments": item["arguments"], - }, - ] - ) - events.append( - {"type": "response.output_item.done", "output_index": index, "item": item} - ) - events.append({"type": "response.completed", "response": completed}) - chunks = [] - for event in events: - chunks.append(f"event: {event['type']}\n".encode()) - chunks.append(b"data: " + _canonical_json(event) + b"\n\n") - chunks.append(b"data: [DONE]\n\n") - return b"".join(chunks) - - -@dataclass -class _PendingTool: - call_id: str - codex_name: str - miles_name: str - raw_arguments: str - - -def _tool_error(message: str) -> dict[str, str]: - compact = message if len(message) <= 1024 else message[:1000] + "...[truncated]" - return {"error": compact} - - -# G6(a): bridge rejections that are, in gateway terms, a session mismatch -# (history is not the expected prefix) or a chain break (reason names from -# ``yeto.rl.engine.timeline.TITO_CHAIN_BREAK_REASONS``). -TITO_SESSION_MISMATCH_MESSAGES = frozenset( - { - "Codex mutated the initial task history", - "Codex truncated or mutated episode history", - "Codex supplied no explicit history", - } -) -TITO_SESSION_MISMATCH_KEY = "tito_session_mismatch" -TITO_CHAIN_BREAKS_KEY = "tito_chain_breaks" -TITO_CHAIN_BREAK_MESSAGES = { - "Codex retried a Responses sample": "retry_fork", - "Codex history compaction is forbidden": "compaction_window", -} - - -def tito_counters(metrics: Any) -> dict[str, Any]: - """The G6(a) counters recorded on a metrics object (zero when none).""" - return { - TITO_SESSION_MISMATCH_KEY: int(getattr(metrics, TITO_SESSION_MISMATCH_KEY, 0) or 0), - TITO_CHAIN_BREAKS_KEY: dict(getattr(metrics, TITO_CHAIN_BREAKS_KEY, {}) or {}), - } - - -class _ResponsesBridge: - def __init__( - self, - miles_base_url: str, - prompt: str, - request_kwargs: dict[str, Any], - metrics: legacy.AgentMetrics, - *, - max_seq_len: int | None, - ) -> None: - self._miles_url = ( - f"{legacy._session_url(miles_base_url).rstrip('/')}/chat/completions" - ) - self._prompt = prompt - self._request_kwargs = dict(request_kwargs) - self._metrics = metrics - self._max_seq_len = max_seq_len - self._model = BACKEND_MODEL - self._max_turns = legacy._positive_int_env("SECRLENV_MAX_TURNS", 40) - self._backend_max_tokens = _positive_int("YETO_CODEX_BACKEND_MAX_TOKENS") - unexpected_sampling = set(self._request_kwargs) - ( - _SAMPLING_FIELDS | _STRUCTURAL_REQUEST_FIELDS - ) - if unexpected_sampling: - raise CodexHarnessError("Miles supplied an unknown sampling field") - forbidden_structural = set(self._request_kwargs) & ( - _STRUCTURAL_REQUEST_FIELDS - - {"best_of", "max_tokens", "n", "stream", "stream_options"} - ) - if forbidden_structural: - raise CodexHarnessError("Miles attempted to override a signed request field") - requested_max_tokens = self._request_kwargs.get( - "max_tokens", self._backend_max_tokens - ) - if ( - isinstance(requested_max_tokens, bool) - or not isinstance(requested_max_tokens, int) - or not 0 < requested_max_tokens <= self._backend_max_tokens - ): - raise CodexHarnessError( - "Miles sampling max_tokens exceeds the signed backend budget" - ) - self._requested_max_tokens = requested_max_tokens - for key in ("n", "best_of"): - if self._request_kwargs.get(key) not in {None, 1}: - raise CodexHarnessError("Miles requested more than one sample") - self._compaction_enabled = _strict_env_flag("YETO_CODEX_COMPACTION_ENABLED") - self._compaction_trigger_tokens = legacy._positive_int_env( - "YETO_CODEX_COMPACTION_TRIGGER_TOKENS", 6_144 - ) - self._compaction_summary_max_tokens = legacy._positive_int_env( - "YETO_CODEX_COMPACTION_SUMMARY_MAX_TOKENS", 1_024 - ) - self._max_compactions = legacy._positive_int_env("YETO_CODEX_MAX_COMPACTIONS", 3) - if self._compaction_enabled: - if self._max_seq_len is None: - raise CodexHarnessError("Codex compaction requires max_seq_len") - if self._compaction_trigger_tokens >= self._max_seq_len: - raise CodexHarnessError( - "Codex compaction trigger must be smaller than max_seq_len" - ) - summary_generation_tokens = min( - self._requested_max_tokens, - self._compaction_summary_max_tokens, - ) - if ( - self._compaction_trigger_tokens - + COMPACTION_SUMMARY_INSTRUCTION_TOKEN_RESERVE - + summary_generation_tokens - > self._max_seq_len - ): - raise CodexHarnessError( - "Codex compaction trigger leaves insufficient room for the " - "summary instruction and generation" - ) - self._token = secrets.token_urlsafe(32) - self._seen_bodies: set[str] = set() - self._expected_input: list[dict[str, Any]] | None = None - self._history_after_model: list[dict[str, Any]] | None = None - self._pending: _PendingTool | None = None - self._messages: list[dict[str, Any]] = [] - self._atomic_steps: list[tuple[dict[str, Any], dict[str, Any]]] = [] - self._last_assistant_message: dict[str, Any] | None = None - self._context_window = 0 - self._segment_index = 0 - self._segment_type = "execution" - self._compaction_count = 0 - self._current_context_tokens = 0 - self._resume_summary: str | None = None - self._retained_step_count = COMPACTION_KEEP_ATOMIC_STEPS - self._unobserved_tool_token_upper_bound = 0 - self._request_count = 0 - self._terminal = False - self._runner: web.AppRunner | None = None - self._session: aiohttp.ClientSession | None = None - self._lock = asyncio.Lock() - self._fatal: asyncio.Future[BaseException] | None = None - self.url: str | None = None - self._gate = SuspendGate(os.environ.get(SUSPEND_GATE_ENV), metrics) - - @property - def token(self) -> str: - return self._token - - @property - def fatal(self) -> asyncio.Future[BaseException]: - if self._fatal is None: - raise RuntimeError("bridge is not started") - return self._fatal - - async def __aenter__(self) -> Self: - loop = asyncio.get_running_loop() - self._fatal = loop.create_future() - self._session = aiohttp.ClientSession( - timeout=aiohttp.ClientTimeout(total=legacy._positive_env( - "SECRLENV_MAX_ROLLOUT_TIME_SECONDS", 3600.0 - )), - connector=aiohttp.TCPConnector(limit=1), - ) - app = web.Application(client_max_size=2 * 1024 * 1024) - app.router.add_post("/v1/responses", self._handle_responses) - app.router.add_route("*", "/{tail:.*}", self._reject_route) - self._runner = web.AppRunner(app, access_log=None) - await self._runner.setup() - site = web.TCPSite(self._runner, "127.0.0.1", 0) - await site.start() - sockets = site._server.sockets if site._server is not None else [] - if len(sockets) != 1: - raise CodexHarnessError("Responses bridge did not bind one loopback socket") - host, port = sockets[0].getsockname()[:2] - if host != "127.0.0.1": - raise CodexHarnessError("Responses bridge escaped loopback") - self.url = f"http://127.0.0.1:{port}" - return self - - async def __aexit__(self, *_exc: object) -> None: - if self._session is not None: - await self._session.close() - if self._runner is not None: - await self._runner.cleanup() - - def _fail(self, exc: BaseException) -> None: - if self._fatal is not None and not self._fatal.done(): - self._fatal.set_result(exc) - - async def _reject_route(self, _request: web.Request) -> web.Response: - self._fail(CodexHarnessError("Codex called an unexpected bridge route")) - return web.json_response({"error": "unsupported route"}, status=404) - - def expect_tool_output(self, call_id: str, output: str) -> None: - if self._terminal or self._pending is None or self._history_after_model is None: - raise CodexHarnessError("Codex tool output arrived outside a model call") - if not hmac.compare_digest(self._pending.call_id, call_id): - raise CodexHarnessError("Codex tool output call ID drifted") - self._expected_input = self._history_after_model + [ - {"type": "function_call_output", "call_id": call_id, "output": output} - ] - self._messages.append( - {"role": "tool", "tool_call_id": call_id, "content": output} - ) - if self._last_assistant_message is None: - raise CodexHarnessError("Codex lost the assistant half of an atomic tool step") - self._atomic_steps.append( - ( - copy.deepcopy(self._last_assistant_message), - copy.deepcopy(self._messages[-1]), - ) - ) - self._last_assistant_message = None - self._unobserved_tool_token_upper_bound += len(output.encode("utf-8")) + 256 - _note_turn_tool_output(self._metrics, output) - self._pending = None - - def mark_terminal(self) -> None: - self._terminal = True - - def take_pending(self) -> _PendingTool: - if self._pending is None: - raise CodexHarnessError("Codex requested a tool without a model call") - return self._pending - - async def _handle_responses(self, request: web.Request) -> web.Response: - authorization = request.headers.get("Authorization", "") - if not hmac.compare_digest(authorization, f"Bearer {self._token}"): - self._fail(CodexHarnessError("Codex bridge authentication failed")) - return web.json_response({"error": "unauthorized"}, status=401) - async with self._lock: - if self._terminal: - self._fail(CodexHarnessError("Codex sampled after terminal submit")) - return web.json_response({"error": "episode is terminal"}, status=409) - try: - body = await request.json() - fingerprint = hashlib.sha256(_canonical_json(body)).hexdigest() - if fingerprint in self._seen_bodies: - raise CodexHarnessError("Codex retried a Responses sample") - self._seen_bodies.add(fingerprint) - history = _validate_codex_request(body, self._model) - if self._request_count >= self._max_turns: - raise CodexTurnLimit("Codex reached the signed turn budget") - if self._expected_input is None: - expected = [ - { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": self._prompt}], - } - ] - if history != expected: - raise CodexHarnessError("Codex mutated the initial task history") - self._messages = [ - {"role": "system", "content": base_instructions()}, - {"role": "user", "content": self._prompt}, - ] - elif history != self._expected_input: - raise CodexHarnessError("Codex truncated or mutated episode history") - estimated_context_tokens = ( - self._current_context_tokens - + self._unobserved_tool_token_upper_bound - ) - if self._should_compact(estimated_context_tokens): - await self._compact_context() - estimated_context_tokens = 0 - if self._max_seq_len is not None and estimated_context_tokens >= self._max_seq_len: - self._metrics.max_seq_len_hit = 1 - raise CodexSequenceLimit("Miles sequence limit reached") - completion = await self._sample_execution_with_tail_fallback() - self._request_count += 1 - self._metrics.turns = self._request_count - usage, total = _usage(completion) - if total is None: - if self._max_seq_len is not None: - self._metrics.usage_missing += 1 - raise CodexHarnessError("Miles omitted required token usage") - else: - self._metrics.max_model_total_tokens = max( - self._metrics.max_model_total_tokens, total - ) - self._current_context_tokens = total - self._unobserved_tool_token_upper_bound = 0 - output = self._translate_completion(completion) - self._history_after_model = history + [ - _canonical_history_item(item) for item in output - ] - return web.Response( - body=_sse_response( - request=body, - output=output, - usage=usage, - response_number=self._request_count, - ), - headers={ - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - "X-Accel-Buffering": "no", - }, - ) - except (CodexSequenceLimit, CodexTurnLimit, CodexModelFailure) as exc: - if isinstance(exc, CodexSequenceLimit): - self._metrics.max_seq_len_hit = 1 - elif isinstance(exc, CodexModelFailure): - self._metrics.parse_failures += 1 - self._fail(exc) - return web.json_response({"error": "episode boundary"}, status=400) - except Exception as exc: # noqa: BLE001 - sanitize all parser failures - failure = exc if isinstance(exc, CodexHarnessError) else CodexHarnessError( - "Codex Responses bridge rejected the request" - ) - self._note_tito_rejection(failure) - self._fail(failure) - return web.json_response({"error": "protocol violation"}, status=400) - - def _note_tito_rejection(self, failure: CodexHarnessError) -> None: - """G6(a): count a rejection the way the gateway would (9.2 judgement). - - The bridge never forks a chain: a request whose history is not the - exact expected prefix is a session mismatch (``tito_session_mismatch``); - a byte-identical retry is a ``retry_fork`` chain break and a compaction - attempt a ``compaction_window`` one. The trusted layer mirrors these - onto the HarnessBoard (``record_session_mismatch`` / ``record_chain_break``). - """ - # ``AgentMetrics`` is sha-pinned legacy source, so the counters ride on - # the instance as extra attributes (read back by ``tito_counters``). - message = str(failure) - if message in TITO_SESSION_MISMATCH_MESSAGES: - setattr(self._metrics, TITO_SESSION_MISMATCH_KEY, getattr(self._metrics, TITO_SESSION_MISMATCH_KEY, 0) + 1) - reason = TITO_CHAIN_BREAK_MESSAGES.get(message) - if reason is not None: - breaks = dict(getattr(self._metrics, TITO_CHAIN_BREAKS_KEY, {}) or {}) - breaks[reason] = breaks.get(reason, 0) + 1 - setattr(self._metrics, TITO_CHAIN_BREAKS_KEY, breaks) - - def _compaction_headers(self) -> dict[str, str]: - if not self._compaction_enabled or self._max_seq_len is None: - return {} - return { - "X-Miles-Compaction-Schema-Version": str(COMPACTION_SCHEMA_VERSION), - "X-Miles-Compaction-Context-Window": str(self._context_window), - "X-Miles-Compaction-Segment-Index": str(self._segment_index), - "X-Miles-Compaction-Segment-Type": self._segment_type, - "X-Miles-Compaction-Context-Budget": str(self._max_seq_len), - } - - def _should_compact(self, estimated_context_tokens: int) -> bool: - """Return whether consumed context reached the configured boundary.""" - return ( - self._compaction_enabled - and self._compaction_count < self._max_compactions - and self._max_seq_len is not None - and estimated_context_tokens >= self._compaction_trigger_tokens - ) - - def _rebuild_compacted_messages(self, retained_step_count: int) -> None: - if self._resume_summary is None: - raise CodexHarnessError("compaction resume summary is missing") - retained = self._atomic_steps[-retained_step_count:] if retained_step_count else [] - messages: list[dict[str, Any]] = [ - {"role": "system", "content": base_instructions()}, - { - "role": "user", - "content": COMPACTION_RESUME_TEMPLATE.format(summary=self._resume_summary), - }, - ] - for assistant, tool in retained: - messages.extend((copy.deepcopy(assistant), copy.deepcopy(tool))) - self._messages = messages - self._atomic_steps = copy.deepcopy(retained) - self._retained_step_count = retained_step_count - - async def _compact_context(self) -> None: - """Sample a trainable summary, then start one reconstructed window.""" - if self._pending is not None or self._last_assistant_message is not None: - raise CodexHarnessError("compaction attempted inside an atomic tool step") - if not self._atomic_steps: - raise CodexHarnessError("compaction attempted without a complete tool step") - if self._max_seq_len is None: - raise CodexHarnessError("compaction requires a context budget") - summary_generation_tokens = min( - self._requested_max_tokens, - self._compaction_summary_max_tokens, - ) - if ( - self._current_context_tokens - + self._unobserved_tool_token_upper_bound - + COMPACTION_SUMMARY_INSTRUCTION_TOKEN_RESERVE - + summary_generation_tokens - > self._max_seq_len - ): - raise CodexSequenceLimit( - "compaction was triggered too late to reserve a complete summary" - ) - self._segment_index += 1 - self._segment_type = "summary" - summary_messages = copy.deepcopy(self._messages) - summary_messages.append( - {"role": "user", "content": COMPACTION_SUMMARY_INSTRUCTION} - ) - try: - completion = await self._sample_miles( - messages=summary_messages, - summary=True, - ) - except _CompactionContextDoesNotFit as exc: - raise CodexSequenceLimit( - "compaction summary instruction no longer fits the context budget" - ) from exc - usage, total = _usage(completion) - del usage - if total is None: - self._metrics.usage_missing += 1 - raise CodexHarnessError("Miles omitted summary token usage") - self._metrics.max_model_total_tokens = max( - self._metrics.max_model_total_tokens, total - ) - self._current_context_tokens = total - self._unobserved_tool_token_upper_bound = 0 - summary = self._translate_summary(completion) - - self._resume_summary = summary - self._context_window += 1 - self._segment_index += 1 - self._segment_type = "execution" - self._compaction_count += 1 - self._current_context_tokens = 0 - self._rebuild_compacted_messages( - min(COMPACTION_KEEP_ATOMIC_STEPS, len(self._atomic_steps)) - ) - - async def _sample_execution_with_tail_fallback(self) -> dict[str, Any]: - """Reduce the exact recent-step tail only after Miles proves it cannot fit.""" - while True: - try: - return await self._sample_miles() - except _CompactionContextDoesNotFit as exc: - if self._resume_summary is None or self._retained_step_count == 0: - raise CodexSequenceLimit( - "compacted resume context does not fit the context budget" - ) from exc - self._rebuild_compacted_messages(self._retained_step_count - 1) - - async def _sample_miles( - self, - *, - messages: list[dict[str, Any]] | None = None, - summary: bool = False, - ) -> dict[str, Any]: - if self._session is None: - raise RuntimeError("bridge is not started") - sampling = { - key: copy.deepcopy(value) - for key, value in self._request_kwargs.items() - if key in _SAMPLING_FIELDS - } - max_tokens = self._requested_max_tokens - if self._max_seq_len is not None: - remaining = ( - self._max_seq_len - - self._current_context_tokens - - self._unobserved_tool_token_upper_bound - ) - if summary: - remaining -= COMPACTION_SUMMARY_INSTRUCTION_TOKEN_RESERVE - if remaining <= 0: - raise CodexSequenceLimit("Miles sequence limit reached") - max_tokens = min(max_tokens, remaining) - if summary: - max_tokens = min(max_tokens, self._compaction_summary_max_tokens) - payload = { - **sampling, - "model": self._model, - "messages": copy.deepcopy(self._messages if messages is None else messages), - "parallel_tool_calls": False, - # Pinned Miles fake-streams one compact client event while retaining - # the full token/logprob response in its session record for TITO. - "stream": True, - "max_tokens": max_tokens, - "reasoning_effort": BACKEND_REASONING_EFFORT, - "thinking": {"type": "enabled"}, - "chat_template_kwargs": copy.deepcopy(BACKEND_CHAT_TEMPLATE_KWARGS), - } - if summary: - payload["tools"] = [] - payload["tool_choice"] = "none" - else: - payload["tools"] = copy.deepcopy(miles_tools()) - payload["tool_choice"] = "auto" - started = time.monotonic() - suspended_before = self._metrics.suspended_seconds - try: - while True: - # 5.1: a suspended trajectory waits here, between two model turns. - await self._gate.wait_open() - async with self._session.post( - self._miles_url, - json=payload, - headers=self._compaction_headers(), - ) as response: - status = response.status - if status == 503 and self._gate.closed(): - # The cut-off aborted this model turn in the engine (the gate - # closed first); the session server did not record it. Redo - # the same turn once the trajectory is resumed. - await _read_bounded_miles_response(response) - self._metrics.suspend_retried_turns += 1 - continue - if status == 413: - await _read_bounded_miles_response(response) - raise _CompactionContextDoesNotFit( - "Miles proved the compacted context exceeds its budget" - ) - if status != 200: - snippet = "" - if status != 503: - # bounded diagnostic read (503 keeps the retry-path - # no-read contract pinned by the sec harness tests) - detail = await _read_bounded_miles_response(response) - snippet = detail[:500].decode("utf-8", errors="replace") if isinstance(detail, bytes) else str(detail)[:500] - raise CodexHarnessError( - f"Miles session returned HTTP {status}" + (f": {snippet}" if snippet else "")) - if response.content_type != "text/event-stream": - raise CodexHarnessError( - "Miles session returned an unexpected media type" - ) - raw = await _read_bounded_miles_response(response) - break - except (aiohttp.ClientError, TimeoutError) as exc: - raise CodexHarnessError("Miles session transport failed") from exc - finally: - # time parked at the gate is not generation time - self._metrics.total_generation_time += ( - time.monotonic() - started - (self._metrics.suspended_seconds - suspended_before)) - completion = _parse_miles_sse(raw) - if not summary and self._resume_summary is not None: - # The session server tokenized this exact candidate before sampling. - # A 413 above is the only allowed reason to reduce k. - self._resume_summary = None - return completion - - def _translate_summary(self, completion: dict[str, Any]) -> str: - choices = completion.get("choices") - if not isinstance(choices, list) or len(choices) != 1: - raise CodexHarnessError("Miles returned extra or missing summary samples") - choice = choices[0] - if not isinstance(choice, dict) or choice.get("finish_reason") != "stop": - raise CodexModelFailure("compaction summary did not finish cleanly") - message = choice.get("message") - if not isinstance(message, dict) or message.get("role") != "assistant": - raise CodexModelFailure("compaction summary is not an assistant message") - content = message.get("content") - calls = message.get("tool_calls") - if ( - not isinstance(content, str) - or not content.strip() - or (calls is not None and calls != []) - ): - raise CodexModelFailure("compaction summary must be non-empty and tool-free") - return content.strip() - - def _translate_completion(self, completion: dict[str, Any]) -> list[dict[str, Any]]: - choices = completion.get("choices") - if not isinstance(choices, list) or len(choices) != 1: - raise CodexHarnessError("Miles returned extra or missing samples") - choice = choices[0] - if not isinstance(choice, dict): - raise CodexHarnessError("Miles returned an invalid sample") - _note_last_completion(self._metrics, choice, completion) - if choice.get("finish_reason") == "length": - raise CodexSequenceLimit("Miles returned a truncated sample") - message = choice.get("message") - if not isinstance(message, dict) or message.get("role") != "assistant": - raise CodexHarnessError("Miles returned an invalid assistant message") - content = message.get("content", "") - reasoning = message.get("reasoning_content") - calls = message.get("tool_calls") - if not isinstance(content, str) or not isinstance(reasoning, str) or not reasoning: - raise CodexModelFailure("DSV4 did not return thinking content") - if content.strip() and not ALLOW_PROSE_WITH_TOOL_CALL: - raise CodexModelFailure("DSV4 mixed prose with the required tool call") - if not isinstance(calls, list) or len(calls) != 1: - raise CodexModelFailure("DSV4 must return exactly one tool call") - call = calls[0] - function = call.get("function") if isinstance(call, dict) else None - call_id = call.get("id") if isinstance(call, dict) else None - if not isinstance(function, dict) or not isinstance(call_id, str) or not call_id: - raise CodexModelFailure("DSV4 returned an invalid tool call") - name = function.get("name") - raw_arguments = function.get("arguments") - if name not in {"terminal.exec", "submit"} or not isinstance(raw_arguments, str): - raise CodexModelFailure("DSV4 returned a forbidden tool") - try: - arguments = json.loads(raw_arguments) - except json.JSONDecodeError as exc: - raise CodexModelFailure("DSV4 returned malformed raw tool arguments") from exc - if not isinstance(arguments, dict): - raise CodexModelFailure("DSV4 tool arguments are not an object") - codex_name = "terminal_exec" if name == "terminal.exec" else "submit" - if self._pending is not None: - raise CodexHarnessError("DSV4 emitted a tool before prior output") - self._pending = _PendingTool(call_id, codex_name, name, raw_arguments) - self._metrics.tool_calls += 1 - self._messages.append(copy.deepcopy(message)) - self._last_assistant_message = copy.deepcopy(message) - index = self._request_count + 1 - return [ - { - "id": f"rs_yeto_{index}", - "type": "reasoning", - "summary": [], - "content": [{"type": "reasoning_text", "text": reasoning}], - "encrypted_content": None, - "status": "completed", - }, - { - "id": f"fc_yeto_{index}", - "type": "function_call", - "status": "completed", - "name": codex_name, - "call_id": call_id, - "arguments": raw_arguments, - }, - ] - - -LAST_COMPLETION_HEAD_CHARS = 200 - -# S17 C9 (observe only): why a Codex episode ended, finer than the signed -# ``exit_status`` (which folds protocol errors into ``max_turns`` and reply -# truncation into ``max_seq_len``). -END_KIND_SUBMIT = "submit" # episode ended after a submit call -END_KIND_COMPLETED = "completed_no_submit" # Codex stopped on its own, no submit -END_KIND_TURN_LIMIT = "turn_limit" # signed turn budget reached -END_KIND_PROTOCOL_ERROR = "protocol_error" # model reply violated the tool protocol -END_KIND_RESPONSE_TRUNCATED = "response_truncated" # one reply hit the per-reply cap -END_KIND_CONTEXT_LIMIT = "context_limit" # whole trajectory hit the sequence cap -END_KIND_TIMEOUT = "timeout" -END_KINDS = (END_KIND_SUBMIT, END_KIND_COMPLETED, END_KIND_TURN_LIMIT, END_KIND_PROTOCOL_ERROR, - END_KIND_RESPONSE_TRUNCATED, END_KIND_CONTEXT_LIMIT, END_KIND_TIMEOUT) -_TRUNCATED_SAMPLE_MESSAGE = "Miles returned a truncated sample" -# Per-turn length lists are capped so one runaway episode cannot bloat events. -TURN_LENGTHS_MAX = 256 - - -def end_kind_for(exc: BaseException) -> str: - if isinstance(exc, CodexSequenceLimit): - return END_KIND_RESPONSE_TRUNCATED if str(exc) == _TRUNCATED_SAMPLE_MESSAGE else END_KIND_CONTEXT_LIMIT - if isinstance(exc, CodexTurnLimit): - return END_KIND_TURN_LIMIT - if isinstance(exc, CodexModelFailure): - return END_KIND_PROTOCOL_ERROR - raise TypeError(f"not a policy boundary: {type(exc).__name__}") - - -def _append_capped(metrics: Any, name: str, value: Any) -> None: - values = getattr(metrics, name, None) - if not isinstance(values, list): - values = [] - setattr(metrics, name, values) - if len(values) < TURN_LENGTHS_MAX: - values.append(value) - - -def _note_turn_tool_output(metrics: Any, output: str) -> None: - """Observe only (S17 C9): bytes of each tool result fed back to the model.""" - _append_capped(metrics, "turn_tool_output_bytes", len(output.encode("utf-8"))) - - -def _note_last_completion(metrics: Any, choice: dict[str, Any], completion: dict[str, Any]) -> None: - """Observe only (S15 r2 follow-up): shape of the latest model reply, so a - tape can tell thinking exhaustion from prose / missing tool calls.""" - message = choice.get("message") if isinstance(choice.get("message"), dict) else {} - content = message.get("content") if isinstance(message.get("content"), str) else "" - reasoning = message.get("reasoning_content") if isinstance(message.get("reasoning_content"), str) else "" - calls = message.get("tool_calls") if isinstance(message.get("tool_calls"), list) else [] - usage = completion.get("usage") if isinstance(completion.get("usage"), dict) else {} - names = [] - for call in calls: - function = call.get("function") if isinstance(call, dict) else None - names.append(str(function.get("name"))[:32] if isinstance(function, dict) else "?") - # S17 C9: one entry per model reply (None when the backend sent no usage). - tokens = usage.get("completion_tokens") - _append_capped(metrics, "turn_completion_tokens", - tokens if isinstance(tokens, int) and not isinstance(tokens, bool) else None) - total = usage.get("total_tokens") - _append_capped(metrics, "turn_context_tokens", - total if isinstance(total, int) and not isinstance(total, bool) else None) - metrics.last_completion = { # type: ignore[attr-defined] - "finish_reason": str(choice.get("finish_reason"))[:16], - "content_chars": len(content), - "reasoning_chars": len(reasoning), - "tool_calls": len(calls), - "tool_names": ",".join(names)[:96], - "completion_tokens": usage.get("completion_tokens") if isinstance(usage.get("completion_tokens"), int) else None, - "content_head": content[:LAST_COMPLETION_HEAD_CHARS], - "reasoning_tail": reasoning[-LAST_COMPLETION_HEAD_CHARS:], - } - - -def _toml_string(value: str) -> str: - return json.dumps(value, ensure_ascii=True) - - -def _suspend_settings() -> list[str]: - """5.1: a suspended trajectory's model request is held by the bridge for up to - the survival limit (plus a training round); Codex must not give up on it.""" - if not os.environ.get(SUSPEND_GATE_ENV): - return [] # default command line unchanged - idle_ms = int((suspend_max_seconds() + 600.0) * 1000) - return [f"model_providers.miles.stream_idle_timeout_ms={idle_ms}"] - - -def _codex_argv(binary: Path, bridge: _ResponsesBridge) -> list[str]: - if bridge.url is None: - raise RuntimeError("bridge is not started") - model = BACKEND_MODEL - settings = [ - f"model={_toml_string(model)}", - 'model_provider="miles"', - 'model_reasoning_effort="xhigh"', - 'model_reasoning_summary="none"', - 'model_providers.miles.name="Yeto Miles"', - f"model_providers.miles.base_url={_toml_string(bridge.url + '/v1')}", - 'model_providers.miles.env_key="YETO_CODEX_BRIDGE_TOKEN"', - 'model_providers.miles.wire_api="responses"', - "model_providers.miles.requires_openai_auth=false", - "model_providers.miles.request_max_retries=0", - "model_providers.miles.stream_max_retries=0", - *_suspend_settings(), - "features.shell_tool=false", - "features.unified_exec=false", - "features.multi_agent=false", - "features.multi_agent_v2=false", - "features.apps=false", - "features.plugins=false", - "features.browser_use=false", - "features.computer_use=false", - "features.image_generation=false", - "features.default_mode_request_user_input=false", - "features.skill_search=false", - "features.workspace_dependencies=false", - "agents.enabled=false", - "orchestrator.skills.enabled=false", - "orchestrator.mcp.enabled=false", - "tools.experimental_request_user_input.enabled=false", - 'web_search="disabled"', - "include_permissions_instructions=false", - "include_environment_context=false", - "include_apps_instructions=false", - "include_collaboration_mode_instructions=false", - "skills.include_instructions=false", - "analytics.enabled=false", - "check_for_update_on_startup=false", - "project_doc_max_bytes=0", - f"tool_output_token_limit={CODEX_TOOL_OUTPUT_TOKEN_LIMIT}", - ] - values = [str(binary), "app-server", "--stdio", "--strict-config"] - for setting in settings: - values.extend(("-c", setting)) - return values - - -class _AppServerDriver: - def __init__( - self, - binary: Path, - bridge: _ResponsesBridge, - client: EpisodeClient, - episode_id: str, - prompt: str, - metrics: legacy.AgentMetrics, - ) -> None: - self._binary = binary - self._bridge = bridge - self._client = client - self._episode_id = episode_id - self._prompt = prompt - self._metrics = metrics - self._process: asyncio.subprocess.Process | None = None - self._process_group_id: int | None = None - self._request_id = 0 - self._thread_id: str | None = None - self._turn_id: str | None = None - self._stderr_task: asyncio.Task[None] | None = None - - async def __aenter__(self) -> Self: - isolated_home = tempfile.TemporaryDirectory(prefix="yeto-codex-") - self._isolated_home = isolated_home - Path(isolated_home.name).chmod(0o700) - environment = { - "CODEX_HOME": isolated_home.name, - "PATH": "/usr/bin:/bin", - "LANG": "C.UTF-8", - "TZ": "UTC", - "YETO_CODEX_BRIDGE_TOKEN": self._bridge.token, - } - self._process = await asyncio.create_subprocess_exec( - *_codex_argv(self._binary, self._bridge), - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - env=environment, - limit=MAX_APP_SERVER_FRAME_BYTES + 1, - start_new_session=True, - ) - self._process_group_id = self._process.pid - self._stderr_task = asyncio.create_task(self._drain_stderr()) - return self - - async def __aexit__( - self, - exc_type: object, - _exc: object, - _traceback: object, - ) -> None: - try: - await self._stop() - except Exception: - if exc_type is None: - raise - LOGGER.warning( - "Codex app-server shutdown also failed while handling an episode error", - exc_info=True, - ) - finally: - await self._cleanup_isolated_home() - - async def _drain_stderr(self) -> None: - if self._process is None or self._process.stderr is None: - return - while await self._process.stderr.read(8192): - pass - - async def _stop(self) -> None: - process = self._process - if process is None: - return - - process_group_id = self._process_group_id - group_signalled = self._signal_process_group(signal.SIGTERM) - if process.returncode is None and not group_signalled: - process.terminate() - if process.returncode is None: - try: - await asyncio.wait_for( - process.wait(), timeout=_CODEX_PROCESS_STOP_TIMEOUT_SECONDS - ) - except asyncio.TimeoutError: - if not self._signal_process_group(signal.SIGKILL): - process.kill() - await process.wait() - - # The app-server may have spawned plugin/helper processes which inherited - # CODEX_HOME. Waiting for only the direct child does not reap them. Give - # the isolated process group a short TERM grace period, then make sure no - # descendant remains able to recreate files while the home is removed. - if process_group_id is not None: - await asyncio.sleep(_CODEX_DESCENDANT_TERM_GRACE_SECONDS) - self._signal_process_group(signal.SIGKILL) - if self._stderr_task is not None: - await self._stderr_task - self._process_group_id = None - - def _signal_process_group(self, sig: signal.Signals) -> bool: - process_group_id = self._process_group_id - if process_group_id is None: - return False - try: - os.killpg(process_group_id, sig) - except ProcessLookupError: - return False - except OSError: - LOGGER.warning( - "failed to signal Codex app-server process group %s", - process_group_id, - exc_info=True, - ) - return False - return True - - async def _cleanup_isolated_home(self) -> None: - isolated_home = getattr(self, "_isolated_home", None) - if isolated_home is None: - return - - last_error: Exception | None = None - for delay_seconds in _CODEX_HOME_CLEANUP_RETRY_DELAYS_SECONDS: - if delay_seconds: - await asyncio.sleep(delay_seconds) - try: - isolated_home.cleanup() - return - except FileNotFoundError: - return - except OSError as exc: - last_error = exc - if exc.errno != errno.ENOTEMPTY: - break - except Exception as exc: # noqa: BLE001 - teardown must not mask outcome - last_error = exc - break - - # Temp-home teardown is resource hygiene, not episode validity. Once the - # app-server group has been reaped, make one non-throwing final removal so - # cleanup can neither invalidate a successful rollout nor mask its primary - # failure. A warning retains evidence if the path somehow remains. - if last_error is not None: - LOGGER.warning( - "Codex isolated HOME cleanup did not complete normally: %s", - last_error, - ) - try: - shutil.rmtree(isolated_home.name, ignore_errors=True) - if Path(isolated_home.name).exists(): - LOGGER.warning("Codex isolated HOME remains at %s", isolated_home.name) - except Exception: - LOGGER.warning( - "Codex isolated HOME final best-effort cleanup failed", - exc_info=True, - ) - - async def _send(self, value: dict[str, Any]) -> None: - if self._process is None or self._process.stdin is None: - raise CodexHarnessError("Codex app-server stdin is unavailable") - raw = _canonical_json(value) + b"\n" - self._process.stdin.write(raw) - await self._process.stdin.drain() - - async def _request(self, method: str, params: dict[str, Any]) -> dict[str, Any]: - self._request_id += 1 - request_id = self._request_id - await self._send({"id": request_id, "method": method, "params": params}) - while True: - message = await self._read_message() - if message.get("id") == request_id: - if "error" in message or not isinstance(message.get("result"), dict): - raise CodexHarnessError(f"Codex rejected app-server method {method}") - return message["result"] - await self._handle_message(message, allow_tools=False) - - async def _read_message(self) -> dict[str, Any]: - if self._process is None or self._process.stdout is None: - raise CodexHarnessError("Codex app-server stdout is unavailable") - while True: - if self._bridge.fatal.done(): - error = self._bridge.fatal.result() - raise error - if self._process.returncode is not None: - raise CodexHarnessError("Codex app-server exited before submit") - try: - line = await asyncio.wait_for( - self._process.stdout.readline(), timeout=0.1 - ) - except asyncio.TimeoutError: - continue - except (ValueError, asyncio.LimitOverrunError) as exc: - raise CodexHarnessError( - "Codex app-server emitted an oversized frame" - ) from exc - if self._bridge.fatal.done(): - raise self._bridge.fatal.result() - if not line: - raise CodexHarnessError("Codex app-server closed stdout before submit") - if len(line) > MAX_APP_SERVER_FRAME_BYTES: - raise CodexHarnessError("Codex app-server emitted an oversized frame") - try: - value = json.loads(line) - except json.JSONDecodeError as exc: - raise CodexHarnessError("Codex app-server emitted invalid JSON") from exc - if not isinstance(value, dict): - raise CodexHarnessError("Codex app-server emitted a non-object frame") - return value - - async def _handle_message(self, value: dict[str, Any], *, allow_tools: bool) -> str | None: - if self._bridge.fatal.done(): - raise self._bridge.fatal.result() - method = value.get("method") - if isinstance(value.get("id"), (int, str)) and method is not None: - if method != "item/tool/call" or not allow_tools: - raise CodexHarnessError("Codex app-server made an unexpected request") - return await self._handle_tool_request(value) - if method in {"thread/compacted", "error"}: - raise CodexHarnessError("Codex app-server reported a forbidden lifecycle event") - if method == "turn/completed": - raise CodexHarnessError("Codex completed without terminal submit") - return None - - async def _handle_tool_request(self, value: dict[str, Any]) -> str | None: - params = value.get("params") - request_id = value.get("id") - if not isinstance(params, dict): - raise CodexHarnessError("Codex supplied invalid tool request params") - pending = self._bridge.take_pending() - if ( - params.get("threadId") != self._thread_id - or params.get("turnId") != self._turn_id - or params.get("callId") != pending.call_id - or params.get("namespace") is not None - or params.get("tool") != pending.codex_name - ): - raise CodexHarnessError("Codex tool request identity drifted") - arguments = params.get("arguments") - if not isinstance(arguments, dict) or _canonical_json(arguments) != _canonical_json( - json.loads(pending.raw_arguments) - ): - raise CodexHarnessError("Codex mutated raw tool arguments") - started = time.monotonic() - terminal = False - if pending.codex_name == "terminal_exec": - output = await self._execute_terminal(arguments) - if "error" not in output: - self._metrics.terminal_calls += 1 - elif pending.codex_name == "submit": - output = await self._execute_submit(arguments) - if output.get("accepted") is True: - self._metrics.submit_calls += 1 - terminal = True - else: - raise CodexHarnessError("Codex requested a forbidden tool") - self._metrics.total_tool_time += time.monotonic() - started - rendered = json.dumps(output, ensure_ascii=False, separators=(",", ":")) - self._bridge.expect_tool_output(pending.call_id, rendered) - if terminal: - self._bridge.mark_terminal() - await self._send( - { - "id": request_id, - "result": { - "contentItems": [{"type": "inputText", "text": rendered}], - "success": True, - }, - } - ) - if terminal: - return "completed" - return None - - async def _execute_terminal(self, arguments: dict[str, Any]) -> dict[str, Any]: - if set(arguments) - {"command", "timeout_seconds"}: - return _tool_error("terminal_exec received unexpected arguments") - command = arguments.get("command") - if not isinstance(command, str) or not command.strip(): - return _tool_error("terminal_exec command must be non-empty") - maximum = legacy._positive_env("SECRLENV_TOOL_TIMEOUT_SECONDS", 120.0) - requested = arguments.get("timeout_seconds", maximum) - if isinstance(requested, bool) or not isinstance(requested, (int, float)): - return _tool_error("terminal_exec timeout is invalid") - timeout = float(requested) - if not math.isfinite(timeout) or not 1.0 <= timeout <= maximum: - return _tool_error("terminal_exec timeout is outside the signed range") - try: - result = await self._client.execute( - self._episode_id, - command, - timeout_seconds=timeout, - output_bytes=MAX_TOOL_OUTPUT_BYTES, - ) - except EpisodeAPIError as exc: - if exc.status == 400: - return _tool_error(exc.message) - raise - output = result.get("output", "") - truncated = result.get("truncated", False) - if ( - not isinstance(output, str) - or len(output.encode("utf-8")) > MAX_TOOL_OUTPUT_BYTES - or not isinstance(truncated, bool) - ): - raise CodexHarnessError("terminal_exec output exceeded its signed boundary") - return { - "exit_code": result.get("exit_code"), - "output": output or "(no output)", - "timed_out": bool(result.get("timed_out", False)), - "truncated": truncated, - } - - async def _execute_submit(self, arguments: dict[str, Any]) -> dict[str, Any]: - evidence = arguments.get("evidence") - if not isinstance(evidence, str) or not evidence: - return _tool_error("submit evidence must be non-empty") - try: - result = await self._client.submit(self._episode_id, arguments) - except EpisodeAPIError as exc: - if exc.status == 400: - return _tool_error(exc.message) - raise - if not isinstance(result, dict) or result.get("accepted") is not True: - raise CodexHarnessError("episode daemon did not accept submit") - return result - - async def drive(self) -> str: - try: - status = await self._drive_protocol() - except CodexSequenceLimit as exc: - self._metrics.end_reason = f"{type(exc).__name__}: {exc}"[:160] # type: ignore[attr-defined] - self._metrics.end_kind = end_kind_for(exc) # type: ignore[attr-defined] - return "max_seq_len" - except (CodexTurnLimit, CodexModelFailure) as exc: - # The signed status stays ``max_turns`` (reward 0, verifier still - # runs; changing it would change the signed outcome contract), but - # ``end_kind`` separates a real turn budget from a protocol error. - self._metrics.end_reason = f"{type(exc).__name__}: {exc}"[:160] # type: ignore[attr-defined] - self._metrics.end_kind = end_kind_for(exc) # type: ignore[attr-defined] - return "max_turns" - self._metrics.end_kind = ( # type: ignore[attr-defined] - END_KIND_SUBMIT if self._metrics.submit_calls else END_KIND_COMPLETED - ) - return status - - async def _drive_protocol(self) -> str: - initialized = await self._request( - "initialize", - { - "clientInfo": {"name": "yeto-secrlenv", "version": "1"}, - "capabilities": {"experimentalApi": True}, - }, - ) - user_agent = initialized.get("userAgent") - if not isinstance(user_agent, str) or "0.145.0" not in user_agent: - raise CodexHarnessError("Codex app-server reported the wrong version") - await self._send({"method": "initialized", "params": {}}) - started = await self._request( - "thread/start", - { - "model": BACKEND_MODEL, - "modelProvider": "miles", - "cwd": self._isolated_home.name, - "baseInstructions": base_instructions(), - "developerInstructions": None, - "dynamicTools": copy.deepcopy(dynamic_tools()), - "environments": [], - "ephemeral": True, - "approvalPolicy": "never", - "sandbox": "read-only", - "threadSource": "exec", - }, - ) - thread = started.get("thread") - self._thread_id = thread.get("id") if isinstance(thread, dict) else None - if not isinstance(self._thread_id, str) or not self._thread_id: - raise CodexHarnessError("Codex did not create an ephemeral thread") - turn = await self._request( - "turn/start", - { - "threadId": self._thread_id, - "input": [{"type": "text", "text": self._prompt, "text_elements": []}], - "effort": "xhigh", - "summary": "none", - "environments": [], - }, - ) - turn_value = turn.get("turn") - self._turn_id = turn_value.get("id") if isinstance(turn_value, dict) else None - if not isinstance(self._turn_id, str) or not self._turn_id: - raise CodexHarnessError("Codex did not start one turn") - while True: - message = await self._read_message() - result = await self._handle_message(message, allow_tools=True) - if result is not None: - return result - - -async def _drive_codex( - binary: Path, - miles_base_url: str, - client: EpisodeClient, - episode: dict[str, Any], - request_kwargs: dict[str, Any], - metrics: legacy.AgentMetrics, - *, - max_seq_len: int | None, -) -> str: - episode_id = episode.get("episode_id") - prompt = episode.get("prompt") - if not isinstance(episode_id, str) or not isinstance(prompt, str) or not prompt: - raise EpisodeClientError("episode daemon returned invalid episode identity") - async with _ResponsesBridge( - miles_base_url, - prompt, - request_kwargs, - metrics, - max_seq_len=max_seq_len, - ) as bridge, _AppServerDriver( - binary, bridge, client, episode_id, prompt, metrics - ) as driver: - return await driver.drive() - - -async def run( - base_url: str, - prompt: Any, - request_kwargs: dict[str, Any] | None = None, - metadata: dict[str, Any] | None = None, - **_kwargs: Any, -) -> dict[str, Any] | None: - """Run one stock-Codex/Miles SecRLEnv episode.""" - - del prompt - metadata = dict(metadata or {}) - request_kwargs = dict(request_kwargs or {}) - try: - task_id, tier = legacy._task_identity(metadata) - max_seq_len = legacy._metadata_max_seq_len(metadata) - infrastructure_503_retries = metadata.get( - INFRASTRUCTURE_503_RETRIES_KEY, 0 - ) - if ( - type(infrastructure_503_retries) is not int - or infrastructure_503_retries not in {0, 1} - ): - raise ValueError("invalid SecRLEnv infrastructure retry ledger") - binary = _attest_runtime() - except (ValueError, CodexHarnessError) as exc: - LOGGER.error("invalid signed Codex episode contract: %s", exc) - return None - - metrics = legacy.AgentMetrics() - episode_id: str | None = None - outcome_status = "completed" - infrastructure_failure = False - cleanup_failure = False - evaluation: dict[str, Any] | None = None - policy_task: asyncio.Task[str] | None = None - retry_ledger = { - INFRASTRUCTURE_503_RETRIES_KEY: infrastructure_503_retries - } - try: - async with EpisodeClient() as client: - started = time.monotonic() - episode = await legacy._create_with_capacity_retry( - client, task_id, tier, retry_ledger=retry_ledger - ) - metrics.create_time = time.monotonic() - started - episode_id = episode.get("episode_id") - if not isinstance(episode_id, str): - raise EpisodeClientError("episode daemon returned no episode ID") - policy_task = asyncio.create_task( - legacy._await_policy_and_claim_finalization( - episode_id, - _drive_codex( - binary, - base_url, - client, - episode, - request_kwargs, - metrics, - max_seq_len=max_seq_len, - ), - ) - ) - legacy._register_episode(episode_id, policy_task) - try: - outcome_status = await asyncio.wait_for( - policy_task, - timeout=legacy._positive_env( - "SECRLENV_MAX_ROLLOUT_TIME_SECONDS", 3600.0 - ), - ) - except asyncio.TimeoutError: - metrics.timed_out = 1 - metrics.end_kind = END_KIND_TIMEOUT # type: ignore[attr-defined] - outcome_status = "timeout" - except asyncio.CancelledError: - raise - except Exception: - infrastructure_failure = True - LOGGER.exception("stock Codex policy/tool loop failed") - - if legacy._claim_episode_finalization(episode_id): - started = time.monotonic() - try: - evaluation = await client.evaluate(episode_id) - except Exception: - infrastructure_failure = True - LOGGER.exception("secrlenv evaluation failed") - finally: - metrics.evaluate_time = time.monotonic() - started - started = time.monotonic() - try: - await client.close(episode_id) - except asyncio.CancelledError: - raise - except EpisodeAPIError as exc: - if exc.status == 404: - legacy._release_episode(episode_id) - else: - infrastructure_failure = True - cleanup_failure = ( - not await legacy._recover_failed_normal_close( - client, episode_id, policy_task - ) - ) - except Exception: # noqa: BLE001 - lifecycle cleanup is best-effort - infrastructure_failure = True - cleanup_failure = ( - not await legacy._recover_failed_normal_close( - client, episode_id, policy_task - ) - ) - else: - legacy._release_episode(episode_id) - finally: - metrics.close_time = time.monotonic() - started - metrics.total_tool_time += ( - metrics.create_time + metrics.evaluate_time + metrics.close_time - ) - else: - infrastructure_failure = True - except asyncio.CancelledError: - cleanup_claimed = episode_id is not None and ( - legacy._claim_episode_cleanup(episode_id) - or legacy._claim_cancelled_policy_cleanup(episode_id, policy_task) - ) - if cleanup_claimed: - try: - async with EpisodeClient(total_timeout_seconds=180.0) as cleanup_client: - await legacy._close_aborted_episode( - cleanup_client, - episode_id, - retry_timeout_seconds=legacy._positive_env( - "SECRLENV_ABORT_CLOSE_RETRY_SECONDS", 180.0 - ), - ) - except Exception: - LOGGER.warning("cancelled Codex episode cleanup failed", exc_info=True) - finally: - legacy._mark_episode_cleanup_pending(episode_id) - raise - except legacy._AdmissionInfrastructureExhausted: - if ( - episode_id is None - and retry_ledger[INFRASTRUCTURE_503_RETRIES_KEY] == 1 - ): - try: - outcome = legacy._admission_infrastructure_outcome( - task_id, - retry_ledger[INFRASTRUCTURE_503_RETRIES_KEY], - ) - return legacy._agent_result( - outcome, INFRASTRUCTURE_STATUS, metrics - ) - except Exception: - LOGGER.exception("refusing untrusted Codex admission metadata") - return None - LOGGER.error("invalid stock Codex admission infrastructure state") - return None - except EpisodeAPIError: - LOGGER.exception("stock Codex episode failed before a trustworthy verdict") - return None - except Exception: - LOGGER.exception("stock Codex episode failed before a trustworthy verdict") - return None - - if episode_id is None: - return None - try: - if cleanup_failure: - outcome_status = legacy.CLEANUP_ERROR_STATUS - outcome = legacy._cleanup_error_outcome( - task_id, - episode_id, - retry_ledger[INFRASTRUCTURE_503_RETRIES_KEY], - ) - elif infrastructure_failure or evaluation is None: - outcome_status = INFRASTRUCTURE_STATUS - outcome = legacy._infrastructure_outcome( - task_id, - episode_id, - retry_ledger[INFRASTRUCTURE_503_RETRIES_KEY], - ) - else: - outcome = legacy._validated_outcome( - evaluation, - task_id, - episode_id, - outcome_status, - infrastructure_503_retries=retry_ledger[ - INFRASTRUCTURE_503_RETRIES_KEY - ], - ) - result = { - OUTCOME_KEY: outcome, - "exit_status": outcome_status, - "agent_metrics": metrics.to_dict(), - } - capture_agent_metadata(result) - except Exception: - LOGGER.exception("refusing untrusted secrlenv evaluation metadata") - return None - return result - - -abort = legacy.abort diff --git a/yeto/rl/harness/codex/codex_openenv_agent_function.py b/yeto/rl/harness/codex/codex_openenv_agent_function.py deleted file mode 100644 index 74145a81a..000000000 --- a/yeto/rl/harness/codex/codex_openenv_agent_function.py +++ /dev/null @@ -1,681 +0,0 @@ -"""Terminal-Bench (OpenEnv) Codex adapter for the ports path. - -REWRITE, not a move: the legacy ``codex_openenv_*`` modules lived in a private -``agentenv/miles`` ``examples/experimental/openenv`` copy and were never found -(CODEX-PROGRESS §1.3). The interface below is the one the yeto consumers -expect (``yeto/rl/adapters/miles/island_entry.py``, ``yeto/rl/tbench_direct_preflight.py``, -``tests/test_rl_codex_schema.py``): ``stock``, ``_OPENENV_IDENTITY_ENV``, -``codex_openenv_harness_identity()`` and a module-level ``run``. - -Trust split (design D7): -- ``drive_untrusted`` runs Codex against a terminal environment and never runs - the verifier; -- ``finish_trusted`` runs the verifier and builds the ``tbench_outcome``. - -The stock driver and Responses bridge are reused unchanged from -``codex_harness_agent`` (moved from ``yeto_miles_secrlenv@5bfc011``). -""" - -from __future__ import annotations - -import asyncio -import json -import os -import secrets -import time -import uuid -from dataclasses import asdict -from pathlib import Path -from types import SimpleNamespace -from typing import Any, Awaitable, Callable, Mapping - -from yeto.rl.codex_backend import ( - QWEN3_30B_A3B_MODEL, - QWEN3_30B_A3B_REVISION, - QWEN35_08B_MODEL, - QWEN35_MODEL, - QWEN35_REVISION, - QWEN35_08B_REVISION, - QWEN38_MODEL, - QWEN38_REVISION, - QWEN38_NEXT_4LAYER_MODEL, - QWEN38_NEXT_4LAYER_REVISION, - QWEN38_NEXT_MODEL, - QWEN38_NEXT_REVISION, - stock_codex_backend_profile, -) -from yeto.rl.tbench_outcome import ( - NATIVE_VERIFIER, - TEST_SH_VERIFIER, - TIMEOUT_VERIFIER, - build_outcome_metadata, -) - -from . import agent as legacy -from . import codex_harness_agent as harness -from . import compaction_bridge -from .environment import TerminalEnvironment, TrustedVerifier - -# S17 M1: observe-only timing line appended to verifier_log (the tape keeps its 2000-char tail). -VERIFIER_LOG_TAIL_CHARS = 4000 -from .pins import OPENENV_BACKEND_PROFILE, OPENENV_BACKEND_PROFILES - -# rl-fn-codex-rollout 1.0: the backend profile is no longer pinned to the -# image default. Two distinct notions: -# * IMAGE_BACKEND_PROFILE_NAME -- the build-time default recorded in the image -# env (``YETO_CODEX_OPENENV_BACKEND_PROFILE`` / ``_MODEL_ID`` / ``_MODEL_REVISION``); -# * BACKEND_PROFILE_NAME -- the runtime profile declared by the launch -# (``--codex-backend-profile``), which the learner publishes to the container -# as ``YETO_CODEX_CHAT_TEMPLATE`` (``stock_codex_backend_contract(profile)["chat_template"]``, -# the same declaration ``codex_harness_agent`` reads) and ``validate_stock_codex_fields`` -# already checks against ``--model`` / ``--model-revision``. -# Why no image rebuild: the adapter's tool surface (``*_SHA256`` pins) does not -# depend on the model, so the image record only has to stay a faithful record of -# the build; the model identity of the runtime profile is derived here from the -# ``yeto.rl.codex_backend`` allowlist (``profile_identity``) instead of being -# compared against the image env. Boundary: the runtime profile must belong to -# ``pins.OPENENV_BACKEND_PROFILES`` and declare the exact HF identity listed in -# ``_PROFILE_IDENTITY``; anything else is refused at import / preflight. -IMAGE_BACKEND_PROFILE_NAME = OPENENV_BACKEND_PROFILE # "qwen35_08b" -RUNTIME_PROFILE_ENV = "YETO_CODEX_CHAT_TEMPLATE" -_PROFILE_IDENTITY: dict[str, tuple[str, str]] = { - "qwen35_08b": (QWEN35_08B_MODEL, QWEN35_08B_REVISION), - "qwen35": (QWEN35_MODEL, QWEN35_REVISION), - "qwen38": (QWEN38_MODEL, QWEN38_REVISION), - "qwen38_next": (QWEN38_NEXT_MODEL, QWEN38_NEXT_REVISION), - "qwen38_next_4layer": (QWEN38_NEXT_4LAYER_MODEL, QWEN38_NEXT_4LAYER_REVISION), - "qwen3_30b_a3b": (QWEN3_30B_A3B_MODEL, QWEN3_30B_A3B_REVISION), -} -assert set(_PROFILE_IDENTITY) == set(OPENENV_BACKEND_PROFILES) -INFRASTRUCTURE_STATUS = "infrastructure" -POLICY_STATUSES = frozenset({"completed", "timeout", "max_turns", "max_seq_len"}) -REWARD_SCOPE = "trajectory" - - -def profile_identity(profile_name: str) -> tuple[str, str]: - """``(model_identifier, model_revision)`` the adapter derives for one runtime profile. - - Raises ``ValueError`` ("requires backend profile") for a profile outside - ``OPENENV_BACKEND_PROFILES`` and ``RuntimeError`` ("profile drifted") when - the allowlisted ``codex_backend`` profile no longer declares the identity - pinned here (the legacy module-level qwen35_08b assertion, per profile). - """ - expected = _PROFILE_IDENTITY.get(profile_name) - if expected is None: - raise ValueError( - "the Codex OpenEnv adapter requires backend profile in " - + ", ".join(OPENENV_BACKEND_PROFILES) + f" (got {profile_name!r})" - ) - profile = stock_codex_backend_profile(profile_name) - if (profile.get("model_identifier"), profile.get("model_revision")) != expected: - raise RuntimeError(f"the Codex OpenEnv {profile_name} profile drifted") - return expected - - -def resolve_backend_profile(name: str | None = None, env: Mapping[str, str] | None = None) -> str: - """The runtime backend profile: explicit ``name``, else the launch declaration - ``YETO_CODEX_CHAT_TEMPLATE`` when it names an allowlisted profile (for every - member of ``OPENENV_BACKEND_PROFILES`` the chat template equals the profile - name), else the image default. An explicit name is validated.""" - if name is not None: - profile_identity(name) - return name - env = os.environ if env is None else env - declared = env.get(RUNTIME_PROFILE_ENV) - if declared in _PROFILE_IDENTITY: - return declared - return IMAGE_BACKEND_PROFILE_NAME - - -BACKEND_PROFILE_NAME = resolve_backend_profile() -profile_identity(BACKEND_PROFILE_NAME) # import-time drift check, per selected profile - -stock = SimpleNamespace( - _BACKEND_PROFILE=stock_codex_backend_profile(BACKEND_PROFILE_NAME), - BACKEND_MODEL=stock_codex_backend_profile(BACKEND_PROFILE_NAME)["model"], - _attest_runtime=harness._attest_runtime, -) - - -def codex_openenv_harness_identity() -> dict[str, str]: - """Live identity of the Codex surface this adapter drives (same tools as the stock harness).""" - live = harness.codex_harness_identity() - return {key: value for key, value in live.items() if key.endswith("_sha256")} - - -# Image record (what the image line bakes into the container env and -# ``yeto.rl.CODEX_OPENENV_IDENTITY_ENV`` mirrors): the *build-time* profile and -# its HF identity plus the profile-independent tool-surface hashes. It is -# deliberately NOT derived from BACKEND_PROFILE_NAME so the container env drift -# check stays a check of the image, whichever runtime profile the launch declares. -_OPENENV_IDENTITY_ENV: dict[str, str] = { - "YETO_CODEX_OPENENV_BACKEND_PROFILE": IMAGE_BACKEND_PROFILE_NAME, - "YETO_CODEX_OPENENV_MODEL_ID": _PROFILE_IDENTITY[IMAGE_BACKEND_PROFILE_NAME][0], - "YETO_CODEX_OPENENV_MODEL_REVISION": _PROFILE_IDENTITY[IMAGE_BACKEND_PROFILE_NAME][1], - **{ - f"YETO_CODEX_OPENENV_{name.upper()}": value - for name, value in codex_openenv_harness_identity().items() - }, -} - - -class ToolWaitEmitter: - """Per-tool-call tool-wait hooks (3.1). ``enter``/``exit`` are sync callables.""" - - def __init__(self, enter: Callable[[], None], exit: Callable[[], None]) -> None: - self.enter = enter - self.exit = exit - - -class _ToolWaitEnvironment: - """Wrap an environment so every tool execution is one tool-wait interval.""" - - def __init__(self, env: TerminalEnvironment, emitter: ToolWaitEmitter | None) -> None: - self._env = env - self._emitter = emitter - - async def _scoped(self, coro: Awaitable[dict[str, Any]]) -> dict[str, Any]: - if self._emitter is None: - return await coro - self._emitter.enter() - try: - return await coro - finally: - self._emitter.exit() - - async def execute(self, episode_id: str, command: str, timeout_seconds: float, output_bytes: int): - return await self._scoped( - self._env.execute(episode_id, command, timeout_seconds=timeout_seconds, output_bytes=output_bytes) - ) - - async def submit(self, episode_id: str, submission: dict[str, Any]): - return await self._scoped(self._env.submit(episode_id, submission)) - - -def _max_rollout_seconds() -> float: - return legacy._positive_env("OPENENV_MAX_ROLLOUT_TIME_SECONDS", 1800.0) - - -async def drive_untrusted( - job: dict[str, Any], - env: TerminalEnvironment, - *, - tool_wait: ToolWaitEmitter | None = None, - binary: Path | None = None, -) -> dict[str, Any]: - """Untrusted layer: run one Codex episode. Returns an UNSIGNED result. - - ``job`` = {base_url, prompt, request_kwargs, episode_id, max_seq_len?}. - Policy boundaries (turn/seq/timeout) are reported as statuses; harness - protocol failures raise ``CodexHarnessError`` and are infrastructure errors. - """ - metrics = legacy.AgentMetrics() - scoped_env = _ToolWaitEnvironment(env, tool_wait) - episode = {"episode_id": job["episode_id"], "prompt": job["prompt"]} - codex = binary if binary is not None else harness._attest_runtime() - segment_urls = job.get(compaction_bridge.SEGMENT_URLS_KEY) - if compaction_bridge.compactionrl_enabled(): - if not isinstance(segment_urls, list) or not all(isinstance(u, str) for u in segment_urls): - raise harness.CodexHarnessError("CompactionRL job has no pre-created segment sessions") - driving = compaction_bridge.drive_codex_compactionrl( - codex, - job["base_url"], - scoped_env, - episode, - dict(job.get("request_kwargs") or {}), - metrics, - max_seq_len=job.get("max_seq_len"), - segment_base_urls=segment_urls, - ) - elif segment_urls is not None: - raise harness.CodexHarnessError("segment sessions supplied without CompactionRL") - else: - driving = harness._drive_codex( - codex, - job["base_url"], - scoped_env, - episode, - dict(job.get("request_kwargs") or {}), - metrics, - max_seq_len=job.get("max_seq_len"), - ) - try: - status = await asyncio.wait_for(driving, timeout=_max_rollout_seconds()) - except asyncio.TimeoutError: - metrics.timed_out = 1 - metrics.end_kind = harness.END_KIND_TIMEOUT # type: ignore[attr-defined] - status = "timeout" - except harness.CodexHarnessError as exc: - # The rejection counters (G6a) must survive the failure: the trusted - # layer mirrors them onto the HarnessBoard before aborting the sample. - exc.metrics = _metrics_dict(metrics) # type: ignore[attr-defined] - raise - if status not in POLICY_STATUSES: - raise harness.CodexHarnessError(f"Codex driver returned unknown status {status!r}") - return {"status": status, "metrics": _metrics_dict(metrics), "episode_id": job["episode_id"]} - - -def _metrics_dict(metrics: legacy.AgentMetrics) -> dict[str, Any]: - # compaction_counters is {} unless the CompactionRL bridge ran. - extra: dict[str, Any] = {} - # Observe only (S15 r2 follow-up): why the Codex episode ended and the shape - # of the last model reply; set by the stock driver, absent otherwise. - # S17 C9: finer end reason and per-turn lengths (model reply tokens, context - # tokens after the reply, tool-result bytes). - for name in ("end_reason", "last_completion", "end_kind", "turn_completion_tokens", - "turn_context_tokens", "turn_tool_output_bytes", - # agentic-rollout-utilization 5.1/5.3: only when a gate was in use - *(harness.SUSPEND_METRIC_FIELDS if getattr(metrics, "suspensions", 0) else ())): - value = getattr(metrics, name, None) - if value is not None: - extra[name] = value - return {**asdict(metrics), **harness.tito_counters(metrics), **compaction_bridge.compaction_counters(metrics), **extra} - - -# --- CompactionRL segment sessions (trusted side; design D8) ---------------- -_SESSION_FIELDS = ("temperature", "top_p", "top_k") - - -def split_session_url(base_url: str) -> tuple[str, str]: - """``{router}/sessions/{id}`` -> (router, id).""" - head, sep, session_id = base_url.rstrip("/").rpartition("/sessions/") - if not sep or not head or not session_id or "/" in session_id: - raise ValueError(f"not a session-server session URL: {base_url!r}") - return head, session_id - - -def segment_session_body(request_kwargs: dict[str, Any]) -> dict[str, Any]: - """``CreateSessionRequest`` body with the rollout's own sampling defaults.""" - body: dict[str, Any] = {"evaluation": False} - for key in _SESSION_FIELDS: - value = request_kwargs.get(key) - if value is not None: - body[key] = int(value) if key == "top_k" else float(value) - return body - - -async def _post_json(url: str, body: dict[str, Any]) -> dict[str, Any]: - import aiohttp - - async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120)) as session: - async with session.post(url, json=body) as response: - if response.status != 200: - raise RuntimeError(f"session server returned HTTP {response.status}") - return await response.json() - - -async def create_segment_sessions( - base_url: str, - request_kwargs: dict[str, Any], - count: int, - *, - post: Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] | None = None, - delete: Callable[[str], Awaitable[None]] | None = None, -) -> tuple[list[str], list[str]]: - """Pre-create one session per allowed compaction; returns (ids, agent base URLs). - - Created here, in the trusted layer, so the untrusted worker can only use - sessions this rollout owns; ``codex_openenv_generate`` collects (and so - deletes) every listed session, used or not. - """ - router, _ = split_session_url(base_url) - body = segment_session_body(request_kwargs) - ids: list[str] = [] - try: - for _ in range(count): - reply = await (post or _post_json)(f"{router}/sessions", dict(body)) - session_id = reply.get("session_id") if isinstance(reply, dict) else None - if not isinstance(session_id, str) or not session_id or "/" in session_id: - raise RuntimeError("session server returned no session_id") - ids.append(session_id) - except BaseException: - # Pre-creation failed half-way: nobody will collect the ones already made. - await delete_segment_sessions(base_url, ids, delete=delete) - raise - return ids, [f"{router}/sessions/{session_id}" for session_id in ids] - - -async def _delete_session(url: str) -> None: - import aiohttp - - async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60)) as session: - async with session.delete(url) as response: - # Pinned session server (miles/rollout/session/sessions.py): 204 on - # success; an already collected/deleted session is a 404. - if response.status not in (200, 204, 404): - raise RuntimeError(f"session server returned HTTP {response.status}") - - -async def delete_segment_sessions( - base_url: str, - session_ids: Any, - *, - delete: Callable[[str], Awaitable[None]] | None = None, -) -> list[str]: - """Best-effort ``DELETE {router}/sessions/{id}`` for every id; returns the failures. - - Used when pre-created segment sessions will not reach - ``codex_openenv_generate`` (which otherwise collects and so deletes them). - Never raises: a cleanup failure must not mask the original error. - """ - if not isinstance(session_ids, (list, tuple)) or not session_ids: - return [] - try: - router, _ = split_session_url(base_url) - except ValueError: - return [str(session_id) for session_id in session_ids] - failed: list[str] = [] - for session_id in session_ids: - try: - await (delete or _delete_session)(f"{router}/sessions/{session_id}") - except Exception: # noqa: BLE001 - keep deleting the rest - failed.append(str(session_id)) - return failed - - -async def release_unreturned_segments( - base_url: str, - segments: dict[str, Any], - *, - delete: Callable[[str], Awaitable[None]] | None = None, -) -> list[str]: - """Delete the sessions listed in ``segments`` (the trusted metadata from - :func:`prepare_segment_sessions`) when it is not being returned.""" - return await delete_segment_sessions( - base_url, segments.get(compaction_bridge.SESSIONS_METADATA_KEY), delete=delete - ) - - -def _rlf_grade_from_evaluation(evaluation: dict[str, Any]) -> Any: - """RLF grade interface (rlf-codex-rl breakpoint 5); lazy import keeps the - neutral RLF package out of the TB2 import graph. A malformed grade raises - ``RolloutRejected`` (fail closed) instead of downgrading to the TB2 path.""" - from yeto.rl.rlf.preflight import trusted_grade_from_evaluation - - return trusted_grade_from_evaluation(evaluation) - - -def _finish_rlf( - untrusted: dict[str, Any], - grade: Any, - *, - task_id: str, - sample_id: str, - evaluate_time: float, - evaluation: dict[str, Any], - status: str, -) -> dict[str, Any]: - """Trusted finish for an RLF grade: close the episode via ``finish_episode`` - (reward/success/grader evidence + the RLF identity quadruple, written by the - trusted side from the per-lease verifier's identity) on top of the usual - outcome envelope, so tape consumers keep working. The ``policy_id`` - leg of the quadruple is not produced here: it rides the input metadata and - is copied onto the chains by ``codex_openenv_generate`` bookkeeping.""" - from yeto.rl.rlf.contracts import EpisodeLease, finish_episode - from yeto.rl.rlf.preflight import RolloutRejected - - evidence_map = grade.grader_evidence - rlf_task_id = evidence_map.get("rlf_task_id") - rlf_episode_id = evidence_map.get("rlf_episode_id") - rlf_trajectory_id = evidence_map.get("rlf_trajectory_id") - if not all(isinstance(value, str) and value - for value in (rlf_task_id, rlf_episode_id, rlf_trajectory_id)): - raise RolloutRejected("RLF grade is missing the identity quadruple") - if rlf_task_id != task_id: - raise RolloutRejected(f"RLF grade task {rlf_task_id!r} does not match rollout task {task_id!r}") - lease = EpisodeLease(task_id=rlf_task_id, trajectory_id=rlf_trajectory_id, - episode_id=rlf_episode_id) - reward = finish_episode(lease, grade) - metadata = build_outcome_metadata( - task_id=task_id, - sample_id=sample_id, - episode_id=untrusted["episode_id"], - status=status, - reward=float(grade.reward), - verifier=NATIVE_VERIFIER, - testsh_rc=None, - ) - metadata.update(reward.metadata) # success, reward, quadruple, grader_evidence, rlf_status - metadata["agent_metrics"] = dict(untrusted.get("metrics") or {}) - metadata["agent_metrics"]["evaluate_time"] = evaluate_time - metadata["exit_status"] = status - if isinstance(evaluation.get("log"), str): - # Observe only; the grader contract forbids the flag in it. - metadata["verifier_log"] = evaluation["log"][-VERIFIER_LOG_TAIL_CHARS:] - return metadata - - -async def finish_trusted( - untrusted: dict[str, Any], - verifier: TrustedVerifier, - *, - task_id: str, - sample_id: str, -) -> dict[str, Any]: - """Trusted layer: verify and build the outcome (reward in {0,1}). - - When the verifier's evaluation carries an RLF grade the episode is closed - via ``finish_episode`` (``_finish_rlf``); anything else takes the TB2 - path below, byte-identical to before the RLF branch existed. - """ - status = untrusted["status"] - if status not in POLICY_STATUSES: - raise ValueError(f"untrusted result has invalid status {status!r}") - evaluation: dict[str, Any] = {} - evaluate_time: float | None = None - if status == "timeout": - passed, testsh_rc, verifier_name = False, None, TIMEOUT_VERIFIER - else: - started = time.monotonic() - evaluation = await verifier.evaluate(untrusted["episode_id"]) - evaluate_time = verify_seconds = time.monotonic() - started # N13 metric + M1 YETO_TIMING - grade = _rlf_grade_from_evaluation(evaluation) - if grade is not None: - return _finish_rlf(untrusted, grade, task_id=task_id, sample_id=sample_id, - evaluate_time=evaluate_time, evaluation=evaluation, status=status) - passed = bool(evaluation.get("passed")) - testsh_rc = evaluation.get("testsh_rc") - verifier_name = ( - TEST_SH_VERIFIER if isinstance(testsh_rc, int) else NATIVE_VERIFIER - ) - metadata = build_outcome_metadata( - task_id=task_id, - sample_id=sample_id, - episode_id=untrusted["episode_id"], - status=status, - reward=1.0 if passed else 0.0, - verifier=verifier_name, - testsh_rc=testsh_rc if isinstance(testsh_rc, int) else None, - ) - metadata["agent_metrics"] = dict(untrusted.get("metrics") or {}) - # S17 N13: the trusted verifier owns grading, so it owns the grading time - # (seconds, observe only). The worker's own ``evaluate_time`` is - # 0.0 on this path; a timeout skips grading and drops the key. - if evaluate_time is None: - metadata["agent_metrics"].pop("evaluate_time", None) - else: - metadata["agent_metrics"]["evaluate_time"] = evaluate_time - metadata["exit_status"] = status - if status != "timeout" and isinstance(evaluation.get("log"), str): - metadata["verifier_log"] = evaluation["log"] # observe only - # S17 M1 (observe only): verifier wall time and agent phase times - # (per-turn shapes would need the pinned codex_harness_agent.py), appended as the log's last line so the tape's - # verifier_log tail carries them. - metrics = untrusted.get("metrics") or {} - timing = {"verify_s": round(verify_seconds, 3)} - for name in ("create_time", "total_generation_time", "total_tool_time", "max_model_total_tokens"): - if isinstance(metrics.get(name), (int, float)) and not isinstance(metrics.get(name), bool): - timing[name] = round(float(metrics[name]), 3) - metadata["verifier_log"] = (metadata["verifier_log"] + "\nYETO_TIMING " - + json.dumps(timing, separators=(",", ":")))[-VERIFIER_LOG_TAIL_CHARS:] - return metadata - - -def infrastructure_metadata(reason: str, *, episode_id: str | None, metrics: dict[str, Any] | None = None) -> dict[str, Any]: - """Infrastructure marker: no reward, sample must be ABORTED (never a 0 reward).""" - return { - "tbench_infrastructure_error": reason, - "exit_status": INFRASTRUCTURE_STATUS, - "episode_id": episode_id, - "agent_metrics": dict(metrics or {}), - } - - -TITO_SESSION_MISMATCH_KEY = harness.TITO_SESSION_MISMATCH_KEY -TITO_CHAIN_BREAKS_KEY = harness.TITO_CHAIN_BREAKS_KEY - - -def mirror_tito_counters(metrics: dict[str, Any] | None, harness_board: Any) -> dict[str, Any]: - """G6(a): replay the bridge's rejection counters onto the HarnessBoard. - - Same method names and counter names as ``gateway/core.py`` so the 1.7 load - sample (``tito_session_mismatch`` / ``tito_chain_breaks``) reads the A path - too. Returns the trajectory-field overrides: the bridge holds one chain and - never forks, so ``chains_total`` stays 1 and ``chain_break_reason`` names the - first break (the sample is infrastructure-aborted, never trained on). - """ - from yeto.rl.engine.rollout_meta import counter_value - - metrics = metrics or {} - mismatches = counter_value(metrics.get(TITO_SESSION_MISMATCH_KEY)) - breaks = dict(metrics.get(TITO_CHAIN_BREAKS_KEY) or {}) - if harness_board is not None: - if mismatches: - _board_call(harness_board, "record_session_mismatch", mismatches) - for reason, count in breaks.items(): - if count: - _board_call(harness_board, "record_chain_break", reason, int(count)) - first_break = next((reason for reason, count in breaks.items() if count), None) - return {"chain_break_reason": first_break} - - -def _board_call(target: Any, method: str, *args: Any) -> Any: - """Call a board method and wait for it (A-T3-7 follow-up, S14/A19). - - The island boards are Ray actors with ``max_concurrency=64``: a - fire-and-forget ``record_session_mismatch`` / ``record_chain_break`` may - land after the round's ``snapshot()`` (counter missing from the rollout it - belongs to) and any error it raises is an unhandled actor error instead of - this trajectory's. Resolving the ObjectRef here keeps the calls in issue - order, like the tool-wait board's enter/exit. A failure of the remote - bookkeeping is logged and never fails the trajectory; a local board's own - exception (e.g. an unknown chain-break reason) propagates as before.""" - import sys - - from yeto.rl.engine.tool_wait import _resolve - - fn = getattr(target, method) - remote = getattr(fn, "remote", None) - if not callable(remote): - return fn(*args) - try: - return _resolve(remote(*args)) - except Exception as exc: # noqa: BLE001 - never fail a trajectory on telemetry - print(f"[codex-harness] harness board {method} failed: {type(exc).__name__}: {exc}", - file=sys.stderr) - return None - - -def trajectory_fields(trajectory_id: str) -> dict[str, Any]: - """D8 / R-D5a bookkeeping carried on every sample of the trajectory.""" - return { - "trajectory_id": trajectory_id, - "chain_index": 0, - "chains_total": 1, - "chain_break_reason": None, - "segment_id": 0, - "segment_boundary_reason": None, - "reward_scope": REWARD_SCOPE, - } - - -def task_identity(metadata: dict[str, Any]) -> tuple[str, str]: - task_id = metadata.get("task_id") - if not isinstance(task_id, str) or not task_id: - raise ValueError("Terminal-Bench sample metadata must carry task_id") - sample_id = metadata.get("sample_id") or metadata.get("trajectory_id") or f"tb-{uuid.uuid4().hex}" - return task_id, str(sample_id) - - -def new_episode_id() -> str: - return f"ep-{secrets.token_hex(8)}" - - -async def run( - base_url: str, - prompt: Any, - request_kwargs: dict[str, Any] | None = None, - metadata: dict[str, Any] | None = None, - *, - environment: TerminalEnvironment | None = None, - verifier: TrustedVerifier | None = None, - tool_wait: ToolWaitEmitter | None = None, - harness_board: Any = None, - **_kwargs: Any, -) -> dict[str, Any] | None: - """In-process entry (both layers in one process; use the subprocess entry in training).""" - metadata = dict(metadata or {}) - if environment is None or verifier is None: - raise ValueError("in-process run requires an environment and a verifier") - task_id, sample_id = task_identity(metadata) - trajectory_id = metadata.get("trajectory_id") or sample_id - episode_id = new_episode_id() - job = { - "base_url": base_url, - "prompt": metadata.get("prompt") if isinstance(metadata.get("prompt"), str) else str(prompt), - "request_kwargs": dict(request_kwargs or {}), - "episode_id": episode_id, - "max_seq_len": metadata.get("max_seq_len"), - } - try: - segments = await prepare_segment_sessions(job) - except Exception as exc: # noqa: BLE001 - session-server failures are infrastructure - return { - **infrastructure_metadata(f"segment sessions: {type(exc).__name__}: {exc}", episode_id=episode_id), - **trajectory_fields(trajectory_id), - } - handed_off = False # True once ``segments`` rides in the returned metadata - try: - try: - untrusted = await drive_untrusted(job, environment, tool_wait=tool_wait) - except (harness.CodexHarnessError, legacy.EpisodeClientError, OSError) as exc: - metrics = getattr(exc, "metrics", None) - tito = mirror_tito_counters(metrics, harness_board) - handed_off = True - return { - **infrastructure_metadata(f"{type(exc).__name__}: {exc}", episode_id=episode_id, metrics=metrics), - **trajectory_fields(trajectory_id), - **tito, - **segments, - } - tito = mirror_tito_counters(untrusted.get("metrics"), harness_board) - finished = await finish_trusted(untrusted, verifier, task_id=task_id, sample_id=sample_id) - handed_off = True - return {**finished, **trajectory_fields(trajectory_id), **tito, **segments} - finally: - if segments and not handed_off: - await release_unreturned_segments(job["base_url"], segments) - - -async def prepare_segment_sessions( - job: dict[str, Any], - *, - post: Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] | None = None, - delete: Callable[[str], Awaitable[None]] | None = None, -) -> dict[str, Any]: - """CompactionRL only: add segment URLs to ``job``; return the trusted metadata. - - Returns ``{}`` (and leaves ``job`` untouched) when CompactionRL is off. - """ - if not compaction_bridge.compactionrl_enabled(): - return {} - if job.get("max_seq_len") is None: - raise harness.CodexHarnessError("CompactionRL requires max_seq_len") - cfg = compaction_bridge.compaction_config(int(job["max_seq_len"])) - ids, urls = await create_segment_sessions( - job["base_url"], dict(job.get("request_kwargs") or {}), cfg.max_compactions, post=post, delete=delete - ) - job[compaction_bridge.SEGMENT_URLS_KEY] = urls - router, _ = split_session_url(job["base_url"]) - return {compaction_bridge.SESSIONS_METADATA_KEY: ids, - compaction_bridge.SESSIONS_ROUTER_METADATA_KEY: router} diff --git a/yeto/rl/harness/codex/codex_openenv_agent_worker.py b/yeto/rl/harness/codex/codex_openenv_agent_worker.py deleted file mode 100644 index f9733be4b..000000000 --- a/yeto/rl/harness/codex/codex_openenv_agent_worker.py +++ /dev/null @@ -1,81 +0,0 @@ -"""Untrusted worker process for the Terminal-Bench Codex adapter (REWRITE). - -Protocol (stdin/stdout, one JSON object per line): -- stdin: one job ``{"base_url", "prompt", "request_kwargs", "episode_id", - "max_seq_len", "env_url", "env_token", "driver": "stock"|"scripted", ...}``; -- stdout events: ``{"event": "tool_wait", "phase": "enter"|"exit"}`` for every - tool execution, ``{"event": "suspend", "phase": "enter"|"exit"}`` around a - wait at the suspended model-turn gate, then exactly one ``{"event": "result", ...}`` or - ``{"event": "error", "reason": ...}``. - -The worker refuses to start when a reward HMAC key is visible (trust split). -``driver: "scripted"`` is a CPU-test driver (no Codex, no Miles) and is only -honoured when ``YETO_CODEX_OPENENV_ALLOW_SCRIPTED_DRIVER=1``; preflight fails -closed when that variable is set. -""" - -from __future__ import annotations - -import asyncio -import json -import os -import sys -from typing import Any - -from . import codex_openenv_agent_function as adapter -from .environment import HttpTerminalEnvironment - -SCRIPTED_DRIVER_ENV = "YETO_CODEX_OPENENV_ALLOW_SCRIPTED_DRIVER" - - -def _emit(payload: dict[str, Any]) -> None: - sys.stdout.write(json.dumps(payload, sort_keys=True) + "\n") - sys.stdout.flush() - - -async def _scripted_drive(job: dict[str, Any], env: Any, tool_wait: adapter.ToolWaitEmitter) -> dict[str, Any]: - """Execute ``job['script']`` commands then submit; mimics the driver's tool loop.""" - scoped = adapter._ToolWaitEnvironment(env, tool_wait) - for command in job.get("script") or []: - await scoped.execute(job["episode_id"], str(command), timeout_seconds=30.0, output_bytes=4096) - if job.get("hang_seconds"): - await asyncio.sleep(float(job["hang_seconds"])) - await scoped.submit(job["episode_id"], {"evidence": "scripted"}) - return {"status": job.get("final_status", "completed"), "metrics": {"turns": len(job.get("script") or []) + 1}, "episode_id": job["episode_id"]} - - -async def _main_async(job: dict[str, Any]) -> dict[str, Any]: - tool_wait = adapter.ToolWaitEmitter( - enter=lambda: _emit({"event": "tool_wait", "phase": "enter"}), - exit=lambda: _emit({"event": "tool_wait", "phase": "exit"}), - ) - # agentic-rollout-utilization 5.1: parked between model turns (gate closed) - adapter.harness.set_suspend_event_sink(lambda phase: _emit({"event": "suspend", "phase": phase})) - async with HttpTerminalEnvironment(job["env_url"], job["env_token"]) as env: - if job.get("driver") == "scripted": - if os.getenv(SCRIPTED_DRIVER_ENV) != "1": - raise RuntimeError("scripted driver is not allowed in this process") - return await _scripted_drive(job, env, tool_wait) - return await adapter.drive_untrusted(job, env, tool_wait=tool_wait) - - -def main(argv: list[str] | None = None) -> int: - del argv - try: - job = json.loads(sys.stdin.readline()) - if not isinstance(job, dict): - raise ValueError("worker job must be a JSON object") - result = asyncio.run(_main_async(job)) - except BaseException as exc: # noqa: BLE001 - every failure is reported on the wire - event = {"event": "error", "reason": f"{type(exc).__name__}: {exc}"} - metrics = getattr(exc, "metrics", None) - if isinstance(metrics, dict): - event["metrics"] = metrics # G6a: rejection counters reach the trusted layer - _emit(event) - return 1 - _emit({"event": "result", **result}) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/yeto/rl/harness/codex/codex_openenv_generate.py b/yeto/rl/harness/codex/codex_openenv_generate.py deleted file mode 100644 index e4f781d5a..000000000 --- a/yeto/rl/harness/codex/codex_openenv_generate.py +++ /dev/null @@ -1,420 +0,0 @@ -"""``--custom-generate-function-path`` wrapper around upstream ``agentic_tool_call.generate``. - -Adds the R-D5a / R-IR-3 bookkeeping that the upstream session path does not -know about: - -- sibling samples of one trajectory get the same ``group_index`` and - ``rollout_id`` so ``train_data_conversion`` shares one reward and counts one - baseline entry per trajectory (never per chain); -- ``expected_policy_version`` (IR-3: ``rollout_meta_hook.expected_policy_version``, - i.e. prompt metadata first, else the driver token published through the - metadata sink by ``MilesRolloutPool.generate``) is compared against the - weight versions SGLang reported for every generation; under the default - policy-age limit 0 any drift aborts the trajectory (metadata ``policy_age_violation=1`` -> ``harness_counters`` -> - driver ``PolicyIdentityError``) instead of giving it a 0 reward; a missing - token refuses generation before upstream is called; under a limit > 0 - (agentic-rollout-utilization 5.2: a trajectory suspended between model turns - continues under the next version) the versions are recorded as segments and - checked against the version current when the trajectory finishes: oldest at - most ``limit`` behind, none newer; outside the window the trajectory is - discarded (infrastructure ``policy_age_exceeded``, not a policy-identity - error); -- mask/token/logprob alignment is asserted on every sample (``alignment``). - -The upstream function is injected (``upstream``) so this wrapper is testable -without the Miles package. -""" - -from __future__ import annotations - -import asyncio -import dataclasses -import os -import uuid -from typing import Any, Awaitable, Callable - -from yeto.rl.engine import rollout_meta -from yeto.rl.engine.policy_versions import engine_version_matches, engine_version_number -from yeto.rl.rlf.contracts import EpisodeLease - -from .alignment import AlignmentError, assert_sample_alignment -from .tbench_reward import INFRASTRUCTURE_KEY - -UPSTREAM_PATH = "miles.rollout.generate_hub.agentic_tool_call.generate" -POLICY_AGE_KEY = rollout_meta.POLICY_AGE_VIOLATION_KEY # "policy_age_violation" -EXPECTED_VERSION_KEY = rollout_meta.EXPECTED_POLICY_VERSION_KEY -ACTUAL_VERSIONS_KEY = "policy_versions_actual" - - -class PolicyVersionMissing(RuntimeError): - """IR-3: no target policy token for this rollout (driver did not publish one).""" - - -def expected_policy_version(input_sample: Any) -> str | None: - """IR-3 target token for ``input_sample`` (metadata first, else the driver sink).""" - try: - return rollout_meta.expected_policy_version(input_sample) - except Exception: # noqa: BLE001 - unreachable sink == nothing published - meta = getattr(input_sample, "metadata", None) - return str(meta[EXPECTED_VERSION_KEY]) if isinstance(meta, dict) and meta.get(EXPECTED_VERSION_KEY) else None - - -def _glue(): - """Backend glue: upstream generate, aborted status, session collection (decoupling 5.7).""" - from yeto.rl.engine import backends - - return backends.module("harness_glue") - - -def _load_upstream() -> Callable[[Any], Awaitable[Any]]: - return _glue().load_agentic_upstream() - - -def _samples_of(output: Any) -> list[Any]: - samples = getattr(output, "samples", output) - return list(samples) if isinstance(samples, list) else [samples] - - -def _mark_aborted(sample: Any, reason: str) -> None: - try: - sample.status = _glue().aborted_status() - except ImportError: # backend not installed (CPU tests) - sample.status = "ABORTED" - metadata = sample.metadata if isinstance(getattr(sample, "metadata", None), dict) else {} - metadata.pop("tbench_trusted_outcome", None) - metadata[INFRASTRUCTURE_KEY] = reason - sample.metadata = metadata - - -def actual_versions(sample: Any) -> list[str]: - versions: list[str] = [] - for call in getattr(sample, "weight_versions", None) or []: - spans = getattr(call, "spans", None) - if spans is None and isinstance(call, dict): - spans = call.get("spans", []) - for span in spans or []: - version = getattr(span, "version", None) - if version is None and isinstance(span, dict): - version = span.get("version") - if version is not None: - versions.append(str(version)) - return versions - - -POLICY_AGE_EXCEEDED = "policy_age_exceeded" -VERSION_SEGMENTS_KEY = "policy_version_segments" - - -def window_problem(versions: list[str], current_token: str | None, limit: int) -> str | None: - """Limit > 0: why ``versions`` (one per generation span) cannot be trained at - ``current_token`` (None = within the window).""" - from yeto.rl.engine.policy_age import token_version - - current = token_version(str(current_token or "")) - if current is None: - return f"current policy token {current_token!r} has no version" - if not versions: - return "no generation version recorded" - numbers = [engine_version_number(v) for v in versions] - if None in numbers: - return f"unparsable generation versions {sorted(set(versions))}" - if max(numbers) > current: - return f"generation version {max(numbers)} is newer than the current {current}" - if current - min(numbers) > limit: - return f"oldest generation version {min(numbers)} is {current - min(numbers)} behind {current} (limit {limit})" - return None - - -RLF_TASK_KEY = "rlf_task_id" -RLF_IDENTITY_KEYS = ("rlf_task_id", "rlf_episode_id", "rlf_trajectory_id", "policy_id") - - -def produce_rlf_identity(input_meta: dict[str, Any], *, expected_version: str | None = None) -> dict[str, Any]: - """Producer for the RLF identity quadruple on catalog-sourced rows. - - Rows emitted by ``yeto.rl.rlf.catalog.prompt_rows`` carry ``rlf_task_id`` - only; the per-rollout episode/trajectory ids are produced here (one - ``EpisodeLease`` per prompt draw) so episodes are never shared across - rollouts, and ``policy_id`` defaults to the IR-3 expected policy version. - Non-RLF rows (no ``rlf_task_id``) pass through unchanged, and ids already - present (e.g. written by the trusted finish) are never overwritten. - """ - if not input_meta.get(RLF_TASK_KEY): - return input_meta - meta = dict(input_meta) - lease = EpisodeLease( - task_id=str(meta[RLF_TASK_KEY]), - trajectory_id=str(meta.get("rlf_trajectory_id") or f"rlf-traj-{uuid.uuid4().hex}"), - ) - if not meta.get("rlf_episode_id"): - meta["rlf_episode_id"] = lease.episode_id - if not meta.get("rlf_trajectory_id"): - meta["rlf_trajectory_id"] = lease.trajectory_id - if not meta.get("policy_id") and expected_version is not None: - meta["policy_id"] = str(expected_version) - return meta - - -def apply_trajectory_bookkeeping(input_sample: Any, samples: list[Any], *, expected_version: str | None = None, - max_policy_age: int = 0, current_version: str | None = None) -> None: - """Siblings share ``group_index``/``rollout_id``; policy version and alignment checks. - - ``max_policy_age`` > 0 (5.2): versions are recorded as segments and checked - against ``current_version`` (the policy token current when the trajectory - finished; default ``expected_version``) instead of requiring one version.""" - group_index = getattr(input_sample, "group_index", None) - rollout_key = getattr(input_sample, "index", None) - input_meta = getattr(input_sample, "metadata", None) - input_meta = input_meta if isinstance(input_meta, dict) else {} - if expected_version is None: - expected_version = expected_policy_version(input_sample) - input_meta = produce_rlf_identity(input_meta, expected_version=expected_version) - for position, sample in enumerate(samples): - if group_index is not None: - sample.group_index = group_index - if rollout_key is not None: - sample.rollout_id = rollout_key - metadata = sample.metadata if isinstance(getattr(sample, "metadata", None), dict) else {} - metadata.setdefault("chain_index", position) - metadata["chains_total"] = len(samples) - # RLF identity (produced above for catalog rows, or written upstream) - # is copied onto every Codex chain so trusted grading and segment - # aggregation cannot depend on framework-only sample fields. - for key in RLF_IDENTITY_KEYS: - if key in input_meta: - metadata.setdefault(key, input_meta[key]) - sample.metadata = metadata - if INFRASTRUCTURE_KEY in metadata: - continue - versions = actual_versions(sample) - metadata[ACTUAL_VERSIONS_KEY] = versions - if max_policy_age > 0: - current = current_version if current_version is not None else expected_version - metadata[EXPECTED_VERSION_KEY] = current - metadata[VERSION_SEGMENTS_KEY] = sorted(set(versions)) - problem = window_problem(versions, current, max_policy_age) - if problem is not None: - _mark_aborted(sample, f"{POLICY_AGE_EXCEEDED}: {problem}") - continue - elif expected_version is not None: - metadata[EXPECTED_VERSION_KEY] = expected_version - if not versions or not all( - engine_version_matches(v, str(expected_version)) for v in versions - ): - metadata[POLICY_AGE_KEY] = 1 - _mark_aborted(sample, f"policy_age_violation: expected {expected_version}, saw {versions}") - continue - try: - assert_sample_alignment(sample) - except AlignmentError as exc: - _mark_aborted(sample, f"alignment: {exc}") - - -# --- CompactionRL segments (design D8; progress.md "S13 Codex 桥压缩拦截") ---- -SESSIONS_KEY = "codex_compaction_sessions" # trusted: pre-created by the agent function -ROUTER_KEY = "codex_compaction_router" # trusted: router those sessions live on -SEGMENT_COLLECT_TIMEOUT_ENV = "YETO_CODEX_SEGMENT_COLLECT_TIMEOUT_S" -DEFAULT_SEGMENT_COLLECT_TIMEOUT_S = 120.0 - - -def segment_collect_timeout_s() -> float: - """Per-session collect budget (env override, positive seconds).""" - raw = os.environ.get(SEGMENT_COLLECT_TIMEOUT_ENV) - try: - value = float(raw) if raw else DEFAULT_SEGMENT_COLLECT_TIMEOUT_S - except ValueError: - value = DEFAULT_SEGMENT_COLLECT_TIMEOUT_S - return value if value > 0 else DEFAULT_SEGMENT_COLLECT_TIMEOUT_S -COMPACTION_METRICS_KEY = "codex_compaction" # untrusted bridge record in agent_metrics -SegmentCollector = Callable[[Any, str, str], Awaitable[tuple[list[Any], dict[str, Any]]]] - - -async def collect_segment_session(input: Any, router: str, session_id: str) -> tuple[list[Any], dict[str, Any]]: - """Collect (and delete) one extra session exactly like upstream's tracer does (backend glue).""" - return await _glue().collect_segment_session(input, router, session_id) - - -def _optimized_tokens(sample: Any) -> int: - mask = getattr(sample, "loss_mask", None) or [] - return int(sum(1 for value in mask if value)) - - -def refresh_segment_counts(samples: Any) -> int: - """Recompute ``segment_tokens``/``tokens_after``/``gae_length`` from the final loss masks. - - The wrapper counts optimised tokens when it assembles the segments; later - hooks still set loss_mask 1 -> 0 (S19 #8 placeholder-logprob stop tokens, - S17 WP6 reasoning tokens) before Miles converts the samples. N_{>s} and l - are optimised-token counts, so they are refreshed after those hooks - (s19-compaction-g1-20261010g: tokens_after was 1-17 tokens high). - Returns the number of segment samples whose counts changed. - """ - - rollouts: dict[Any, list[Any]] = {} - for sample in samples: - meta = getattr(sample, "metadata", None) - if not isinstance(meta, dict) or not isinstance(meta.get("segment_index"), int): - continue - key = (getattr(sample, "rollout_id", None), getattr(sample, "group_index", None), - getattr(sample, "index", None)) - rollouts.setdefault(key, []).append(sample) - changed = 0 - for segments in rollouts.values(): - segments.sort(key=lambda s: s.metadata["segment_index"]) - tokens = [_optimized_tokens(s) for s in segments] - for index, sample in enumerate(segments): - new = {"segment_tokens": tokens[index], "tokens_after": sum(tokens[index + 1:]), - "gae_length": sum(tokens)} - if any(sample.metadata.get(k) != v for k, v in new.items()): - changed += 1 - sample.metadata.update(new) - return changed - - -async def assemble_compaction_segments(input: Any, output: Any, *, collect: SegmentCollector | None = None) -> Any: - """Collect the pre-created segment sessions; one sample per segment. - - Every listed session is collected (which deletes it) even when the - trajectory failed. Segment ``s`` carries the rollout's metadata (signed - outcome -> shared reward) plus ``segment_index``/``num_segments``/ - ``segment_tokens``/``tokens_after`` (N_{>s}, optimised tokens of later - segments)/``gae_length`` (optimised tokens of the whole rollout). - """ - samples = _samples_of(output) - if len(samples) != 1: - raise ValueError("CompactionRL expects one upstream sample for segment 0") - first = samples[0] - meta0 = first.metadata if isinstance(getattr(first, "metadata", None), dict) else {} - session_ids = meta0.get(SESSIONS_KEY) - if session_ids is None: - return output - collect = collect or collect_segment_session - # The agent function records the router it created the sessions on - # (s19-compaction-g1-20261010d: session_server_id is not in the returned - # metadata, the collect URL became "http:///sessions/..." and retried forever). - router = meta0.get(ROUTER_KEY) - if not router and meta0.get("session_server_id"): - router = f"http://{meta0['session_server_id']}" - if not isinstance(router, str) or not router.startswith(("http://", "https://")) or router.endswith("//"): - _mark_aborted(first, "segment collect: no session-server router for the segment sessions") - return output - collected: list[tuple[list[Any], dict[str, Any]]] = [] - failure: str | None = None - for session_id in session_ids if isinstance(session_ids, list) else []: - try: - collected.append(await asyncio.wait_for(collect(input, router, str(session_id)), - timeout=segment_collect_timeout_s())) - except asyncio.TimeoutError: - # s19-compaction-g1-20261010d: Miles' http_utils retries a failing - # collect for a long time; bound it so one rollout cannot stall the round. - failure = failure or f"segment collect: timed out after {segment_collect_timeout_s():g} s" - collected.append(([], {})) - except Exception as exc: # noqa: BLE001 - keep deleting the remaining sessions - failure = failure or f"segment collect: {type(exc).__name__}: {exc}" - collected.append(([], {})) - if INFRASTRUCTURE_KEY in meta0: - return output # already aborted; extra sessions were only drained - if not isinstance(session_ids, list): - failure = "segment sessions metadata is not a list" - record = (meta0.get("agent_metrics") or {}).get(COMPACTION_METRICS_KEY) - compactions = record.get("compactions") if isinstance(record, dict) else None - if failure is None and (type(compactions) is not int or not 0 <= compactions <= len(collected)): - failure = "segment record missing or out of range" - if failure is None: - used, unused = collected[:compactions], collected[compactions:] - if any(len(got) != 1 for got, _ in used) or any(got for got, _ in unused): - failure = f"segment sessions do not match {compactions} compactions" - if failure is not None: - _mark_aborted(first, failure) - return output - segments = [first] - for got, session_metadata in collected[:compactions]: - sample = got[0] - sample.metadata = {**meta0, **session_metadata} - segments.append(sample) - tokens = [_optimized_tokens(sample) for sample in segments] - truncated = meta0.get("exit_status") == "max_seq_len" - for index, sample in enumerate(segments): - sample.metadata.update( - { - "segment_index": index, - "num_segments": len(segments), - "segment_tokens": tokens[index], - "tokens_after": sum(tokens[index + 1 :]), - "gae_length": sum(tokens), - "compactions": compactions, - "truncated": truncated, - "chain_index": index, - "chains_total": len(segments), - "segment_id": index, - "chain_break_reason": "compaction_window" if index else None, - "segment_boundary_reason": "compaction" if index < len(segments) - 1 else None, - } - ) - if len(segments) > 1: - # GenerateFnOutput is a frozen dataclass ("multiple samples" is its documented use). - output = dataclasses.replace(output, samples=segments) if dataclasses.is_dataclass(output) else segments - return output - - -async def generate( - input: Any, - *, - upstream: Callable[[Any], Awaitable[Any]] | None = None, - collect: SegmentCollector | None = None, -) -> Any: - expected_version = expected_policy_version(input.sample) - if expected_version is None: # IR-3: refuse before any generation happens - raise PolicyVersionMissing("expected_policy_version missing: the driver published no policy token") - output = await (upstream or _load_upstream())(input) - samples = _samples_of(output) - if len(samples) == 1 and isinstance(getattr(samples[0], "metadata", None), dict) and SESSIONS_KEY in samples[0].metadata: - output = await assemble_compaction_segments(input, output, collect=collect) - limit = int(getattr(getattr(input, "args", None), "yeto_rl_max_policy_age", 0) or 0) - current = None - if limit > 0: - # 5.2: the trajectory may have been suspended across a publish; it is - # trained at the version current NOW (driver token in the metadata sink). - try: - current = rollout_meta.current_policy_token() or expected_version - except Exception: # noqa: BLE001 - unreachable sink: the start token - current = expected_version - apply_trajectory_bookkeeping(input.sample, _samples_of(output), expected_version=expected_version, - max_policy_age=limit, current_version=current) - if _compactionrl_on(): - output = _as_sample_list(output) - return output - - -def _compactionrl_on() -> bool: - from .compaction_bridge import compactionrl_enabled - - return compactionrl_enabled() - - -def _as_sample_list(output: Any) -> Any: - """CompactionRL: every rollout's output carries a *list* of samples (one per - segment, possibly one), so Miles flattens one uniform nesting level - (s19-compaction-g1-20261010e: compacted rollouts returned lists, the others a - bare Sample -> rollout_data_conversion "'Sample' object is not iterable").""" - samples = getattr(output, "samples", output) - if isinstance(samples, list): - return output - if dataclasses.is_dataclass(output) and hasattr(output, "samples"): - return dataclasses.replace(output, samples=[samples]) - return [samples] - - -def _add_arguments(parser: Any) -> None: - """Miles parse_args registers a generate function's own flags through its - ``add_arguments`` attribute; this wrapper hands the agent flags - (--custom-agent-function-path, --max-seq-len) to upstream - agentic_tool_call.generate, so it must register upstream's flags too - (s19-compaction-g1-20261010c: "unrecognized arguments").""" - upstream = _load_upstream() - add = getattr(upstream, "add_arguments", None) - if callable(add): - add(parser) - - -generate.add_arguments = _add_arguments diff --git a/yeto/rl/harness/codex/codex_openenv_subprocess_agent_function.py b/yeto/rl/harness/codex/codex_openenv_subprocess_agent_function.py deleted file mode 100644 index de052c8f6..000000000 --- a/yeto/rl/harness/codex/codex_openenv_subprocess_agent_function.py +++ /dev/null @@ -1,616 +0,0 @@ -"""Trusted-side entry for the Terminal-Bench Codex adapter (REWRITE). - -``run`` is the ``--custom-agent-function-path`` target under upstream -``agentic_tool_call.generate``. It owns the environment lease, spawns the -untrusted worker with the reward key scrubbed from its environment, maps the -worker's tool-wait events onto the ToolWaitBoard (3.1), runs the verifier and -signs the outcome (D7), and guarantees cleanup on cancellation/timeouts: -worker process group, environment lease, tool-wait scope. - -Environment provisioning is injected (``EnvironmentProvider``) so the sandbox -backend (D9) and the CPU fake share one code path. The boards are installed -by the IR-1 preflight hook (``preflight.harness_preflight``): the island's -``ToolWaitBoard`` (3.1) and ``HarnessBoard`` (IR-2) actors. IR-2 wiring here: -``allow_new_session`` before any environment is acquired (drain closes -admission; fail closed -> infrastructure ABORTED, never a 0 reward), -``enter_session``/``exit_session`` around the worker, ``lease_acquired`` -(with the hard deadline) / ``lease_released`` around the environment lease so -``env_live`` counts idle sandboxes too. IR-3: the target policy token comes -from ``rollout_meta_hook.expected_policy_version`` (prompt metadata first, -else the driver's token in the metadata sink); a missing token refuses the -trajectory before any environment is acquired. -""" - -from __future__ import annotations - -import asyncio -import json -import os -import signal -import sys -import time -from dataclasses import dataclass -from types import SimpleNamespace -from typing import Any, Awaitable, Callable, Protocol - -from yeto.rl.engine import rollout_meta -from yeto.rl.engine.tool_wait import _call as _board_call -from yeto.rl.engine.trajectory_timing import PhaseClock -from yeto.rl.engine.tool_wait import _resolve as _board_resolve - -from . import codex_openenv_agent_function as adapter -from .environment import TerminalEnvironment, TrustedVerifier - -WORKER_MODULE = "yeto.rl.harness.codex.codex_openenv_agent_worker" -WORKER_GRACE_SECONDS = 5.0 - - -@dataclass -class EnvironmentLease: - """One environment bound to one trajectory (R-D9: deadline is a hard cap).""" - - env_url: str - env_token: str - verifier: TrustedVerifier - destroy: Callable[[], Awaitable[None]] - describe: Callable[[], Awaitable[str]] # "live" | "gone" - deadline_seconds: float - # Extra environment for the worker process (e.g. a per-trajectory - # ``SECRLENV_MAX_TURNS`` turn budget injected by the provider); never a - # reward key (the worker refuses to start with one). - worker_env: dict[str, str] | None = None - - -class EnvironmentProvider(Protocol): - async def acquire(self, task_id: str, trajectory_id: str) -> EnvironmentLease: ... - - -_provider: EnvironmentProvider | None = None -_tool_wait_board: Any = None -_harness_board: Any = None -_member: str | None = None - -ADMISSION_CLOSED = "harness_admission_closed" -# Consecutive non-injected acquire failures before the island is torn down. -# Each failure is reported as INFRASTRUCTURE (sample ABORTED); upstream -# Miles then draws the next group forever (``generate_rollout`` has no -# abort cap: codex-smoke-20261003-7 looped 17k samples in six minutes on a -# broken provider), so after this many the worker raises -# ``EnvironmentProviderOutage`` instead. -PROVIDER_OUTAGE_THRESHOLD_ENV = "YETO_HARNESS_PROVIDER_OUTAGE_THRESHOLD" -DEFAULT_PROVIDER_OUTAGE_THRESHOLD = 4 -_acquire_failures = 0 - - -class EnvironmentProviderOutage(BaseException): - """The environment provider failed repeatedly: stop the rollout, do not resample. - - A ``BaseException`` on purpose: upstream Miles wraps the agent function and - each rollout task in ``except Exception`` and keeps generating, so only a - BaseException escapes ``generate_rollout`` and fails the island. - """ - - -def _provider_outage_threshold() -> int: - try: - return max(1, int(os.environ.get(PROVIDER_OUTAGE_THRESHOLD_ENV) or DEFAULT_PROVIDER_OUTAGE_THRESHOLD)) - except ValueError: - return DEFAULT_PROVIDER_OUTAGE_THRESHOLD - - -def _note_acquire_failure(exc: BaseException, trajectory_id: str) -> None: - """Count real provisioning failures; injected ones (fault tests) do not count.""" - global _acquire_failures - if getattr(exc, "injected_fault", False): - return - _acquire_failures += 1 - threshold = _provider_outage_threshold() - if _acquire_failures >= threshold: - raise EnvironmentProviderOutage( - f"environment provider failed {_acquire_failures} consecutive acquires " - f"(threshold {threshold}); last for {trajectory_id!r}: {type(exc).__name__}: {exc}" - ) from exc - - -def _note_acquire_success() -> None: - global _acquire_failures - _acquire_failures = 0 -POLICY_VERSION_MISSING = "expected_policy_version missing" - - -class PolicyVersionMissing(RuntimeError): - """IR-3: the driver published no target policy token for this rollout (refuse).""" - - -def configure( - *, - provider: EnvironmentProvider | None, - tool_wait_board: Any = None, - harness_board: Any = None, - member: str | None = None, -) -> None: - """Install the environment provider and the island boards. - - ``tool_wait_board``: ``ToolWaitBoard`` (local or actor handle, 3.1). - ``harness_board``: ``HarnessBoard`` (local or actor handle, IR-2); None - means the island reports "harness counts unknown" and stays undrainable. - ``member``: the rollout member this process targets (admission key). - """ - global _provider, _tool_wait_board, _harness_board, _member, _acquire_failures - _acquire_failures = 0 - _provider = provider - _tool_wait_board = tool_wait_board - _harness_board = harness_board - _member = member - - -def configured() -> dict[str, Any]: - return {"provider": _provider, "tool_wait_board": _tool_wait_board, "harness_board": _harness_board, - "member": _member} - - -async def _board_await(target: Any, method: str, *args: Any) -> None: - """Call a board method and wait for it, so consecutive calls stay ordered. - - The island boards are Ray actors with ``max_concurrency=64``: back-to-back - fire-and-forget ``exit``/``enter`` of one trajectory's successive tool - calls ran out of order ("already in a tool call" / "is not in a tool call", - A-T3-7, codex-smoke-20261003-11). Board bookkeeping is telemetry, so its - errors are logged and never fail the trajectory. - """ - try: - value = _board_call(target, method, *args) - if type(value).__name__ == "ObjectRef": - await asyncio.to_thread(_board_resolve, value) - except Exception as exc: # noqa: BLE001 - never fail a trajectory on telemetry - print(f"[codex-harness] tool-wait board {method} failed: {type(exc).__name__}: {exc}", file=sys.stderr) - - -def _board_kwcall(target: Any, method: str, *args: Any, **kwargs: Any) -> Any: - """Like ``tool_wait._call`` but forwards keyword arguments (``lease_acquired(deadline=)``).""" - fn = getattr(target, method) - remote = getattr(fn, "remote", None) - return _board_resolve(remote(*args, **kwargs) if callable(remote) else fn(*args, **kwargs)) - - -def resolve_expected_policy_version(metadata: dict[str, Any]) -> str | None: - """IR-3 target token: prompt metadata first, else the driver token in the sink.""" - try: - return rollout_meta.expected_policy_version(SimpleNamespace(metadata=metadata)) - except Exception: # noqa: BLE001 - no sink reachable == no token published - return None - - -async def _terminate(process: asyncio.subprocess.Process) -> None: - if process.returncode is not None: - return - for sig, wait in ((signal.SIGTERM, WORKER_GRACE_SECONDS), (signal.SIGKILL, WORKER_GRACE_SECONDS)): - try: - os.killpg(process.pid, sig) - except ProcessLookupError: - return - # asyncio.wait (not wait_for): must work inside a cancelled task's finally. - waiter = asyncio.ensure_future(process.wait()) - done, _pending = await asyncio.wait({waiter}, timeout=wait) - if done: - return - waiter.cancel() - - -async def _drive_worker( - job: dict[str, Any], trajectory_id: str, board: Any, worker_env: dict[str, str] | None = None -) -> dict[str, Any]: - """Spawn the worker, relay tool-wait events, return its result event.""" - process = await asyncio.create_subprocess_exec( - sys.executable, - "-m", - WORKER_MODULE, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.DEVNULL if os.getenv("YETO_CODEX_WORKER_STDERR") != "1" else None, - env={**os.environ, **{str(k): str(v) for k, v in (worker_env or {}).items()}}, - start_new_session=True, - ) - in_tool = False - phases = PhaseClock() # agentic-rollout-utilization 1.3 - try: - assert process.stdin is not None and process.stdout is not None - process.stdin.write((json.dumps(job) + "\n").encode()) - await process.stdin.drain() - process.stdin.close() - while True: - line = await process.stdout.readline() - if not line: - raise adapter.harness.CodexHarnessError("worker exited without a result") - event = json.loads(line) - kind = event.get("event") - if kind == "tool_wait": - if event.get("phase") == "enter" and not in_tool: - in_tool = True - phases.enter_tool() - if board is not None: - await _board_await(board, "enter", trajectory_id) - elif event.get("phase") == "exit" and in_tool: - in_tool = False - phases.exit_tool() - if board is not None: - await _board_await(board, "exit", trajectory_id) - elif kind == "suspend": # agentic-rollout-utilization 5.1 - if event.get("phase") == "enter": - phases.enter_suspend() - elif event.get("phase") == "exit": - phases.exit_suspend() - elif kind == "result": - if isinstance(event.get("metrics"), dict): # observe only, unsigned - event["metrics"].update(phases.finish()) - return event - elif kind == "error": - error = adapter.harness.CodexHarnessError(str(event.get("reason"))) - if isinstance(event.get("metrics"), dict): - error.metrics = event["metrics"] # type: ignore[attr-defined] - raise error - finally: - if in_tool and board is not None: - await _board_await(board, "exit", trajectory_id) - await _terminate(process) - - -class TaskPromptMissing(ValueError): - """No task statement could be resolved for the episode (fail closed).""" - - -def _last_user_text(prompt: Any) -> str | None: - if not isinstance(prompt, list): - return None - for message in reversed(prompt): - if isinstance(message, dict) and message.get("role") == "user": - content = message.get("content") - if isinstance(content, str) and content.strip(): - return content - return None - - -def task_prompt(metadata: dict[str, Any], prompt: Any, lease: Any) -> str: - """The task statement Codex receives as its first user message. - - Precedence: explicit ``metadata["prompt"]`` string; the provider task's - ``instruction.md``; a plain-string Miles prompt; the last user message of a - chat-format prompt. A chat prompt without a user message (e.g. only a - generic system message) is refused instead of being stringified -- S15 - stage 2 sent ``str([{'role': 'system', ...}])`` to every trajectory, so the - model never saw the task and all 48 rewards were 0. - """ - explicit = metadata.get("prompt") - if isinstance(explicit, str) and explicit.strip(): - return explicit - task = getattr(lease, "task", None) - instruction = getattr(task, "instruction", None) - if isinstance(instruction, str) and instruction.strip(): - return instruction - if isinstance(prompt, str) and prompt.strip(): - return prompt - user = _last_user_text(prompt) - if user is not None: - return user - raise TaskPromptMissing( - f"no task statement for task {metadata.get('task_id')!r}: metadata has no prompt, " - "the environment task has no instruction.md and the Miles prompt has no user message" - ) - - -# S17 N17 (over-sampling A/B): in-flight trajectories of this rollout process. -# Miles calls ``.abort(args)`` (``call_agent_abort_hook``) once -# enough groups are collected; without it the Codex loops keep issuing fresh -# turns after SGLang's abort_all and the rollout waits for every surplus group. -# Each entry: ``stage`` (acquire / segments / worker / verify / teardown), the -# abortable step's task (None in acquire and teardown) and the start time. -# acquire is never cancelled (a sandbox created by a cancelled acquire would -# leak): the trajectory is flagged and stops right after acquire returns; the -# other steps are cancelled and the lease is destroyed in ``run``'s finally. -_INFLIGHT: dict[str, dict[str, Any]] = {} -_ABORTED_IDS: set[str] = set() -ABORTED_STATUS = "aborted_surplus" -LEASE_EXPIRED_ON_BOARD = "lease force-released on the harness board before release" - - -class _AbortedSurplus(Exception): - """abort() stopped this trajectory (surplus of an over-sampled rollout).""" - - -class _LeaseExpiredOnBoard(Exception): - """The board force-released the lease at its hard deadline before ``run`` released it.""" - - -def _stage(trajectory_id: str, stage: str) -> None: - entry = _INFLIGHT.get(trajectory_id) - if entry is not None: - entry["stage"] = stage - - -async def _abortable(trajectory_id: str, stage: str, awaitable: Any) -> Any: - """Run one step of ``run`` as a task abort() may cancel; a cancel that abort() - issued surfaces as :class:`_AbortedSurplus` (any other cancel propagates).""" - if trajectory_id in _ABORTED_IDS: - if asyncio.iscoroutine(awaitable): - awaitable.close() - raise _AbortedSurplus(stage) - task = asyncio.ensure_future(awaitable) - entry = _INFLIGHT.get(trajectory_id) - if entry is not None: - entry.update(stage=stage, task=task) - try: - return await task - except asyncio.CancelledError: - if trajectory_id in _ABORTED_IDS and task.cancelled(): - raise _AbortedSurplus(stage) from None - raise - finally: - if entry is not None: - entry["task"] = None - - -async def abort(args: Any = None) -> int: - """Stop every in-flight trajectory of this process: cancel its current step - (segments / worker / verify) or, while it is still acquiring its sandbox, - flag it so it stops as soon as the sandbox exists. Leases are destroyed in - ``run``'s finally. Returns how many trajectories were stopped and logs one - ``rl_codex_abort`` line (count per stage, elapsed seconds).""" - now = time.monotonic() - stages: dict[str, int] = {} - elapsed = [] - for tid, entry in list(_INFLIGHT.items()): - if entry.get("stage") == "teardown": - continue - _ABORTED_IDS.add(tid) - task = entry.get("task") - if task is not None and not task.done(): - task.cancel() - stage = entry.get("stage", "?") - stages[stage] = stages.get(stage, 0) + 1 - elapsed.append(round(now - entry["t0"], 2)) - if elapsed: - print("YETO_CODEX_ABORT " + json.dumps({ - "event": "rl_codex_abort", "cancelled": len(elapsed), "stages": stages, - "elapsed_s": sorted(elapsed)}), flush=True) - return len(elapsed) - - -# agentic-rollout-utilization 5.1/5.3: suspension between model turns -# (Miles --agentic-suspend-between-turns calls ``resume`` at the start of every -# rollout and ``suspend`` at its cut-off). Each trajectory's worker gets its own -# gate file (``YETO_CODEX_SUSPEND_GATE``); ``suspend`` creates the files of the -# trajectories in flight (closed gate: the worker's bridge holds the next model -# request -- a running tool call is not interrupted), ``resume`` removes them. -# A trajectory parked longer than ``YETO_CODEX_SUSPEND_MAX_SECONDS`` (600 s) is -# discarded by its bridge (CodexSuspendExpired -> infrastructure) and its lease -# is destroyed in ``run``'s finally. -_SUSPEND_DIR: str | None = None -_SUSPEND_STATE: dict[str, Any] = {"suspended_at": None, "gates": {}} -SUSPEND_EXPIRED = "CodexSuspendExpired" - - -def _gate_path(trajectory_id: str) -> str | None: - if _SUSPEND_DIR is None: - return None - safe = "".join(c if c.isalnum() or c in "-_." else "_" for c in trajectory_id) - return os.path.join(_SUSPEND_DIR, f"{safe}.gate") - - -async def suspend(args: Any = None) -> int: - """Close the model-turn gate of every trajectory in flight; returns how many.""" - global _SUSPEND_DIR - if _SUSPEND_DIR is None: - import tempfile - - _SUSPEND_DIR = tempfile.mkdtemp(prefix="yeto-codex-suspend-") - gates = _SUSPEND_STATE["gates"] - stages: dict[str, int] = {} - for tid, entry in list(_INFLIGHT.items()): - if entry.get("stage") == "teardown": - continue - path = _gate_path(tid) - with open(path, "w", encoding="utf-8") as fh: - fh.write("suspended\n") - gates[tid] = path - stage = entry.get("stage", "?") - stages[stage] = stages.get(stage, 0) + 1 - _SUSPEND_STATE["suspended_at"] = time.time() - print("YETO_CODEX_SUSPEND " + json.dumps({ - "event": "rl_codex_suspend", "suspended": len(gates), "stages": stages}), flush=True) - return len(gates) - - -async def resume(args: Any = None) -> int: - """Open every closed gate (start of a rollout); returns how many were closed. - The first call also turns the gates on for the trajectories started after it.""" - global _SUSPEND_DIR - if _SUSPEND_DIR is None: - import tempfile - - _SUSPEND_DIR = tempfile.mkdtemp(prefix="yeto-codex-suspend-") - gates = _SUSPEND_STATE["gates"] - opened = 0 - for tid, path in list(gates.items()): - try: - os.unlink(path) - opened += 1 - except FileNotFoundError: - pass - gates.pop(tid, None) - since = _SUSPEND_STATE.get("suspended_at") - _SUSPEND_STATE["suspended_at"] = None - if opened or since is not None: - print("YETO_CODEX_RESUME " + json.dumps({ - "event": "rl_codex_resume", "resumed": opened, - "suspended_s": None if since is None else round(time.time() - since, 3)}), flush=True) - return opened - - -def _suspend_worker_env(trajectory_id: str) -> dict[str, str]: - path = _gate_path(trajectory_id) - if path is None: - return {} # suspension not in use: worker environment unchanged - if _SUSPEND_STATE.get("suspended_at") is not None and trajectory_id not in _SUSPEND_STATE["gates"]: - # started while the rollout is suspended (should not happen: Miles holds - # new samples) -- park it too rather than generate on offloaded engines - with open(path, "w", encoding="utf-8") as fh: - fh.write("suspended\n") - _SUSPEND_STATE["gates"][trajectory_id] = path - return {adapter.harness.SUSPEND_GATE_ENV: path} - - -async def run( - base_url: str, - prompt: Any, - request_kwargs: dict[str, Any] | None = None, - metadata: dict[str, Any] | None = None, - **kwargs: Any, -) -> dict[str, Any] | None: - metadata = dict(metadata or {}) - try: - return await _run(base_url, prompt, request_kwargs, metadata, **kwargs) - except _LeaseExpiredOnBoard as exc: - # The trajectory outlived its lease's hard deadline (e.g. a detached - # Modal sandbox) and the board already force-released the lease: - # infrastructure, never a reward, and never a failed rollout (S17 N17 A/B a). - _task_id, sample_id = adapter.task_identity(metadata) - trajectory_id = str(metadata.get("trajectory_id") or sample_id) - return {**adapter.infrastructure_metadata(f"{LEASE_EXPIRED_ON_BOARD}: {exc}", episode_id=None), - **adapter.trajectory_fields(trajectory_id)} - - -async def _run( - base_url: str, - prompt: Any, - request_kwargs: dict[str, Any] | None, - metadata: dict[str, Any], - **_kwargs: Any, -) -> dict[str, Any] | None: - if _provider is None: - # Rollout workers are separate Ray actors: the driver's configure() - # never ran here, so build the same wiring from the forwarded env. - from .preflight import configure_rollout_worker - - configure_rollout_worker() - if _provider is None: - raise RuntimeError("codex_openenv_subprocess_agent_function.configure(provider=...) was not called") - task_id, sample_id = adapter.task_identity(metadata) - trajectory_id = str(metadata.get("trajectory_id") or sample_id) - episode_id = adapter.new_episode_id() - fields = adapter.trajectory_fields(trajectory_id) - started_at = time.time() # agentic-rollout-utilization 1.3 (observe only) - expected_version = resolve_expected_policy_version(metadata) - if expected_version is None: # IR-3: refuse before any environment exists - raise PolicyVersionMissing(f"{POLICY_VERSION_MISSING} for trajectory {trajectory_id!r}") - board = _harness_board - if board is not None and not _board_kwcall(board, "allow_new_session", _member): - # IR-2: admission closed (drain of this member) -> infrastructure, not a reward. - return {**adapter.infrastructure_metadata(ADMISSION_CLOSED, episode_id=None), **fields} - lease: EnvironmentLease | None = None - session_open = lease_open = handed_off = False - _INFLIGHT[trajectory_id] = {"stage": "acquire", "task": None, "t0": time.monotonic()} - try: - if board is not None: - _board_kwcall(board, "enter_session", trajectory_id, _member) - session_open = True - try: - acquire_started = time.monotonic() - lease = await _provider.acquire(task_id, trajectory_id) - sandbox_start_seconds = round(time.monotonic() - acquire_started, 3) - except Exception as exc: # noqa: BLE001 - provisioning failures are infrastructure - _note_acquire_failure(exc, trajectory_id) # raises EnvironmentProviderOutage past the threshold - return {**adapter.infrastructure_metadata(f"acquire: {type(exc).__name__}: {exc}", episode_id=None), **fields} - _note_acquire_success() - if trajectory_id in _ABORTED_IDS: - # abort() during acquire: never started, the sandbox is destroyed below - return {**adapter.infrastructure_metadata(ABORTED_STATUS, episode_id=None), **fields} - if board is not None: - _board_kwcall(board, "lease_acquired", trajectory_id, deadline=time.monotonic() + lease.deadline_seconds) - lease_open = True - job = { - "base_url": base_url, - "prompt": task_prompt(metadata, prompt, lease), - "request_kwargs": dict(request_kwargs or {}), - "episode_id": episode_id, - "max_seq_len": metadata.get("max_seq_len"), - "env_url": lease.env_url, - "env_token": lease.env_token, - "driver": metadata.get("codex_openenv_driver", "stock"), - **{k: metadata[k] for k in ("script", "final_status", "hang_seconds") if k in metadata}, - } - try: - segments = await _abortable(trajectory_id, "segments", adapter.prepare_segment_sessions(job)) - except _AbortedSurplus: - return {**adapter.infrastructure_metadata(ABORTED_STATUS, episode_id=episode_id), **fields} - except Exception as exc: # noqa: BLE001 - session-server failures are infrastructure - return {**adapter.infrastructure_metadata(f"segment sessions: {type(exc).__name__}: {exc}", episode_id=episode_id), **fields} - try: - try: - untrusted = await _abortable(trajectory_id, "worker", asyncio.wait_for( - _drive_worker(job, trajectory_id, _tool_wait_board, - {**(getattr(lease, "worker_env", None) or {}), - **_suspend_worker_env(trajectory_id)}), - timeout=lease.deadline_seconds)) - except _AbortedSurplus: - # abort(): surplus trajectory of an over-sampled rollout; Miles - # discards it (no partial rollout). Infrastructure, never a reward. - return {**adapter.infrastructure_metadata(ABORTED_STATUS, episode_id=episode_id), **fields} - except asyncio.TimeoutError: - untrusted = {"status": "timeout", "metrics": {"timed_out": 1}, "episode_id": episode_id} - except adapter.harness.CodexHarnessError as exc: - metrics = getattr(exc, "metrics", None) - tito = adapter.mirror_tito_counters(metrics, board) - handed_off = True - return {**adapter.infrastructure_metadata(str(exc), episode_id=episode_id, metrics=metrics), **fields, **tito, **segments} - except Exception as exc: # noqa: BLE001 - e.g. the sandbox shut down under the worker - # Environment failure, not a policy outcome: infrastructure (no reward, group dropped) - # instead of an exception that fails the whole rollout step. The finally block - # deletes the pre-created sessions (not handed off). - return {**adapter.infrastructure_metadata(f"worker: {type(exc).__name__}: {exc}", episode_id=episode_id), **fields} - tito = adapter.mirror_tito_counters(untrusted.get("metrics"), board) - try: - signed = await _abortable(trajectory_id, "verify", adapter.finish_trusted( - untrusted, lease.verifier, task_id=task_id, sample_id=sample_id)) - except _AbortedSurplus: - return {**adapter.infrastructure_metadata(ABORTED_STATUS, episode_id=episode_id), **fields} - except Exception as exc: # noqa: BLE001 - verifier sandbox gone (Modal ConflictError / NotFoundError) - return {**adapter.infrastructure_metadata(f"verify: {type(exc).__name__}: {exc}", episode_id=episode_id), **fields} - signed["expected_policy_version"] = expected_version - signed.update(trajectory_started_at=started_at, trajectory_ended_at=time.time(), - sandbox_start_seconds=sandbox_start_seconds) - handed_off = True - return {**signed, **fields, **tito, **segments} - finally: - # Worker crash (any non-harness exception), cancellation or a - # finish_trusted error: the metadata never reaches - # codex_openenv_generate, so delete the pre-created sessions here. - if segments and not handed_off: - await adapter.release_unreturned_segments(base_url, segments) - finally: - _stage(trajectory_id, "teardown") - try: - if lease is not None: - await lease.destroy() - if await lease.describe() != "gone": - # Not confirmed gone: keep the lease on the board; it is - # force-released at its hard deadline (leases_expired_total). - raise RuntimeError(f"environment for {trajectory_id} was not confirmed destroyed") - if lease_open: - try: - _board_kwcall(board, "lease_released", trajectory_id) - except Exception as exc: # noqa: BLE001 - ray wraps the board's ToolWaitError - if "is not live" not in str(exc): - raise - # the board force-released it at its hard deadline already - raise _LeaseExpiredOnBoard(str(exc).splitlines()[-1][:300]) from exc - finally: - _INFLIGHT.pop(trajectory_id, None) - _ABORTED_IDS.discard(trajectory_id) - gate = _SUSPEND_STATE["gates"].pop(trajectory_id, None) - if gate is not None: - try: - os.unlink(gate) - except FileNotFoundError: - pass - if session_open: - _board_kwcall(board, "exit_session", trajectory_id) diff --git a/yeto/rl/harness/codex/compaction_bridge.py b/yeto/rl/harness/codex/compaction_bridge.py deleted file mode 100644 index 90a02a18a..000000000 --- a/yeto/rl/harness/codex/compaction_bridge.py +++ /dev/null @@ -1,379 +0,0 @@ -"""CompactionRL context compaction inside the Codex Responses bridge (design D8). - -Opt-in only (``YETO_CODEX_COMPACTIONRL=1``); with it unset nothing here runs and -the stock path (``codex_harness_agent._drive_codex``) is used unchanged. - -Why the bridge can rewrite the context although Codex keeps the full history: -the bridge never forwards Codex's Responses ``input``. It checks that input is -exactly the expected append-only history (``_expected_input``) and samples Miles -from its own Chat Completions message list (``_messages``). Codex requests use -``store:false`` and no ``previous_response_id`` (both enforced by -``_validate_codex_request``), so Codex is never told about the compaction: its -own history stays complete and append-only, while each Miles request carries -only the rebuilt context ``s + u_resume(S) + last k steps``. - -Training shape (progress.md "S13 cross_segment 每段 sample"): one sample per -segment. The fork-pinned v1 session server is append-only and would silently -roll back a rewritten history (R-D5a), so every segment uses its own session: -the summary is sampled in the current session (trainable, logprobs recorded by -the session server like any other turn), then the rebuilt context starts in the -next pre-created session. The trusted parent pre-creates those sessions -(``codex_openenv_agent_function.create_segment_sessions``) and -``codex_openenv_generate`` collects them and writes segment metadata. - -Codex 0.145.0's own auto-compaction is switched off explicitly -(``model_auto_compact_token_limit`` = i64 max); any Codex-side compaction still -fails closed in the stock bridge (history item types / ``thread/compacted``). -""" - -from __future__ import annotations - -import asyncio -import copy -import json -import os -import tempfile -from collections.abc import Mapping, Sequence -from pathlib import Path -from typing import Any, Callable, Self - -from yeto.rl.compaction import ( - MAX_COMPACTIONS, - RESUME_TEMPLATE, - SUMMARY_PROMPT, - T_COMP, - CompactionConfig, - extract_summary, - should_compact, -) - -from . import codex_harness_agent as harness -from . import agent as legacy - -COMPACTIONRL_ENV = "YETO_CODEX_COMPACTIONRL" -T_COMP_ENV = "YETO_CODEX_COMPACTIONRL_T_COMP" -# Explicit compaction budget registration: when set, the named profile fully -# determines the CompactionConfig (its context_budget must equal max_seq_len, -# which the bridge enforces anyway); unset keeps the max_seq_len+T_COMP -# derivation below. Starts with YETO_CODEX_ so workers inherit it. -PROFILE_ENV = "YETO_CODEX_COMPACTIONRL_PROFILE" -SEGMENT_URLS_KEY = "segment_base_urls" -METRICS_KEY = "codex_compaction" -SESSIONS_METADATA_KEY = "codex_compaction_sessions" -# Router URL the segment sessions were created on (trusted, next to the ids); -# the generate wrapper collects them there. -SESSIONS_ROUTER_METADATA_KEY = "codex_compaction_router" -# Codex 0.145.0 accepts this key under --strict-config (checked against the -# pinned binary: a string value fails with "expected i64"). -CODEX_AUTO_COMPACT_DISABLED = "model_auto_compact_token_limit=9223372036854775807" -# Summary request tool table (design D8 "摘要请求工具表"): explicit choice to -# send NO tool table -- exactly the stock ``codex_harness_agent._sample_miles -# (summary=True)`` payload (tools=[], tool_choice="none") -- so the policy -# cannot call a tool while summarising. Risk, unverified until GPU: the -# execution turns carry the Codex tool table, so the summary prompt renders a -# different system/tools prefix and TITO prefix reuse on the session server may -# break or re-tokenise. Changing this needs a stock-sampler change; tests pin -# that the wire payload matches these constants. -SUMMARY_REQUEST_TOOLS: tuple[Any, ...] = () -SUMMARY_REQUEST_TOOL_CHOICE = "none" -SUMMARY_PROMPT_TOKEN_RESERVE = len(SUMMARY_PROMPT.encode("utf-8")) + 256 -_MESSAGE_TOKEN_OVERHEAD = 256 -_FALSE = frozenset({"", "0", "false", "no", "off"}) -_TRUE = frozenset({"1", "true", "yes", "on"}) - - -def compactionrl_enabled(env: Mapping[str, str] | None = None) -> bool: - raw = (os.environ if env is None else env).get(COMPACTIONRL_ENV, "").strip().lower() - if raw in _FALSE: - return False - if raw in _TRUE: - return True - raise harness.CodexHarnessError(f"{COMPACTIONRL_ENV} must be a boolean flag") - - -_COMPACTION_PROFILES: dict[str, Callable[[], CompactionConfig]] = {} - - -def register_compaction_profile(name: str, factory: Callable[[], CompactionConfig]) -> None: - """Register a named, explicit compaction budget selectable via PROFILE_ENV.""" - if not name or not callable(factory): - raise harness.CodexHarnessError("compaction profiles need a name and a factory") - _COMPACTION_PROFILES[name] = factory - - -def compaction_profiles() -> tuple[str, ...]: - return tuple(sorted(_COMPACTION_PROFILES)) - - -def compaction_config(max_seq_len: int, env: Mapping[str, str] | None = None) -> CompactionConfig: - env = os.environ if env is None else env - name = env.get(PROFILE_ENV, "").strip() - if name: - factory = _COMPACTION_PROFILES.get(name) - if factory is None: - raise harness.CodexHarnessError( - f"{PROFILE_ENV}={name!r} is not a registered compaction profile " - f"(registered: {', '.join(compaction_profiles()) or 'none'})" - ) - cfg = factory() - if cfg.context_budget != int(max_seq_len): - raise harness.CodexHarnessError( - f"compaction profile {name!r} fixes context_budget={cfg.context_budget}; " - f"max_seq_len is {max_seq_len} (the bridge requires them equal)" - ) - return cfg - raw = env.get(T_COMP_ENV, "").strip() - t_comp = T_COMP - if raw: - if not raw.isdigit() or int(raw) <= 0: - raise harness.CodexHarnessError(f"{T_COMP_ENV} must be a positive integer") - t_comp = int(raw) - try: - return CompactionConfig(context_budget=int(max_seq_len), t_comp=t_comp) - except ValueError as exc: - raise harness.CodexHarnessError(f"CompactionRL config rejected: {exc}") from exc - - -def qwen3_cyber_compaction_config() -> CompactionConfig: - """Registered fallback for long cyber traces after a 32K baseline truncates. - - Keep 4K tokens available for the next response and permit at most three - segment boundaries; segment identity and shared reward are preserved by - ``assemble_compaction_segments``. Select it on the run path with - ``YETO_CODEX_COMPACTIONRL_PROFILE=qwen3_cyber`` together with - ``--max-seq-len 24576`` (spec 5.7: the Qwen budget is explicit, never a - default of an unregistered model). - """ - return CompactionConfig(context_budget=24_576, t_comp=4_096, max_compactions=3) - - -register_compaction_profile("qwen3_cyber", qwen3_cyber_compaction_config) - - -def _message_bound(message: Mapping[str, Any]) -> int: - """Upper bound on a chat message's tokens: bytes + template overhead.""" - return len(json.dumps(message, ensure_ascii=False).encode("utf-8")) + _MESSAGE_TOKEN_OVERHEAD - - -def compaction_counters(metrics: Any) -> dict[str, Any]: - record = getattr(metrics, METRICS_KEY, None) - return {METRICS_KEY: copy.deepcopy(record)} if isinstance(record, dict) else {} - - -class CompactionRLBridge(harness._ResponsesBridge): - """Stock bridge + D8 trigger/summary/rebuild, one session per segment.""" - - def __init__( - self, - miles_base_url: str, - prompt: str, - request_kwargs: dict[str, Any], - metrics: legacy.AgentMetrics, - *, - max_seq_len: int | None, - segment_base_urls: Sequence[str], - config: CompactionConfig | None = None, - ) -> None: - super().__init__( - miles_base_url, prompt, request_kwargs, metrics, max_seq_len=max_seq_len - ) - if self._compaction_enabled: - raise harness.CodexHarnessError( - "legacy YETO_CODEX_COMPACTION_ENABLED and CompactionRL are exclusive" - ) - if max_seq_len is None: - raise harness.CodexHarnessError("CompactionRL requires max_seq_len (C)") - self._cfg = config or compaction_config(max_seq_len) - if self._cfg.context_budget != max_seq_len: - raise harness.CodexHarnessError("CompactionRL budget must equal max_seq_len") - if len(segment_base_urls) != self._cfg.max_compactions: - raise harness.CodexHarnessError( - "CompactionRL needs one pre-created session per allowed compaction" - ) - self._segment_urls = [ - f"{legacy._session_url(url).rstrip('/')}/chat/completions" - for url in segment_base_urls - ] - # Paper sec. 5.1: one reply cap for every turn, summaries included. - self._compaction_summary_max_tokens = self._requested_max_tokens - self._max_compactions = self._cfg.max_compactions - self._step_bounds: list[int] = [] - self._record: dict[str, Any] = { - "schema": 1, - "context_budget": self._cfg.context_budget, - "t_comp": self._cfg.t_comp, - "max_compactions": self._cfg.max_compactions, - "compactions": 0, - "segments": [], - } - setattr(metrics, METRICS_KEY, self._record) - - # --- stock hooks ----------------------------------------------------- - def _compaction_headers(self) -> dict[str, str]: - return {} # the pinned session server has no compaction headers (R-D5a) - - def _should_compact(self, estimated_context_tokens: int) -> bool: - return ( - self._compaction_count < self._cfg.max_compactions - and bool(self._atomic_steps) - and self._pending is None - and should_compact(estimated_context_tokens, self._cfg) - ) - - def expect_tool_output(self, call_id: str, output: str) -> None: - super().expect_tool_output(call_id, output) - assistant, tool = self._atomic_steps[-1] - self._step_bounds.append(_message_bound(assistant) + _message_bound(tool)) - - def _rebuild_compacted_messages(self, retained_step_count: int) -> None: - if self._resume_summary is None: - raise harness.CodexHarnessError("compaction resume summary is missing") - retained = self._atomic_steps[-retained_step_count:] if retained_step_count else [] - bounds = self._step_bounds[-retained_step_count:] if retained_step_count else [] - messages: list[dict[str, Any]] = [ - {"role": "system", "content": harness.base_instructions()}, - {"role": "user", "content": RESUME_TEMPLATE.format(summary=self._resume_summary)}, - ] - for assistant, tool in retained: - messages.extend((copy.deepcopy(assistant), copy.deepcopy(tool))) - self._messages = messages - self._atomic_steps = copy.deepcopy(retained) - self._step_bounds = list(bounds) - self._retained_step_count = retained_step_count - if self._record["segments"]: - self._record["segments"][-1]["kept_steps"] = retained_step_count - # Nothing of the new window is tokenized yet: bound it from above. - self._current_context_tokens = 0 - self._unobserved_tool_token_upper_bound = sum(_message_bound(m) for m in messages) - - def _choose_kept_steps(self, summary: str) -> int: - """Eq. 9: largest k <= cfg.k whose rebuilt context does not trigger again.""" - head = _message_bound({"role": "system", "content": harness.base_instructions()}) + _message_bound( - {"role": "user", "content": RESUME_TEMPLATE.format(summary=summary)} - ) - for k in range(min(self._cfg.keep_recent_steps, len(self._atomic_steps)), 0, -1): - if not should_compact(head + sum(self._step_bounds[-k:]), self._cfg): - return k - return 0 - - async def _compact_context(self) -> None: - if self._pending is not None or self._last_assistant_message is not None: - raise harness.CodexHarnessError("compaction attempted inside an atomic tool step") - used = self._current_context_tokens + self._unobserved_tool_token_upper_bound - if self._cfg.context_budget - used - SUMMARY_PROMPT_TOKEN_RESERVE <= 0: - raise harness.CodexSequenceLimit("no room left for the compaction summary") - # The stock sampler reserves the legacy instruction size; top it up to ours. - self._unobserved_tool_token_upper_bound += max( - 0, SUMMARY_PROMPT_TOKEN_RESERVE - harness.COMPACTION_SUMMARY_INSTRUCTION_TOKEN_RESERVE - ) - summary_messages = copy.deepcopy(self._messages) - summary_messages.append({"role": "user", "content": SUMMARY_PROMPT}) - try: - # summary=True -> SUMMARY_REQUEST_TOOLS / SUMMARY_REQUEST_TOOL_CHOICE (see above). - completion = await self._sample_miles(messages=summary_messages, summary=True) - except harness._CompactionContextDoesNotFit as exc: - raise harness.CodexSequenceLimit("compaction summary does not fit") from exc - _usage, total = harness._usage(completion) - if total is None: - self._metrics.usage_missing += 1 - raise harness.CodexHarnessError("Miles omitted summary token usage") - self._metrics.max_model_total_tokens = max(self._metrics.max_model_total_tokens, total) - text, finish_reason = self._summary_text(completion) - summary, ok = extract_summary(text) - if not summary: - raise harness.CodexModelFailure("compaction summary is empty") - segment = self._compaction_count - self._record["segments"].append( - { - "segment_index": segment, - "trigger_context_tokens": used, - "summary_total_tokens": total, - "summary_finish_reason": finish_reason, - "summary_ok": ok, - "kept_steps": None, - } - ) - self._miles_url = self._segment_urls[segment] - self._compaction_count += 1 - self._record["compactions"] = self._compaction_count - self._context_window += 1 - self._resume_summary = summary - self._rebuild_compacted_messages(self._choose_kept_steps(summary)) - - @staticmethod - def _summary_text(completion: dict[str, Any]) -> tuple[str, str]: - choices = completion.get("choices") - if not isinstance(choices, list) or len(choices) != 1 or not isinstance(choices[0], dict): - raise harness.CodexHarnessError("Miles returned extra or missing summary samples") - choice = choices[0] - finish_reason = choice.get("finish_reason") - if finish_reason not in {"stop", "length"}: - raise harness.CodexModelFailure("compaction summary ended abnormally") - message = choice.get("message") - if not isinstance(message, dict) or message.get("role") != "assistant": - raise harness.CodexModelFailure("compaction summary is not an assistant message") - content = message.get("content") - if not isinstance(content, str) or message.get("tool_calls") not in (None, []): - raise harness.CodexModelFailure("compaction summary must be tool-free text") - return content, finish_reason - - -class _CompactionAppServerDriver(harness._AppServerDriver): - """Stock driver; only the Codex argv gains the auto-compaction off switch.""" - - async def __aenter__(self) -> Self: - isolated_home = tempfile.TemporaryDirectory(prefix="yeto-codex-") - self._isolated_home = isolated_home - Path(isolated_home.name).chmod(0o700) - environment = { - "CODEX_HOME": isolated_home.name, - "PATH": "/usr/bin:/bin", - "LANG": "C.UTF-8", - "TZ": "UTC", - "YETO_CODEX_BRIDGE_TOKEN": self._bridge.token, - } - self._process = await asyncio.create_subprocess_exec( - *codex_argv(self._binary, self._bridge), - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - env=environment, - limit=harness.MAX_APP_SERVER_FRAME_BYTES + 1, - start_new_session=True, - ) - self._process_group_id = self._process.pid - self._stderr_task = asyncio.create_task(self._drain_stderr()) - return self - - -def codex_argv(binary: Path, bridge: harness._ResponsesBridge) -> list[str]: - return [*harness._codex_argv(binary, bridge), "-c", CODEX_AUTO_COMPACT_DISABLED] - - -async def drive_codex_compactionrl( - binary: Path, - miles_base_url: str, - client: Any, - episode: dict[str, Any], - request_kwargs: dict[str, Any], - metrics: legacy.AgentMetrics, - *, - max_seq_len: int | None, - segment_base_urls: Sequence[str], -) -> str: - """``codex_harness_agent._drive_codex`` with the CompactionRL bridge.""" - episode_id = episode.get("episode_id") - prompt = episode.get("prompt") - if not isinstance(episode_id, str) or not isinstance(prompt, str) or not prompt: - raise legacy.EpisodeClientError("episode daemon returned invalid episode identity") - async with CompactionRLBridge( - miles_base_url, - prompt, - request_kwargs, - metrics, - max_seq_len=max_seq_len, - segment_base_urls=segment_base_urls, - ) as bridge, _CompactionAppServerDriver( - binary, bridge, client, episode_id, prompt, metrics - ) as driver: - return await driver.drive() diff --git a/yeto/rl/harness/codex/environment.py b/yeto/rl/harness/codex/environment.py deleted file mode 100644 index e45d0d8b8..000000000 --- a/yeto/rl/harness/codex/environment.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Terminal environment surface used by the Codex driver (ports path). - -The stock driver (``codex_harness_agent._AppServerDriver``) only needs two -calls from its environment: ``execute`` and ``submit``. Legacy bound those to -the SecRLEnv ``EpisodeClient``; the Terminal-Bench path binds them to a sandbox -HTTP endpoint instead. Verification (``evaluate``) is *not* part of this -surface: it belongs to the trusted layer (design D7 / R-D9). - -``HttpTerminalEnvironment`` is the untrusted-side client (no HMAC key, bearer -token only). ``serve_environment`` exposes any in-process environment over the -same wire so the worker subprocess can be exercised on CPU with a fake. -""" - -from __future__ import annotations - -import asyncio -import hmac -import json -from dataclasses import dataclass, field -from typing import Any, Protocol - -import aiohttp -from aiohttp import web - -from .client import EpisodeAPIError, EpisodeTransportError - -MAX_WIRE_BYTES = 1 << 20 - - -class TerminalEnvironment(Protocol): - async def execute( - self, episode_id: str, command: str, timeout_seconds: float, output_bytes: int - ) -> dict[str, Any]: ... - - async def submit(self, episode_id: str, submission: dict[str, Any]) -> dict[str, Any]: ... - - -class TrustedVerifier(Protocol): - """Runs the task verifier for one episode; trusted layer only.""" - - async def evaluate(self, episode_id: str) -> dict[str, Any]: ... - - -@dataclass -class FakeTerminalEnvironment: - """Scripted environment for CPU tests: ``outputs[command]`` or default echo.""" - - outputs: dict[str, dict[str, Any]] = field(default_factory=dict) - passed: bool = False - calls: list[tuple[str, str]] = field(default_factory=list) - submissions: list[dict[str, Any]] = field(default_factory=list) - evaluations: int = 0 - exec_delay_seconds: float = 0.0 - - async def execute( - self, episode_id: str, command: str, timeout_seconds: float, output_bytes: int - ) -> dict[str, Any]: - self.calls.append((episode_id, command)) - if self.exec_delay_seconds: - await asyncio.sleep(self.exec_delay_seconds) - return dict( - self.outputs.get( - command, - {"exit_code": 0, "output": f"ran:{command}", "timed_out": False, "truncated": False}, - ) - ) - - async def submit(self, episode_id: str, submission: dict[str, Any]) -> dict[str, Any]: - self.submissions.append(dict(submission)) - return {"accepted": True} - - async def evaluate(self, episode_id: str) -> dict[str, Any]: - self.evaluations += 1 - return {"passed": self.passed, "testsh_rc": 0 if self.passed else 1} - - -class HttpTerminalEnvironment: - """Untrusted-side client: ``POST {base}/execute`` and ``POST {base}/submit``.""" - - def __init__(self, base_url: str, token: str, *, total_timeout_seconds: float = 600.0): - self._base = base_url.rstrip("/") - self._token = token - self._timeout = aiohttp.ClientTimeout(total=total_timeout_seconds) - self._session: aiohttp.ClientSession | None = None - - async def __aenter__(self) -> "HttpTerminalEnvironment": - self._session = aiohttp.ClientSession(timeout=self._timeout) - return self - - async def __aexit__(self, *_exc: Any) -> None: - if self._session is not None: - await self._session.close() - self._session = None - - async def _post(self, route: str, payload: dict[str, Any]) -> dict[str, Any]: - if self._session is None: - raise EpisodeTransportError("environment client is not started") - try: - async with self._session.post( - f"{self._base}/{route}", - json=payload, - headers={"Authorization": f"Bearer {self._token}"}, - ) as response: - raw = await response.content.read(MAX_WIRE_BYTES + 1) - if len(raw) > MAX_WIRE_BYTES: - raise EpisodeTransportError("environment response is oversized") - if response.status != 200: - raise EpisodeAPIError( - response.status, "environment_error", raw.decode("utf-8", "replace")[:512] - ) - except (aiohttp.ClientError, TimeoutError) as exc: - raise EpisodeTransportError("environment transport failed") from exc - try: - body = json.loads(raw) - except ValueError as exc: - raise EpisodeTransportError("environment returned invalid JSON") from exc - if not isinstance(body, dict): - raise EpisodeTransportError("environment returned a non-object") - return body - - async def execute( - self, episode_id: str, command: str, timeout_seconds: float, output_bytes: int - ) -> dict[str, Any]: - return await self._post( - "execute", - { - "episode_id": episode_id, - "command": command, - "timeout_seconds": timeout_seconds, - "output_bytes": output_bytes, - }, - ) - - async def submit(self, episode_id: str, submission: dict[str, Any]) -> dict[str, Any]: - return await self._post("submit", {"episode_id": episode_id, "submission": submission}) - - -async def serve_environment( - env: TerminalEnvironment, token: str, *, host: str = "127.0.0.1" -) -> tuple[web.AppRunner, str]: - """Expose ``env`` on the worker wire. ``evaluate`` is deliberately not routed.""" - - def _authorized(request: web.Request) -> bool: - return hmac.compare_digest( - request.headers.get("Authorization", ""), f"Bearer {token}" - ) - - async def execute(request: web.Request) -> web.Response: - if not _authorized(request): - return web.json_response({"error": "unauthorized"}, status=401) - body = await request.json() - try: - result = await env.execute( - str(body["episode_id"]), - str(body["command"]), - timeout_seconds=float(body["timeout_seconds"]), - output_bytes=int(body["output_bytes"]), - ) - except EpisodeAPIError as exc: - return web.json_response({"error": exc.message}, status=exc.status) - return web.json_response(result) - - async def submit(request: web.Request) -> web.Response: - if not _authorized(request): - return web.json_response({"error": "unauthorized"}, status=401) - body = await request.json() - try: - result = await env.submit(str(body["episode_id"]), dict(body["submission"])) - except EpisodeAPIError as exc: - return web.json_response({"error": exc.message}, status=exc.status) - return web.json_response(result) - - app = web.Application(client_max_size=MAX_WIRE_BYTES) - app.router.add_post("/execute", execute) - app.router.add_post("/submit", submit) - runner = web.AppRunner(app, shutdown_timeout=1.0) - await runner.setup() - site = web.TCPSite(runner, host, 0) - await site.start() - port = site._server.sockets[0].getsockname()[1] # noqa: SLF001 - ephemeral port - return runner, f"http://{host}:{port}" diff --git a/yeto/rl/harness/codex/generate.py b/yeto/rl/harness/codex/generate.py deleted file mode 100644 index fec404d21..000000000 --- a/yeto/rl/harness/codex/generate.py +++ /dev/null @@ -1,23 +0,0 @@ -"""SecRLEnv agent-metadata capture (neutral part of the old generate wrapper). - -The custom agent (``agent.py`` / ``codex_harness_agent.py``) records its signed -result here; the Miles generate wrapper that restores it after Miles' aborted -fallback moved to ``yeto/rl/adapters/miles/harness_glue/codex_generate.py`` -(yeto-framework-decoupling 5.7). -""" - -from __future__ import annotations - -import copy -from contextvars import ContextVar -from typing import Any - -_CAPTURED_AGENT_METADATA: ContextVar[dict[str, Any] | None] = ContextVar( - "secrlenv_agent_metadata", default=None -) - - -def capture_agent_metadata(metadata: dict[str, Any]) -> None: - """Retain one agent result inside its generation task only.""" - - _CAPTURED_AGENT_METADATA.set(copy.deepcopy(metadata)) diff --git a/yeto/rl/harness/codex/pins.py b/yeto/rl/harness/codex/pins.py deleted file mode 100644 index bd502f83f..000000000 --- a/yeto/rl/harness/codex/pins.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Self-attestation pins for the Codex harness package (ports path). - -Provenance: moved from yeto commit 5bfc011 ``yeto_miles_secrlenv/`` (see -openspec change rl-codex-harness-rollout, design R-CTX). The legacy pins in -``yeto.rl`` (``SECRLENV_*_SHA256``, ``CODEX_HARNESS_AGENT_SHA256``) are owned by -the image line and still describe the legacy import path; the values below -describe the moved files. ``LEGACY_SOURCE_SHA256`` is kept for comparison. -""" - -AGENT_SHA256 = "84c6c708cba2142e3d043ae047c90f529b90b953b739261466b8a6e28b2c7411" -GENERATE_SHA256 = "1630dc8a3aae3cbbd4708a81c9803920791ee58d6725c3fcaa4736373f8f9e48" # adapters/miles/harness_glue/codex_generate.py (P0-1.7b.4-Q1 failed status + terminal-error record; was a3212183…) - -LEGACY_SOURCE_SHA256 = { - "agent.py": "0f76c7fbd81135bc5b02cab2488629aaff1bb58dc59eae9228ca317583d90c26", - "generate.py": "9e034d6b2e9fec642501ea4a638a8fe196819dacde614ce2903359fc54ea1713", - "codex_harness_agent.py": "94fa4c245b719d236ec1007b70d395adb12456b3ef04278592ae2d3c0d843947", -} -LEGACY_SOURCE_COMMIT = "5bfc011" - -# Build-time default backend profile of the Codex OpenEnv adapter. This is the -# value the image line bakes into the container as -# ``YETO_CODEX_OPENENV_BACKEND_PROFILE`` (``yeto/rl/__init__.py`` -# CODEX_OPENENV_IDENTITY_ENV) and is a *record of the image build*, not the -# profile a launch must use. rl-fn-codex-rollout 1.0: the runtime profile is -# declared per launch (``--codex-backend-profile``) and only has to belong to -# OPENENV_BACKEND_PROFILES; see codex_openenv_agent_function.resolve_backend_profile. -OPENENV_BACKEND_PROFILE = "qwen35_08b" -# Profiles the OpenEnv adapter may be driven with at runtime. Every member must -# declare ``model_identifier``/``model_revision`` in ``yeto.rl.codex_backend`` -# (checked by ``codex_openenv_agent_function.profile_identity``). The tool -# surface pins (``*_SHA256`` identity env) are profile-independent, so adding a -# profile here needs no image rebuild. -OPENENV_BACKEND_PROFILES = ("qwen35_08b", "qwen35", "qwen38", "qwen38_next", "qwen38_next_4layer", "qwen3_30b_a3b") # qwen35 (4B): S17 M1; qwen3_30b_a3b: rlf-codex-rl 5.6; qwen38 (27B): S24 diff --git a/yeto/rl/harness/codex/preflight.py b/yeto/rl/harness/codex/preflight.py deleted file mode 100644 index d6e4f7e94..000000000 --- a/yeto/rl/harness/codex/preflight.py +++ /dev/null @@ -1,360 +0,0 @@ -"""Ports-path preflight for the Codex Terminal-Bench harness (design D3, R-CTX, R-D5a). - -``harness_preflight(miles_args, launch)`` is the IR-1 hook (INFRA signature -``entry.HarnessPreflight = Callable[[miles_args, launch], None]``): ports -``entry.preflight_stage`` calls it before ``connect_island_ray`` / placement / -model allocation, resolved from ``miles_args.yeto_harness_preflight`` or -``YETO_HARNESS_PREFLIGHT`` (= ``HARNESS_PREFLIGHT_SPEC``). It fails closed on: -- ``reward_scope`` other than ``trajectory`` (task 7.1; IR-1 - ``config.check_harness_reward_scope`` on ``miles_args.yeto_harness_reward_scope``); -- an agent function other than this package's ``CODEX_OPENENV_AGENT``; -- no environment provider (``miles_args.yeto_harness_environment_provider`` or - ``YETO_HARNESS_ENVIRONMENT_PROVIDER`` = ``module:callable`` returning an - ``EnvironmentProvider``), then installs it together with the island's - ``ToolWaitBoard`` / ``HarnessBoard`` actors (same names ``entry.harness_source`` uses); -and, through ``preflight_codex_openenv``, on: -- legacy trainable compaction being enabled: the fork pin's session server does - not know the ``X-Miles-Compaction-*`` headers and would silently roll back - window 0 (R-D5a), so ``YETO_CODEX_COMPACTION_ENABLED`` must be unset/false; -- the CPU-only scripted worker driver being allowed in a training process; -- Codex binary / version / identity drift (``codex_harness_agent._attest_runtime``: - signed env, binary sha256 and size, ``codex --version`` == 0.145.0); -- OpenEnv identity env drift against ``codex_openenv_agent_function._OPENENV_IDENTITY_ENV``. -""" - -from __future__ import annotations - -import importlib -import os -from types import SimpleNamespace -from collections.abc import Callable, Mapping -from pathlib import Path -from typing import Any - -from yeto.rl.engine.members import MEMBER_PREFIX, member_id -from yeto.rl.engine.tool_wait import LazyBoardActor, harness_board_actor -from yeto.rl.harness.reward_scope import check_harness_reward_scope - -from . import codex_harness_agent as harness -from . import codex_openenv_agent_function as adapter -from .codex_openenv_agent_worker import SCRIPTED_DRIVER_ENV - -COMPACTION_ENV = ( - "YETO_CODEX_COMPACTION_ENABLED", - "YETO_CODEX_COMPACTION_TRIGGER_TOKENS", - "YETO_CODEX_COMPACTION_SUMMARY_MAX_TOKENS", - "YETO_CODEX_MAX_COMPACTIONS", -) -_FALSE = frozenset({"", "0", "false", "no", "off"}) - - -class PreflightError(RuntimeError): - pass - - -def assert_compaction_disabled(env: Mapping[str, str]) -> None: - if env.get("YETO_CODEX_COMPACTION_ENABLED", "").strip().lower() not in _FALSE: - raise PreflightError( - "YETO_CODEX_COMPACTION_ENABLED is set: legacy trainable compaction is " - "incompatible with the fork pin session server (R-D5a) and must stay off" - ) - extra = [name for name in COMPACTION_ENV[1:] if env.get(name)] - if extra: - raise PreflightError("compaction tuning without support: " + ", ".join(extra)) - - -def assert_compactionrl_consistent(miles_args: Any, env: Mapping[str, str]) -> None: - """CompactionRL switch vs the trainer's ``--gae-variant`` (progress "S13 Codex 压缩接线"). - - This hook only runs for the Codex OpenEnv agent, which can compact, so: - ``cross_segment_per_sample`` requires ``YETO_CODEX_COMPACTIONRL`` on; any - other variant requires it off; ``cross_segment_whole_rollout`` (control arm) - is rejected outright (per-segment sessions cannot form one sample). - """ - from yeto.rl.algos import compactionrl as crl - - try: - switch = crl.compaction_switch(env.get(crl.COMPACTION_SWITCH_ENV)) - crl.check_rollout_compaction( - getattr(miles_args, "gae_variant", None), - switch, - harness_compacts=True, - t_comp=env.get(crl.COMPACTION_T_COMP_ENV), - ) - except ValueError as exc: - raise PreflightError(f"CompactionRL: {exc}") from exc - - -def assert_no_scripted_driver(env: Mapping[str, str]) -> None: - if env.get(SCRIPTED_DRIVER_ENV): - raise PreflightError(f"{SCRIPTED_DRIVER_ENV} must not be set in a training process") - - -def assert_openenv_identity(env: Mapping[str, str]) -> None: - drift = [name for name, value in adapter._OPENENV_IDENTITY_ENV.items() if env.get(name) != value] - if drift: - raise PreflightError("Codex OpenEnv identity drifted: " + ", ".join(drift)) - - -def preflight_codex_openenv( - env: Mapping[str, str] | None = None, - *, - attest_runtime: Callable[[], Path] = harness._attest_runtime, -) -> dict[str, Any]: - """Return a small report on success; raise ``PreflightError`` otherwise.""" - env = os.environ if env is None else env - assert_compaction_disabled(env) - assert_no_scripted_driver(env) - assert_openenv_identity(env) - try: - binary = attest_runtime() - except harness.CodexHarnessError as exc: - raise PreflightError(f"Codex runtime attestation failed: {exc}") from exc - return { - "codex_binary": str(binary), - "codex_version": harness.CODEX_CLI_VERSION, - "identity": adapter.codex_openenv_harness_identity(), - "compaction": "disabled", - } - - -# --------------------------------------------------------------------------- -# IR-1 hook: (miles_args, launch) -> None, run by entry.preflight_stage - -HARNESS_PREFLIGHT_SPEC = "yeto.rl.harness.codex.preflight:harness_preflight" -ENVIRONMENT_PROVIDER_ENV = "YETO_HARNESS_ENVIRONMENT_PROVIDER" -EXPECTED_AGENT_FUNCTION = "yeto.rl.harness.codex.codex_openenv_subprocess_agent_function.run" - - -def _load_dotted(spec: Any, what: str) -> Any: - if callable(spec): - return spec - module, sep, name = str(spec).partition(":") - if not sep: - module, _, name = module.rpartition(".") - if not module or not name: - raise PreflightError(f"{what} {spec!r} is not module:callable") - try: - target = getattr(importlib.import_module(module), name) - except (ImportError, AttributeError) as exc: - raise PreflightError(f"{what} {spec!r} cannot be imported: {exc}") from exc - if not callable(target): - raise PreflightError(f"{what} {spec!r} is not callable") - return target - - -def resolve_environment_provider(miles_args: Any, env: Mapping[str, str]) -> Any: - """The sandbox provider factory for this island (fail closed when absent).""" - spec = getattr(miles_args, "yeto_harness_environment_provider", None) or env.get(ENVIRONMENT_PROVIDER_ENV) - if not spec: - raise PreflightError( - f"no environment provider: set miles_args.yeto_harness_environment_provider or {ENVIRONMENT_PROVIDER_ENV}" - ) - factory = _load_dotted(spec, "environment provider") - provider = factory(miles_args) if not hasattr(factory, "acquire") else factory - if not hasattr(provider, "acquire"): - raise PreflightError(f"environment provider {spec!r} has no acquire()") - return provider - - -def island_boards(miles_args: Any) -> tuple[Any, Any]: - """``(tool_wait_board, harness_board)``: the island's named actors, looked up lazily.""" - learner_id = int(getattr(miles_args, "yeto_rl_learner_id", 0) or 0) - return LazyBoardActor(learner_id), LazyBoardActor(learner_id, factory=harness_board_actor) - - -def assert_task_prompts(miles_args: Any, provider: Any) -> None: - """2026-10-07 (S15 stage 2): every row must resolve to a real task statement. - Runs when the provider exposes its TB2 ``tasks_dir`` and the prompt data file - is readable here; the launch script's PLAN_ONLY runs the same check locally.""" - tasks_dir = getattr(provider, "tasks_dir", None) - data = getattr(miles_args, "prompt_data", None) - if tasks_dir is None or not isinstance(data, str) or not Path(data).expanduser().is_file(): - return - from .tb2_provider import TaskPromptPreflightError, preflight_task_prompts - - try: - preflight_task_prompts(Path(data).expanduser(), Path(tasks_dir)) - except TaskPromptPreflightError as exc: - raise PreflightError(str(exc)) from exc - - -def harness_preflight(miles_args: Any, launch: Any, *, env: Mapping[str, str] | None = None) -> None: - """IR-1 hook body. Raises ``PreflightError``/``HarnessConfigError`` before any allocation.""" - del launch # identity / binary / key checks do not depend on the launch args - env = os.environ if env is None else env - # 7.1: a per-segment reward scope must not start (IR-1 path, same check as validate_parsed_args). - check_harness_reward_scope(getattr(miles_args, "yeto_harness_reward_scope", None)) - agent = getattr(miles_args, "custom_agent_function_path", None) - if agent != EXPECTED_AGENT_FUNCTION: - raise PreflightError(f"custom_agent_function_path={agent!r}; the Codex harness preflight expects {EXPECTED_AGENT_FUNCTION}") - preflight_codex_openenv(env) - assert_compactionrl_consistent(miles_args, env) - provider = resolve_environment_provider(miles_args, env) - assert_task_prompts(miles_args, provider) - from . import codex_openenv_subprocess_agent_function as subprocess_agent - - tool_wait_board, harness_board = island_boards(miles_args) - subprocess_agent.configure( - provider=provider, - tool_wait_board=tool_wait_board, - harness_board=harness_board, - member=resolve_member(miles_args, env), - ) - - -MEMBER_CELL_ENV = "YETO_RL_CELL_ID" -LEARNER_ID_ENV = "YETO_RL_LEARNER_ID" -# Harness env the island driver forwards into Ray's job runtime_env so the -# rollout workers (which inherit the raylet's environment, not the driver's) -# can self-configure (A-T3-4, codex-smoke-20261003-6). -WORKER_PASSTHROUGH_ENV = ( - ENVIRONMENT_PROVIDER_ENV, - # secret-handling-hardening: the sandbox-only Modal token, never MODAL_TOKEN_*. - "YETO_SANDBOX_MODAL_TOKEN_ID", - "YETO_SANDBOX_MODAL_TOKEN_SECRET", - "OPENENV_RUN_ID", - "SECRLENV_MAX_TURNS", - MEMBER_CELL_ENV, - LEARNER_ID_ENV, -) -WORKER_PASSTHROUGH_ENV_PREFIXES = ("YETO_HARNESS_TB2_", "YETO_CODEX_") - - -def worker_runtime_env(miles_args: Any, env: Mapping[str, str] | None = None) -> dict[str, str]: - """Env vars every Ray actor of this island must see for the agent function to run.""" - env = os.environ if env is None else env - forwarded = { - key: value - for key, value in env.items() - if value and (key in WORKER_PASSTHROUGH_ENV or key.startswith(WORKER_PASSTHROUGH_ENV_PREFIXES)) - } - learner_id = getattr(miles_args, "yeto_rl_learner_id", None) - if learner_id is not None and str(learner_id) != "": - forwarded[LEARNER_ID_ENV] = str(int(learner_id)) - cell = getattr(miles_args, "yeto_rl_cell_id", None) - if cell is not None and str(cell) != "": - forwarded[MEMBER_CELL_ENV] = str(cell) - return forwarded - - -def configure_rollout_worker(env: Mapping[str, str] | None = None) -> bool: - """Install the provider and boards in a rollout worker from its environment. - - ``harness_preflight`` configures the driver process only; upstream Miles - calls the agent function inside ``RolloutExecutor`` actors, where the - module globals start empty. Returns False (configures nothing) when the - environment names no provider, so ``run`` keeps failing closed. - """ - env = os.environ if env is None else env - if not env.get(ENVIRONMENT_PROVIDER_ENV): - return False - miles_args = SimpleNamespace( - yeto_harness_environment_provider=None, - yeto_rl_learner_id=int(env.get(LEARNER_ID_ENV) or 0), - yeto_rl_member_id=None, - # A16 (D5): the island entry exports its cell (publish_member_cell) and - # worker_runtime_env forwards it; read it back into the field so the - # worker resolves the same member key as the driver. - yeto_rl_cell_id=env.get(MEMBER_CELL_ENV) or None, - ) - provider = resolve_environment_provider(miles_args, env) - assert_task_prompts(miles_args, provider) - from . import codex_openenv_subprocess_agent_function as subprocess_agent - - tool_wait_board, harness_board = island_boards(miles_args) - subprocess_agent.configure( - provider=provider, - tool_wait_board=tool_wait_board, - harness_board=harness_board, - member=resolve_member(miles_args, env), - ) - return True - - -def resolve_member(miles_args: Any, env: Mapping[str, str] | None = None) -> str | None: - """The INFRA rollout member key this island's sessions are admitted under. - - INFRA keys members as ``rollout.member_id(cell_id)`` (= ``engine:``, - ``MilesRolloutPool.members`` / ``close_admission``). Sources, in order: - ``miles_args.yeto_rl_member_id`` (a full member key, or a bare cell id), - ``miles_args.yeto_rl_cell_id`` / ``YETO_RL_CELL_ID`` (a cell id). None - (no source) keeps the global admission key, which single-island smoke - runs rely on. - """ - env = os.environ if env is None else env - member = getattr(miles_args, "yeto_rl_member_id", None) - if member is not None and str(member) != "": - member = str(member) - return member if member.startswith(MEMBER_PREFIX) else member_id(member) - cell = getattr(miles_args, "yeto_rl_cell_id", None) - if cell is None or str(cell) == "": - cell = env.get(MEMBER_CELL_ENV) or None - return member_id(cell) if cell is not None else None - - -# --------------------------------------------------------------------------- -# Task 2.2: legacy preflight forwarding (injectable; legacy modules are not edited here) -# -# ``yeto/rl/adapters/miles/island_entry.py::_preflight_codex_openenv_adapter`` and -# ``yeto/rl/tbench_direct_preflight.py::_attest_adapter`` currently import the -# adapter from ``/examples/experimental/openenv`` and compare it -# against ``yeto.rl.CODEX_OPENENV_IDENTITY_ENV`` (image-line pins). Once the -# image line repoints those pins (see ``required_pin_updates``), the legacy -# functions become one-liners calling ``forward_legacy_openenv_preflight``. - -def forward_legacy_openenv_preflight(args: Any, profile_name: str, env: Mapping[str, str] | None = None) -> None: - """Drop-in body for ``learner._preflight_codex_openenv_adapter(args, profile)``. - - Same failure classes/messages as legacy so existing callers keep their - behaviour: ``ValueError`` with "requires backend profile", "identity drifted", - "environment drifted". - """ - del args # the adapter no longer lives under miles_root - # rl-fn-codex-rollout 1.0: the runtime profile only has to be an allowlisted - # OpenEnv profile with a declared HF identity (``adapter.profile_identity``); - # ``learner._preflight_codex_harness`` has already matched it against - # ``--model`` / ``--model-revision`` (``validate_stock_codex_fields``). The - # identity env below is the image's build-time record and is compared as - # such, so a new profile needs no image rebuild. - adapter.profile_identity(profile_name) - live = adapter.codex_openenv_harness_identity() - expected = { - name.removeprefix("YETO_CODEX_OPENENV_").lower(): value - for name, value in adapter._OPENENV_IDENTITY_ENV.items() - if name.endswith("_SHA256") - } - if live != expected: - raise ValueError("the Codex OpenEnv surface identity drifted") - env = os.environ if env is None else env - mismatched = [name for name, value in adapter._OPENENV_IDENTITY_ENV.items() if env.get(name) != value] - if mismatched: - raise ValueError("Codex OpenEnv container environment drifted: " + ", ".join(mismatched)) - - -def required_pin_updates() -> dict[str, Any]: - """What ``yeto/rl/__init__.py`` (image line) must change to point at this package.""" - import hashlib - - here = Path(__file__).resolve().parent - sha = lambda name: hashlib.sha256((here / name).read_bytes()).hexdigest() # noqa: E731 - return { - "SECRLENV_AGENT_PATH": "yeto/rl/harness/codex/agent.py", - "SECRLENV_AGENT_SHA256": sha("agent.py"), - "SECRLENV_AGENT": "yeto.rl.harness.codex.agent.run", - "SECRLENV_REWARD": "yeto.rl.harness.codex.reward:reward_func", - "SECRLENV_GROUP_FILTER": "yeto.rl.harness.codex.reward.check_group", - "SECRLENV_GENERATE": "yeto.rl.adapters.miles.harness_glue.codex_generate.generate", - "SECRLENV_GENERATE_SHA256": hashlib.sha256( # decoupling 5.7: glue lives in the adapter - (here.parents[1] / "adapters" / "miles" / "harness_glue" / "codex_generate.py").read_bytes() - ).hexdigest(), - "CODEX_HARNESS_AGENT": "yeto.rl.harness.codex.codex_harness_agent.run", - "CODEX_HARNESS_AGENT_PATH": "yeto/rl/harness/codex/codex_harness_agent.py", - "CODEX_HARNESS_AGENT_SHA256": sha("codex_harness_agent.py"), - "CODEX_OPENENV_AGENT": "yeto.rl.harness.codex.codex_openenv_subprocess_agent_function.run", - "CODEX_OPENENV_AGENT_MODULES": ( - "codex_openenv_subprocess_agent_function.py", "codex_openenv_agent_worker.py", "codex_openenv_agent_function.py", - "compaction_bridge.py", - ), - "CODEX_OPENENV_IDENTITY_ENV": dict(adapter._OPENENV_IDENTITY_ENV), - } diff --git a/yeto/rl/harness/codex/reward.py b/yeto/rl/harness/codex/reward.py deleted file mode 100644 index cc3a27762..000000000 --- a/yeto/rl/harness/codex/reward.py +++ /dev/null @@ -1,377 +0,0 @@ -"""Validated reward extraction for secrlenv Miles samples.""" - -from __future__ import annotations - -import json -import math -import re -import statistics -from typing import Any - -from yeto.rl.rewards.types import RewardResult, Trajectory - -OUTCOME_KEY = "secrlenv_trusted_outcome" -INFRASTRUCTURE_STATUS = "infrastructure_error" -CLEANUP_ERROR_STATUS = "cleanup_error" -INFRASTRUCTURE_503_RETRIES_KEY = "secrlenv_infrastructure_503_retries" -ADMISSION_SCHEMA = 2 -ADMISSION_PHASE_KEY = "admission_phase" -ADMISSION_ERROR_CODE_KEY = "admission_error_code" -ADMISSION_NONCE_KEY = "admission_nonce" -ADMISSION_PHASE = "pre_create" -ADMISSION_ERROR_CODE = "infrastructure_error" -_ADMISSION_NONCE = re.compile(r"[0-9a-f]{32}") -_OUTCOME_STATUSES = { - "completed", - "timeout", - "max_turns", - "max_seq_len", - INFRASTRUCTURE_STATUS, - CLEANUP_ERROR_STATUS, -} - - -class SecRLEnvRolloutTerminalError(RuntimeError): - """Sanitized base class for non-recoverable SecRLEnv rollout failures.""" - - -class SecRLEnvReplacementExhausted(SecRLEnvRolloutTerminalError): - """The one authenticated same-task replacement did not restore progress.""" - - -class SecRLEnvUntrustedEvidence(SecRLEnvRolloutTerminalError): - """A rollout group did not carry authentic SecRLEnv evidence.""" - - -class SecRLEnvRolloutCleanupError(SecRLEnvRolloutTerminalError): - """Pinned Miles could not drain a terminal SecRLEnv rollout.""" - - -class UntrustedOutcome(RuntimeError): - pass - - -def _canonical(value: dict[str, Any]) -> bytes: - return json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8") - - -def _verified_outcome(metadata: Any) -> tuple[dict[str, Any], float]: - if not isinstance(metadata, dict): - raise UntrustedOutcome("sample metadata is not an object") - outcome = metadata.get(OUTCOME_KEY) - if not isinstance(outcome, dict): - raise UntrustedOutcome("sample has no secrlenv outcome") - schema = outcome.get("schema") - status = outcome.get("status") - if schema not in {1, ADMISSION_SCHEMA} or status not in _OUTCOME_STATUSES: - raise UntrustedOutcome("secrlenv outcome has an invalid schema or status") - infrastructure_503_retries = outcome.get(INFRASTRUCTURE_503_RETRIES_KEY) - if ( - type(infrastructure_503_retries) is not int - or infrastructure_503_retries not in {0, 1} - ): - raise UntrustedOutcome("secrlenv outcome has an invalid retry ledger") - reward = outcome.get("reward") - if isinstance(reward, bool) or not isinstance(reward, (int, float)): - raise UntrustedOutcome("secrlenv outcome reward is not numeric") - value = float(reward) - if not math.isfinite(value) or not 0.0 <= value <= 1.0: - raise UntrustedOutcome("secrlenv outcome reward is outside [0, 1]") - if not isinstance(outcome.get("task_id"), str): - raise UntrustedOutcome("secrlenv outcome has no episode/task identity") - if schema == ADMISSION_SCHEMA: - expected_keys = { - "schema", - "status", - "episode_id", - "task_id", - "reward", - "passed", - "class", - INFRASTRUCTURE_503_RETRIES_KEY, - ADMISSION_PHASE_KEY, - ADMISSION_ERROR_CODE_KEY, - ADMISSION_NONCE_KEY, - } - if ( - set(outcome) != expected_keys - or status != INFRASTRUCTURE_STATUS - or outcome.get("episode_id") is not None - or not outcome["task_id"] - or outcome.get(ADMISSION_PHASE_KEY) != ADMISSION_PHASE - or outcome.get(ADMISSION_ERROR_CODE_KEY) != ADMISSION_ERROR_CODE - or not isinstance(outcome.get(ADMISSION_NONCE_KEY), str) - or _ADMISSION_NONCE.fullmatch(outcome[ADMISSION_NONCE_KEY]) is None - or infrastructure_503_retries != 1 - or value != 0.0 - or outcome.get("passed") is not False - or outcome.get("class") is not None - ): - raise UntrustedOutcome("secrlenv admission outcome is invalid") - elif ( - not isinstance(outcome.get("episode_id"), str) - or any( - key in outcome - for key in ( - ADMISSION_PHASE_KEY, - ADMISSION_ERROR_CODE_KEY, - ADMISSION_NONCE_KEY, - ) - ) - ): - raise UntrustedOutcome("secrlenv outcome has no episode/task identity") - return outcome, value - - -def verify_outcome(metadata: Any) -> float: - outcome, value = _verified_outcome(metadata) - if outcome.get("status") in {INFRASTRUCTURE_STATUS, CLEANUP_ERROR_STATUS}: - raise UntrustedOutcome( - "secrlenv non-verdict outcome is not a grader verdict" - ) - return value - - -INFRASTRUCTURE_ABORT_REASON = "secrlenv_infrastructure_failure" - - -def secrlenv_reward(trajectory: Trajectory) -> RewardResult: - """Neutral form (decoupling 3.3): only structurally valid rewards. - - Missing or invalid evidence raises :class:`UntrustedOutcome` (fatal). - An infrastructure marker is not a - verdict: it returns value 0.0 with ``aborted=True`` so the backend adapter - marks the sample with its abort status and the group is replaced. - """ - - outcome, value = _verified_outcome(trajectory.metadata) - if outcome.get("status") != INFRASTRUCTURE_STATUS: - return RewardResult(value=value) - return RewardResult(value=0.0, aborted=True, reason=INFRASTRUCTURE_ABORT_REASON) - - -def _sample_reward(sample: Any) -> float: - result = secrlenv_reward(Trajectory(metadata=getattr(sample, "metadata", None))) - if not result.aborted: - return result.value - - # Miles calls the reward function before the dynamic group filter. Mark an - # infrastructure failure FAILED here so the group check replaces it, while - # still returning the numeric value required by the reward-function API. - # The group check never permits this transient zero into training data. - # FAILED, not ABORTED: Miles drops a group with an aborted sample before - # the group check runs, so its same-task retry would never be scheduled. - from miles.utils.types import Sample - - sample.status = Sample.Status.FAILED - return result.value - - -async def reward_func(args: Any, samples: Any, **_kwargs: Any) -> float | list[float]: - """Return only structurally valid host-side rewards. - - Missing or invalid metadata remains fatal. An infrastructure marker - is returned as the numeric zero required by Miles only after marking its - sample FAILED, so the group check replaces it and it never enters training. - """ - - try: - if isinstance(samples, list): - return [_sample_reward(sample) for sample in samples] - return _sample_reward(samples) - except UntrustedOutcome as error: - remember_terminal_error(args, error) - raise - - -def sample_statuses(samples: Any) -> list[str | None]: - """Status value of every sample in ``samples`` (nested lists flattened).""" - - values: list[str | None] = [] - pending = list(samples) if isinstance(samples, list) else [] - while pending: - item = pending.pop(0) - if isinstance(item, list): - pending[0:0] = item - continue - status = getattr(item, "status", None) - value = getattr(status, "value", status) - values.append(value if isinstance(value, str) else None) - return values - - -# Miles catches an exception raised inside a generate task and only counts the -# group as failed; the rollout then fails later with a different error. The -# SecRLEnv code records the first terminal error here so the backend can -# raise it instead. -TERMINAL_ERROR_ATTR = "_yeto_secrlenv_terminal_error" - - -def remember_terminal_error(args: Any, error: BaseException) -> None: - if args is not None and getattr(args, TERMINAL_ERROR_ATTR, None) is None: - try: - setattr(args, TERMINAL_ERROR_ATTR, error) - except AttributeError: - pass - - -def _flatten_samples(samples: list[Any]) -> list[Any]: - flattened: list[Any] = [] - for item in samples: - flattened.extend(item if isinstance(item, list) else [item]) - return flattened - - -def _outcome_fingerprint(outcome: Any) -> bytes | None: - """Return a hashable outcome identity; malformed outcomes fail later.""" - - if not isinstance(outcome, dict): - return None - try: - return _canonical(outcome) - except (TypeError, ValueError): - return None - - -def _group_key(samples: list[Any]) -> tuple[Any, ...]: - """Identify a generated group across Miles' repeated callbacks. - - Sample indexes are stable across the generation-time and all-samples - callbacks. The episode identity and canonical outcome prevent an index - collision from reusing a decision for different external-environment - evidence. - """ - - values = [] - for sample in samples: - metadata = getattr(sample, "metadata", None) - outcome = metadata.get(OUTCOME_KEY) if isinstance(metadata, dict) else None - index = getattr(sample, "index", None) - values.append( - ( - index if index is not None else id(sample), - outcome.get("episode_id") if isinstance(outcome, dict) else None, - outcome.get("task_id") if isinstance(outcome, dict) else None, - _outcome_fingerprint(outcome), - ) - ) - return tuple(values) - - -def _replacement_limit(args: Any) -> int | None: - value = getattr(args, "yeto_rl_dynamic_sampling_max_replacements", None) - try: - limit = int(value) if value is not None else None - except (TypeError, ValueError): - return None - return limit if limit is not None and limit >= 0 else None - - -def check_group(args: Any, samples: list[Any], **_kwargs: Any): - """Validate groups and coordinate bounded same-task replacements.""" - - from miles.rollout.filter_hub.base_types import DynamicFilterOutput - - flattened = _flatten_samples(samples) - rollout_id = getattr(args, "yeto_rl_policy_version", None) - state = getattr(args, "_yeto_secrlenv_filter_state", None) - if state is None or state.get("rollout_id") != rollout_id: - state = { - "rollout_id": rollout_id, - "rejections": 0, - "nontrainable_rejections": 0, - "forced": 0, - "decisions": {}, - } - args._yeto_secrlenv_filter_state = state - - key = _group_key(flattened) - previous = state["decisions"].get(key) - if previous is not None: - return DynamicFilterOutput(keep=previous[0], reason=previous[1]) - - try: - verified = [_verified_outcome(sample.metadata) for sample in flattened] - except UntrustedOutcome: - raise SecRLEnvUntrustedEvidence( - "SecRLEnv rollout evidence validation failed" - ) from None - - statuses = sample_statuses(flattened) - infrastructure = any( - outcome.get("status") == INFRASTRUCTURE_STATUS - for outcome, _ in verified - ) - cleanup_failure = any( - outcome.get("status") == CLEANUP_ERROR_STATUS - for outcome, _ in verified - ) - if cleanup_failure: - raise SecRLEnvRolloutCleanupError( - "SecRLEnv episode cleanup could not be verified" - ) - incomplete = any( - status not in {"completed", "truncated"} for status in statuses - ) - if infrastructure or incomplete: - if infrastructure: - reason = "secrlenv_infrastructure_failure" - elif "pending" in statuses: - reason = "secrlenv_pending" - elif "failed" in statuses: - reason = "secrlenv_failed" - else: - reason = "secrlenv_incomplete_status" - decision = (False, reason) - state["decisions"][key] = decision - state["nontrainable_rejections"] += 1 - retry = getattr(args, "_yeto_secrlenv_retry_callback", None) - if not callable(retry): - raise SecRLEnvReplacementExhausted( - "SecRLEnv rollout replacement scheduler is unavailable" - ) - retry_ledger: dict[int, int] = {} - for sample, (outcome, _) in zip(flattened, verified, strict=True): - try: - index = int(sample.index) - except (AttributeError, TypeError, ValueError): - raise SecRLEnvReplacementExhausted( - "SecRLEnv rollout replacement task identity is unavailable" - ) from None - value = outcome[INFRASTRUCTURE_503_RETRIES_KEY] - previous_retry_count = retry_ledger.setdefault(index, value) - if previous_retry_count != value: - raise SecRLEnvUntrustedEvidence( - "SecRLEnv rollout retry ledger is inconsistent" - ) - retry( - flattened, - evidence_key=key, - reason=reason, - infrastructure_503_retries=retry_ledger, - ) - return DynamicFilterOutput(keep=False, reason=reason) - - rewards = [value for _, value in verified] - - std = statistics.pstdev(rewards) if rewards else 0.0 - if math.isfinite(std) and std > 1e-8: - decision = (True, None) - else: - limit = _replacement_limit(args) - if limit is None: - raise SecRLEnvReplacementExhausted( - "SecRLEnv zero-variance replacement contract is missing" - ) - if state["rejections"] < limit: - state["rejections"] += 1 - value = round(rewards[0], 3) if rewards else 0.0 - decision = (False, f"secrlenv_zero_std_{value}") - else: - state["forced"] += 1 - decision = ( - True, - f"secrlenv_bounded_fallback_after_{limit}_replacements", - ) - state["decisions"][key] = decision - return DynamicFilterOutput(keep=decision[0], reason=decision[1]) diff --git a/yeto/rl/harness/codex/tb2_network_grants.py b/yeto/rl/harness/codex/tb2_network_grants.py deleted file mode 100644 index d29bcac2c..000000000 --- a/yeto/rl/harness/codex/tb2_network_grants.py +++ /dev/null @@ -1,104 +0,0 @@ -"""Explicit network grants for the 89 Terminal-Bench-2 tasks (rl-agentic-reward-env task 5.1). - -Sandbox egress is closed unless a task is listed in a network policy. Every -TB2 task declares ``allow_internet = true`` upstream and its verifier -(``tests/test.sh``) installs uv/pip/apt packages at grading time, and the -agentic runs that passed so far (S17 G3 positive control, FN codex r3) ran -with open egress. To keep them working, each task is granted ``open`` -explicitly here. Narrowing these to a domain allowlist needs a Modal CPU -check first (not done; see design.md D-NET). -""" - -from __future__ import annotations - -TB2_OPEN_EGRESS_TASKS: tuple[str, ...] = ( - "adaptive-rejection-sampler", - "bn-fit-modify", - "break-filter-js-from-html", - "build-cython-ext", - "build-pmars", - "build-pov-ray", - "caffe-cifar-10", - "cancel-async-tasks", - "chess-best-move", - "circuit-fibsqrt", - "cobol-modernization", - "code-from-image", - "compile-compcert", - "configure-git-webserver", - "constraints-scheduling", - "count-dataset-tokens", - "crack-7z-hash", - "custom-memory-heap-crash", - "db-wal-recovery", - "distribution-search", - "dna-assembly", - "dna-insert", - "extract-elf", - "extract-moves-from-video", - "feal-differential-cryptanalysis", - "feal-linear-cryptanalysis", - "filter-js-from-html", - "financial-document-processor", - "fix-code-vulnerability", - "fix-git", - "fix-ocaml-gc", - "gcode-to-text", - "git-leak-recovery", - "git-multibranch", - "gpt2-codegolf", - "headless-terminal", - "hf-model-inference", - "install-windows-3.11", - "kv-store-grpc", - "large-scale-text-editing", - "largest-eigenval", - "llm-inference-batching-scheduler", - "log-summary-date-ranges", - "mailman", - "make-doom-for-mips", - "make-mips-interpreter", - "mcmc-sampling-stan", - "merge-diff-arc-agi-task", - "model-extraction-relu-logits", - "modernize-scientific-stack", - "mteb-leaderboard", - "mteb-retrieve", - "multi-source-data-merger", - "nginx-request-logging", - "openssl-selfsigned-cert", - "overfull-hbox", - "password-recovery", - "path-tracing", - "path-tracing-reverse", - "polyglot-c-py", - "polyglot-rust-c", - "portfolio-optimization", - "protein-assembly", - "prove-plus-comm", - "pypi-server", - "pytorch-model-cli", - "pytorch-model-recovery", - "qemu-alpine-ssh", - "qemu-startup", - "query-optimize", - "raman-fitting", - "regex-chess", - "regex-log", - "reshard-c4-data", - "rstan-to-pystan", - "sam-cell-seg", - "sanitize-git-repo", - "schemelike-metacircular-eval", - "sparql-university", - "sqlite-db-truncate", - "sqlite-with-gcov", - "torch-pipeline-parallelism", - "torch-tensor-parallelism", - "train-fasttext", - "tune-mjcf", - "video-processing", - "vulnerable-secret", - "winning-avg-corewars", - "write-compressor", -) diff --git a/yeto/rl/harness/codex/tb2_provider.py b/yeto/rl/harness/codex/tb2_provider.py deleted file mode 100644 index 5a46718e8..000000000 --- a/yeto/rl/harness/codex/tb2_provider.py +++ /dev/null @@ -1,980 +0,0 @@ -"""Terminal-Bench-2 ``EnvironmentProvider`` (design R-TB / gap G1). - -One sandbox per trajectory, started from the task's **official image** -(``task.toml [environment].docker_image``, the exact image the TB2 harness -itself runs). The sandbox runs nothing but ``sleep infinity``; every agent -command is relayed from the *trusted* side through the backend's ``exec``. -The worker-facing wire (``/execute`` ``/submit``) is therefore served by -``environment.serve_environment`` **in the trusted process**, and the untrusted -worker only ever sees ``http://127.0.0.1:`` plus a bearer token — exactly -the surface ``HttpTerminalEnvironment`` already speaks. Compared with the -upstream recipe (``tb2_sandbox_modal`` + an OpenEnv ``tbench2_env`` server -layer baked into every task image) this needs no extra image layers, no -tunnel and no Python inside the task image. - -Verification mirrors the official harness's stage-at-verify model: ``tests/`` -never enters the sandbox before ``evaluate`` (the agent cannot read the -checks), is staged at ``/tests`` from the trusted side when the episode is -over, ``bash /tests/test.sh`` runs in the task workdir and the reward is the -``/logs/verifier/reward.txt`` the script writes (``1``/``0``). ``solution/`` -is never shipped at all (it exists only for the R-TB positive-path proof, run -by the operator through the same relay). - -Backends implement ``SandboxBackend``; ``LocalProcessBackend`` (CPU tests and -dev) runs commands in a scratch directory on this host, ``ModalSandboxBackend`` -materialises a Modal Sandbox. ``module:callable`` factories for -``YETO_HARNESS_ENVIRONMENT_PROVIDER``:: - - yeto.rl.harness.codex.tb2_provider:modal_provider - yeto.rl.harness.codex.tb2_provider:local_provider - -Configuration (environment variables, read when the factory runs): - -``YETO_HARNESS_TB2_TASKS_DIR`` terminal-bench-2 checkout root (required); - ``task_id`` = task directory name (the - ``metadata.task_id`` make_tbench2_data.py writes) -``YETO_HARNESS_TB2_LEASE_SECONDS`` hard per-trajectory wall clock (default: - the task's ``[agent].timeout_sec``, 900) -``YETO_HARNESS_TB2_MODAL_APP`` Modal app the sandboxes are created under - (default ``yeto-tbench2``; sweep scope) -``YETO_HARNESS_TB2_SANDBOX_TTL_S`` / ``YETO_HARNESS_TB2_IDLE_TIMEOUT_S`` - Modal hard ceiling / idle reclamation -``YETO_HARNESS_TB2_NETWORK_POLICY`` path of a JSON network policy that - replaces the built-in one (see - ``NetworkPolicy``); sandbox egress is closed - for every task the policy does not grant -``YETO_HARNESS_TB2_FAULT`` fault injection for the R-TB coverage matrix, - comma separated ``kind:selector[=value]``: - ``create_fail:`` acquire raises - (-> INFRASTRUCTURE / ABORTED); - ``deadline:=`` lease deadline of - ``s`` seconds (wall-clock cancel -> timeout, - destroy, env_live back to 0); - ``max_turns:=`` worker gets - ``SECRLENV_MAX_TURNS=n`` (turn budget). - ```` is the 0-based acquire ordinal in - this process or an exact trajectory id. -""" - -from __future__ import annotations - -import asyncio -import base64 -import io -import os -import re -import secrets -import shlex -import shutil -import signal -import subprocess -import sys -import tarfile -import tempfile -import threading -import time -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Callable, Protocol - -try: - import tomllib -except ImportError: # pragma: no cover - Python < 3.11 - import tomli as tomllib # type: ignore[no-redef] - -from .client import EpisodeAPIError -from .environment import serve_environment - -TASKS_DIR_ENV = "YETO_HARNESS_TB2_TASKS_DIR" -LEASE_SECONDS_ENV = "YETO_HARNESS_TB2_LEASE_SECONDS" -FAULT_ENV = "YETO_HARNESS_TB2_FAULT" -MODAL_APP_ENV = "YETO_HARNESS_TB2_MODAL_APP" -SANDBOX_TTL_ENV = "YETO_HARNESS_TB2_SANDBOX_TTL_S" -IDLE_TIMEOUT_ENV = "YETO_HARNESS_TB2_IDLE_TIMEOUT_S" -ENV_HOST_ENV = "YETO_HARNESS_TB2_ENV_HOST" -NETWORK_POLICY_ENV = "YETO_HARNESS_TB2_NETWORK_POLICY" - -DEFAULT_MODAL_APP = "yeto-tbench2" -DEFAULT_SANDBOX_TTL_S = 1800 -DEFAULT_IDLE_TIMEOUT_S = 600 -DEFAULT_WORKDIR = "/app" -TESTS_PATH = "/tests" -VERIFIER_LOGS_PATH = "/logs/verifier" -MAX_VERIFIER_OUTPUT_BYTES = 64 * 1024 -FAULT_KINDS = ("create_fail", "deadline", "max_turns") -TURN_BUDGET_ENV = "SECRLENV_MAX_TURNS" - - -# ----------------------------------------------------------------------------- task - - -@dataclass(frozen=True) -class Tb2Task: - task_id: str - task_dir: Path - docker_image: str - cpus: int - memory_mb: int - workdir: str - agent_timeout_s: float - verifier_timeout_s: float - # Task statement from ``instruction.md`` (the Terminal-Bench prompt). The - # smoke datasets carry only ``task_id`` and a generic system message, so - # this is the only place the model can learn what to do (S15 stage 2: all - # 48 trajectories ran without it and scored 0). - instruction: str | None = None - - @property - def tests_dir(self) -> Path: - return self.task_dir / "tests" - - @property - def solution_dir(self) -> Path: - return self.task_dir / "solution" - - -_TASK_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") - - -def resolve_task(task_ref: str, tasks_dir: Path) -> Tb2Task: - """``task_ref`` is the task directory name (``metadata.task_id``); fail closed otherwise.""" - if not isinstance(task_ref, str) or not _TASK_ID_RE.match(task_ref): - raise ValueError(f"invalid TB2 task reference {task_ref!r}") - task_dir = Path(tasks_dir) / task_ref - toml_path = task_dir / "task.toml" - if not toml_path.is_file(): - raise FileNotFoundError(f"{task_ref}: no task.toml under {tasks_dir}") - config = tomllib.loads(toml_path.read_text()) - env_cfg = config.get("environment", {}) - image = env_cfg.get("docker_image") - if not image: - raise ValueError(f"{task_ref}: task.toml has no [environment].docker_image") - if not (task_dir / "tests" / "test.sh").is_file(): - raise FileNotFoundError(f"{task_ref}: tests/test.sh is missing") - return Tb2Task( - task_id=task_ref, - task_dir=task_dir, - docker_image=str(image), - cpus=max(1, int(env_cfg.get("cpus", 1))), - memory_mb=max(2048, int(env_cfg.get("memory_mb", 2048))), - workdir=_dockerfile_workdir(task_dir / "environment" / "Dockerfile"), - agent_timeout_s=float(config.get("agent", {}).get("timeout_sec", 900.0)), - verifier_timeout_s=float(config.get("verifier", {}).get("timeout_sec", 900.0)), - instruction=_read_instruction(task_dir / "instruction.md"), - ) - - -def tb2_instructions_env() -> dict[str, str]: - """Worker env selecting the signed TB2 Codex system prompt.""" - from . import codex_harness_agent as harness - - return { - harness.INSTRUCTIONS_FAMILY_ENV: harness.TB2_INSTRUCTIONS_FAMILY, - harness.TB2_INSTRUCTIONS_SHA_ENV: harness.TB2_BASE_INSTRUCTIONS_SHA256, - } - - -class TaskPromptPreflightError(ValueError): - """A dataset row would reach Codex without a real task statement.""" - - -def preflight_task_prompts(data_path: Path, tasks_dir: Path) -> list[tuple[str, str]]: - """Resolve the first user message of every dataset row exactly as the - subprocess agent will (``task_prompt``) and refuse rows without a task - statement or whose prompt is a stringified chat list (S15 stage 2 root - cause). Returns ``[(task_id, prompt), ...]`` in row order.""" - import json - from types import SimpleNamespace - - from .codex_openenv_subprocess_agent_function import TaskPromptMissing, task_prompt - - out: list[tuple[str, str]] = [] - for lineno, line in enumerate(Path(data_path).read_text().splitlines(), 1): - if not line.strip(): - continue - row = json.loads(line) - metadata = dict(row.get("metadata") or {}) - task_id = metadata.get("task_id") - try: - task = resolve_task(task_id, Path(tasks_dir)) - prompt = task_prompt(metadata, row.get("prompt"), SimpleNamespace(task=task)) - except (TaskPromptMissing, ValueError, FileNotFoundError) as exc: - raise TaskPromptPreflightError(f"{data_path}:{lineno} task {task_id!r}: {exc}") from exc - head = prompt.lstrip()[:2] - if head in ("[{", "{'", '{"') or "'role':" in prompt or '"role":' in prompt: - raise TaskPromptPreflightError(f"{data_path}:{lineno} task {task_id!r}: prompt is a stringified chat message") - out.append((str(task_id), prompt)) - if not out: - raise TaskPromptPreflightError(f"{data_path}: no dataset rows") - return out - - -def _read_instruction(path: Path) -> str | None: - if not path.is_file(): - return None - text = path.read_text(errors="replace").strip() - return text or None - - -def _dockerfile_workdir(dockerfile: Path) -> str: - workdir = DEFAULT_WORKDIR - if dockerfile.is_file(): - for line in dockerfile.read_text(errors="replace").splitlines(): - m = re.match(r"^\s*WORKDIR\s+(\S+)", line) - if m: - workdir = m.group(1).strip('"') - return workdir - - -def tests_tar_b64(tests_dir: Path) -> str: - buf = io.BytesIO() - with tarfile.open(fileobj=buf, mode="w:gz") as tar: - for path in sorted(tests_dir.rglob("*")): - if path.is_file() and "__pycache__" not in path.parts: - tar.add(path, arcname=str(path.relative_to(tests_dir))) - return base64.b64encode(buf.getvalue()).decode() - - -# ----------------------------------------------------------------------------- network policy - - -@dataclass(frozen=True) -class NetworkGrant: - """Outbound network access of one sandbox. - - The default grant is closed (no egress). ``open`` allows all egress. - Otherwise only ``domains`` (``*.`` wildcard prefixes allowed) and - ``cidrs`` are reachable. - """ - - open: bool = False - domains: tuple[str, ...] = () - cidrs: tuple[str, ...] = () - - @property - def closed(self) -> bool: - return not self.open and not self.domains and not self.cidrs - - @classmethod - def parse(cls, value: Any, *, where: str) -> "NetworkGrant": - if value == "open": - return cls(open=True) - if value in ("closed", None): - return cls() - if not isinstance(value, dict) or set(value) - {"domains", "cidrs"}: - raise ValueError(f"{where}: grant must be 'open', 'closed' or {{domains, cidrs}}, got {value!r}") - domains = tuple(value.get("domains") or ()) - cidrs = tuple(value.get("cidrs") or ()) - for item in domains + cidrs: - if not isinstance(item, str) or not item or item == "*": - raise ValueError(f"{where}: bad allowlist entry {item!r} (use 'open' for all egress)") - import ipaddress - - for cidr in cidrs: - ipaddress.ip_network(cidr, strict=False) - return cls(domains=domains, cidrs=cidrs) - - def modal_kwargs(self) -> dict[str, Any]: - """``modal.Sandbox.create`` network arguments (modal >= 1.5.5).""" - if self.open: - return {} - if self.closed: - return {"block_network": True} - kwargs: dict[str, Any] = {} - if self.cidrs: - kwargs["outbound_cidr_allowlist"] = list(self.cidrs) - if self.domains: - kwargs["outbound_domain_allowlist"] = list(self.domains) - return kwargs - - -CLOSED = NetworkGrant() - - -@dataclass(frozen=True) -class NetworkPolicy: - """Per-task network grants; every task not listed gets ``default`` (closed). - - JSON form (``YETO_HARNESS_TB2_NETWORK_POLICY``):: - - {"default": "closed", - "tasks": {"pip-task": {"domains": ["pypi.org", "files.pythonhosted.org"]}, - "web-task": "open"}} - """ - - tasks: dict[str, NetworkGrant] = field(default_factory=dict) - default: NetworkGrant = CLOSED - - def grant(self, task_id: str) -> NetworkGrant: - return self.tasks.get(task_id, self.default) - - @classmethod - def from_mapping(cls, data: Any, *, where: str = "network policy") -> "NetworkPolicy": - if not isinstance(data, dict) or set(data) - {"default", "tasks"}: - raise ValueError(f"{where}: expected {{'default', 'tasks'}}") - tasks = data.get("tasks", {}) - if not isinstance(tasks, dict): - raise ValueError(f"{where}: 'tasks' must be an object") - return cls( - tasks={str(k): NetworkGrant.parse(v, where=f"{where} task {k}") for k, v in tasks.items()}, - default=NetworkGrant.parse(data.get("default", "closed"), where=f"{where} default"), - ) - - @classmethod - def builtin(cls) -> "NetworkPolicy": - from .tb2_network_grants import TB2_OPEN_EGRESS_TASKS - - return cls(tasks={t: NetworkGrant(open=True) for t in TB2_OPEN_EGRESS_TASKS}) - - @classmethod - def from_env(cls) -> "NetworkPolicy": - path = os.environ.get(NETWORK_POLICY_ENV) - if not path: - return cls.builtin() - import json - - return cls.from_mapping(json.loads(Path(path).read_text()), where=path) - - -# ----------------------------------------------------------------------------- backend protocol - - -@dataclass(frozen=True) -class ExecResult: - exit_code: int - output: str - timed_out: bool = False - - -class SandboxHandle(Protocol): - """One running sandbox. ``exec`` is synchronous (run via ``asyncio.to_thread``).""" - - workdir: str - - def exec(self, command: str, *, timeout_s: float, workdir: str | None = None) -> ExecResult: ... - - def terminate(self) -> None: ... - - def alive(self) -> bool: ... - - -class SandboxBackend(Protocol): - def create(self, task: Tb2Task, trajectory_id: str) -> SandboxHandle: ... - - -def shell_with_limits(command: str, *, timeout_s: float, output_bytes: int) -> str: - """Wrap an agent command: wall clock via coreutils ``timeout``, output capped. - - Exit code 124 (``timeout``) / 137 (KILL after grace) mean the clock expired; - stdout and stderr are merged and capped at ``output_bytes + 1`` so the relay - can tell truncation from an exact fit without ever reading more. - """ - seconds = max(1, int(timeout_s)) - return ( - f"timeout -k 2 {seconds} bash -c {shlex.quote(command)} 2>&1 | head -c {int(output_bytes) + 1}; " - "exit ${PIPESTATUS[0]}" - ) - - -_TIMEOUT_EXIT_CODES = frozenset({124, 137}) - - -# ----------------------------------------------------------------------------- local backend - - -# Only these parent variables reach a local sandbox command (never tokens). -LOCAL_ENV_PASSTHROUGH = ("PATH", "LANG", "LC_ALL", "TZ") -DEFAULT_PATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" - - -def minimal_env(parent: dict[str, str] | None = None) -> dict[str, str]: - """The base environment of a local sandbox command. - - The parent environment is not inherited: it holds Modal / HF / W&B - tokens and the reward HMAC key. Only ``LOCAL_ENV_PASSTHROUGH`` is - copied; ``PATH`` falls back to a standard value. - """ - parent = os.environ if parent is None else parent - env = {k: parent[k] for k in LOCAL_ENV_PASSTHROUGH if parent.get(k)} - env.setdefault("PATH", DEFAULT_PATH) - env.setdefault("LANG", "C.UTF-8") - return env - - -class LocalProcessSandbox: - """Commands run on this host in a scratch root; ``/tests`` and ``/logs`` are - redirected into that root through ``TB2_TESTS_DIR`` / ``TB2_VERIFIER_LOGS_DIR``. - - Commands get ``minimal_env`` only. The network grant is recorded but NOT - enforced: this backend is for CPU tests and dev on a trusted host. - """ - - def __init__(self, root: Path, workdir: str, network: NetworkGrant = CLOSED) -> None: - self.root = root - self.workdir = workdir - self.network = network - (root / workdir.lstrip("/")).mkdir(parents=True, exist_ok=True) - self._alive = True - self._processes: set[subprocess.Popen] = set() - self._lock = threading.Lock() - - def host_path(self, sandbox_path: str) -> Path: - return self.root / sandbox_path.lstrip("/") - - def exec(self, command: str, *, timeout_s: float, workdir: str | None = None) -> ExecResult: - if not self._alive: - raise RuntimeError("sandbox is gone") - cwd = self.host_path(workdir or self.workdir) - cwd.mkdir(parents=True, exist_ok=True) - env = { - **minimal_env(), - "TB2_TESTS_DIR": str(self.host_path(TESTS_PATH)), - "TB2_VERIFIER_LOGS_DIR": str(self.host_path(VERIFIER_LOGS_PATH)), - "HOME": str(self.root), - } - proc = subprocess.Popen( - ["bash", "-lc", command], cwd=cwd, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, - start_new_session=True, - ) - with self._lock: - self._processes.add(proc) - try: - try: - out, _ = proc.communicate(timeout=timeout_s + 5) - except subprocess.TimeoutExpired: - self._kill(proc) - out, _ = proc.communicate() - return ExecResult(exit_code=124, output=(out or b"").decode("utf-8", "replace"), timed_out=True) - finally: - with self._lock: - self._processes.discard(proc) - output = (out or b"").decode("utf-8", "replace") - return ExecResult(exit_code=proc.returncode, output=output, timed_out=proc.returncode in _TIMEOUT_EXIT_CODES) - - @staticmethod - def _kill(proc: subprocess.Popen) -> None: - try: - os.killpg(proc.pid, signal.SIGKILL) - except ProcessLookupError: - pass - - def terminate(self) -> None: - """Like a sandbox kill: every command still running dies with the sandbox.""" - self._alive = False - with self._lock: - running = list(self._processes) - for proc in running: - self._kill(proc) - shutil.rmtree(self.root, ignore_errors=True) - - def alive(self) -> bool: - return self._alive and self.root.exists() - - -class LocalProcessBackend: - """CPU backend: a scratch directory per sandbox, commands run as subprocesses. - - ``setup`` (optional) seeds the workdir from the task (tests use it to put - the fixture the verifier checks into place). - """ - - def __init__(self, base_dir: Path | None = None, *, setup: Callable[[Tb2Task, LocalProcessSandbox], None] | None = None, - network_policy: NetworkPolicy | None = None) -> None: - self.base_dir = Path(base_dir) if base_dir else Path(tempfile.gettempdir()) / "yeto-tb2-local" - self.network_policy = network_policy or NetworkPolicy() - self.setup = setup - self.created: list[LocalProcessSandbox] = [] - - def create(self, task: Tb2Task, trajectory_id: str) -> LocalProcessSandbox: - self.base_dir.mkdir(parents=True, exist_ok=True) - root = Path(tempfile.mkdtemp(prefix=f"{task.task_id}-", dir=self.base_dir)) - sandbox = LocalProcessSandbox(root, task.workdir, self.network_policy.grant(task.task_id)) - if self.setup is not None: - self.setup(task, sandbox) - self.created.append(sandbox) - return sandbox - - -# ----------------------------------------------------------------------------- Modal backend - - -def _decode_output(data: bytes | str | None) -> str: - if not data: - return "" - return data if isinstance(data, str) else data.decode("utf-8", errors="replace") - - -class ModalSandbox: - def __init__(self, sandbox: Any, workdir: str) -> None: - self._sandbox = sandbox - self.workdir = workdir - self.object_id = getattr(sandbox, "object_id", "") - - def exec(self, command: str, *, timeout_s: float, workdir: str | None = None) -> ExecResult: - # Bytes, decoded leniently: command output cut by the byte limit (or plain - # binary output) can end mid UTF-8 sequence, and Modal's text mode then - # raises UnicodeDecodeError inside the relay (S17 M1, 5 tool calls failed). - process = self._sandbox.exec( - "bash", "-lc", command, workdir=workdir or self.workdir, timeout=int(timeout_s) + 30, text=False - ) - output = _decode_output(process.stdout.read()) - err = _decode_output(process.stderr.read()) - code = process.wait() - return ExecResult(exit_code=int(code), output=output + err, timed_out=code in _TIMEOUT_EXIT_CODES) - - def terminate(self) -> None: - try: - self._sandbox.terminate() - finally: - detach = getattr(self._sandbox, "detach", None) - if callable(detach): - detach() - - def alive(self) -> bool: - try: - return self._sandbox.poll() is None - except Exception: # noqa: BLE001 - a sandbox the API no longer knows is gone - return False - - -SANDBOX_MODAL_TOKEN_ID_ENV = "YETO_SANDBOX_MODAL_TOKEN_ID" -SANDBOX_MODAL_TOKEN_SECRET_ENV = "YETO_SANDBOX_MODAL_TOKEN_SECRET" - - -def sandbox_modal_client(environ=None) -> Any: - """secret-handling-hardening: a Modal client built from the sandbox-only - token (YETO_SANDBOX_MODAL_TOKEN_*), or None when it is not set (then - Modal's default credentials apply, e.g. on a developer machine).""" - environ = os.environ if environ is None else environ - token_id = environ.get(SANDBOX_MODAL_TOKEN_ID_ENV) - token_secret = environ.get(SANDBOX_MODAL_TOKEN_SECRET_ENV) - if not (token_id and token_secret): - return None - from modal import Client - - return Client.from_credentials(token_id, token_secret) - - -class ModalSandboxBackend: - """Modal Sandbox from the task's official image. All ``modal`` imports are lazy.""" - - def __init__( - self, - *, - app_name: str = DEFAULT_MODAL_APP, - ttl_s: int = DEFAULT_SANDBOX_TTL_S, - idle_timeout_s: int = DEFAULT_IDLE_TIMEOUT_S, - run_id: str | None = None, - network_policy: NetworkPolicy | None = None, - ) -> None: - self.app_name = app_name - # Closed for every task unless a policy grants it (task 5.1). - self.network_policy = network_policy or NetworkPolicy() - self.ttl_s = int(ttl_s) - self.idle_timeout_s = int(idle_timeout_s) - self.run_id = run_id or os.environ.get("OPENENV_RUN_ID") or "" - self._app: Any = None - self._client: Any = None - self._lock = threading.Lock() - - def _client_kwargs(self) -> dict[str, Any]: - """``{"client": c}`` for the sandbox-only token, ``{}`` otherwise.""" - with self._lock: - if self._client is None: - self._client = sandbox_modal_client() or False - return {"client": self._client} if self._client else {} - - def _get_app(self) -> Any: - kwargs = self._client_kwargs() - with self._lock: - if self._app is None: - from modal import App - - self._app = App.lookup(self.app_name, create_if_missing=True, **kwargs) - return self._app - - def create(self, task: Tb2Task, trajectory_id: str) -> ModalSandbox: - from modal import Image, Sandbox - - tags = {"yeto-tb2-task": task.task_id, "yeto-trajectory": trajectory_id[:63]} - if self.run_id: - tags["yeto-run-id"] = self.run_id[:63] - sandbox = Sandbox.create( - "sleep", - "infinity", - app=self._get_app(), - image=Image.from_registry(task.docker_image), - timeout=self.ttl_s, - idle_timeout=self.idle_timeout_s, - cpu=float(task.cpus), - memory=int(task.memory_mb), - workdir=task.workdir, - tags=tags, - **self._client_kwargs(), - **self.network_policy.grant(task.task_id).modal_kwargs(), - ) - return ModalSandbox(sandbox, task.workdir) - - -# ----------------------------------------------------------------------------- relay environment + verifier - - -class EpisodeBinding: - """The one episode a sandbox serves: bound by the first request on the wire. - - The agent function mints the episode id after ``acquire`` (the worker job - carries it), so the lease cannot know it up front; it binds to the first id - it sees and refuses every other one (one sandbox == one episode). - """ - - def __init__(self) -> None: - self.episode_id: str | None = None - - def bind(self, episode_id: str) -> None: - if self.episode_id is None: - self.episode_id = episode_id - elif episode_id != self.episode_id: - raise EpisodeAPIError(404, "unknown_episode", "unknown episode") - - -class RelayTerminalEnvironment: - """``TerminalEnvironment`` served in the trusted process; commands go to the sandbox.""" - - def __init__(self, handle: SandboxHandle, binding: EpisodeBinding | None = None) -> None: - self._handle = handle - self.binding = binding or EpisodeBinding() - self.commands: list[str] = [] - self.submissions: list[dict[str, Any]] = [] - self.submitted = False - - def _check_episode(self, episode_id: str) -> None: - self.binding.bind(episode_id) - if self.submitted: - raise EpisodeAPIError(409, "episode_terminal", "episode already submitted") - - async def execute(self, episode_id: str, command: str, timeout_seconds: float, output_bytes: int) -> dict[str, Any]: - self._check_episode(episode_id) - self.commands.append(command) - result = await asyncio.to_thread( - self._handle.exec, - shell_with_limits(command, timeout_s=timeout_seconds, output_bytes=output_bytes), - timeout_s=timeout_seconds, - ) - truncated = len(result.output.encode("utf-8")) > output_bytes - return { - "exit_code": result.exit_code, - "output": result.output[:output_bytes], - "timed_out": result.timed_out, - "truncated": truncated, - } - - async def submit(self, episode_id: str, submission: dict[str, Any]) -> dict[str, Any]: - self._check_episode(episode_id) - self.submissions.append(dict(submission)) - self.submitted = True - return {"accepted": True} - - -# S17 G3: Modal rejects an exec whose argv exceeds 64 KiB (ARG_MAX), and 7 of -# the 89 TB2 tasks ship tests/ larger than that once tar+base64'd. Small tests -# stay inline (command unchanged); larger ones are first appended to a staging -# file in the sandbox in chunks (``verifier_stage_commands``), and the verifier -# command decodes that file instead. -MAX_INLINE_TESTS_B64 = 48_000 -TESTS_STAGE_CHUNK = 48_000 -TESTS_STAGE_PATH = "/tmp/yeto-tb2-tests.tgz.b64" - - -def verifier_stage_commands(task: Tb2Task) -> list[str]: - """Commands to run (in order) before ``verifier_command``; empty when inline.""" - b64 = tests_tar_b64(task.tests_dir) - if len(b64) <= MAX_INLINE_TESTS_B64: - return [] - commands = [f": > {TESTS_STAGE_PATH}"] - commands += [f"printf %s {b64[i:i + TESTS_STAGE_CHUNK]} >> {TESTS_STAGE_PATH}" - for i in range(0, len(b64), TESTS_STAGE_CHUNK)] - return commands - - -def verifier_command(task: Tb2Task) -> str: - """Stage ``tests/`` at ``/tests``, run ``test.sh`` in the workdir, print the reward. - - The reward line is tagged so it survives arbitrary test output in front of it. - Large ``tests/`` must be staged first with ``verifier_stage_commands``. - """ - b64 = tests_tar_b64(task.tests_dir) - source = (f"echo {b64} | base64 -d" if len(b64) <= MAX_INLINE_TESTS_B64 - else f"base64 -d {TESTS_STAGE_PATH}") - return ( - f"T=${{TB2_TESTS_DIR:-{TESTS_PATH}}}; L=${{TB2_VERIFIER_LOGS_DIR:-{VERIFIER_LOGS_PATH}}}; " - f"rm -rf \"$T\" && mkdir -p \"$T\" \"$L\" && {source} | tar xzf - -C \"$T\" && " - f"bash \"$T/test.sh\"; rc=$?; echo; echo \"YETO_TB2_TESTSH_RC=$rc\"; " - f"echo \"YETO_TB2_REWARD=$(cat \"$L/reward.txt\" 2>/dev/null | tr -d '[:space:]')\"" - ) - - -def stage_tests(handle: "SandboxHandle", task: Tb2Task) -> None: - """Run ``verifier_stage_commands`` in ``handle``; raise if any chunk fails.""" - for command in verifier_stage_commands(task): - result = handle.exec(command, timeout_s=120) - if result.exit_code != 0: - raise RuntimeError(f"staging tests/ failed (rc={result.exit_code}): {result.output[-200:]}") - - -_REWARD_RE = re.compile(r"YETO_TB2_REWARD=(\S*)") -_RC_RE = re.compile(r"YETO_TB2_TESTSH_RC=(\d+)") - - -VERIFIER_LOG_CHARS = 2000 - - -def verifier_log_excerpt(output: str, limit: int = VERIFIER_LOG_CHARS) -> str: - """Last ``limit`` chars of the verifier output (pytest prints its verdict last).""" - text = output if isinstance(output, str) else "" - return text if len(text) <= limit else "…" + text[-(limit - 1):] - - -class Tb2Verifier: - """Trusted-side verifier: ``tests/test.sh`` inside the sandbox, reward from ``reward.txt``.""" - - def __init__(self, handle: SandboxHandle, task: Tb2Task, binding: EpisodeBinding | None = None) -> None: - self._handle = handle - self._task = task - self.binding = binding or EpisodeBinding() - self.evaluations = 0 - self.last_output = "" - - async def evaluate(self, episode_id: str) -> dict[str, Any]: - bound = self.binding.episode_id - if bound is not None and episode_id != bound: - raise ValueError(f"verifier bound to {bound!r}, asked for {episode_id!r}") - self.evaluations += 1 - await asyncio.to_thread(stage_tests, self._handle, self._task) - result = await asyncio.to_thread( - self._handle.exec, verifier_command(self._task), timeout_s=self._task.verifier_timeout_s - ) - self.last_output = result.output[-MAX_VERIFIER_OUTPUT_BYTES:] - reward = _REWARD_RE.search(result.output) - rc = _RC_RE.search(result.output) - testsh_rc = int(rc.group(1)) if rc else result.exit_code - return { - "passed": bool(reward) and reward.group(1) == "1", - "testsh_rc": testsh_rc, - "timed_out": result.timed_out, - # Observe only (S15 verifier probe): the tail of test.sh output - # (pytest summary / install errors); never part of the signed outcome. - "log": verifier_log_excerpt(result.output), - } - - -# ----------------------------------------------------------------------------- faults - - -@dataclass(frozen=True) -class FaultSpec: - kind: str - selector: str - value: str | None = None - - def matches(self, ordinal: int, trajectory_id: str) -> bool: - return self.selector == str(ordinal) or self.selector == trajectory_id - - -def parse_faults(spec: str | None) -> tuple[FaultSpec, ...]: - faults: list[FaultSpec] = [] - for raw in (spec or "").split(","): - item = raw.strip() - if not item: - continue - kind, sep, rest = item.partition(":") - if not sep or kind not in FAULT_KINDS: - raise ValueError(f"{FAULT_ENV}: invalid fault {item!r} (kinds: {', '.join(FAULT_KINDS)})") - selector, _sep, value = rest.partition("=") - if not selector: - raise ValueError(f"{FAULT_ENV}: fault {item!r} needs a selector (acquire ordinal or trajectory id)") - faults.append(FaultSpec(kind, selector, value or None)) - return tuple(faults) - - -class InjectedCreateFailure(RuntimeError): - """Raised by ``acquire`` when ``create_fail`` selects the trajectory.""" - - injected_fault = True # not a provider outage (subprocess agent fail-fast counter) - - -# ----------------------------------------------------------------------------- provider - - -@dataclass -class Tb2Lease: - """Returned by ``acquire``: field names match ``subprocess_agent.EnvironmentLease``.""" - - env_url: str - env_token: str - verifier: Tb2Verifier - destroy: Callable[[], Any] - describe: Callable[[], Any] - deadline_seconds: float - worker_env: dict[str, str] | None = None - task: Tb2Task | None = None - environment: RelayTerminalEnvironment | None = None - handle: Any = None - - -class Tb2EnvironmentProvider: - def __init__( - self, - backend: SandboxBackend, - tasks_dir: Path, - *, - lease_seconds: float | None = None, - faults: tuple[FaultSpec, ...] = (), - host: str = "127.0.0.1", - ) -> None: - self.backend = backend - self.tasks_dir = Path(tasks_dir) - self.lease_seconds = lease_seconds - self.faults = tuple(faults) - self.host = host - self.acquired = 0 - self.destroyed = 0 - self.live: dict[str, Tb2Lease] = {} - self.fault_log: list[tuple[str, str, int]] = [] - - def _faults_for(self, ordinal: int, trajectory_id: str) -> dict[str, str | None]: - hits = {f.kind: f.value for f in self.faults if f.matches(ordinal, trajectory_id)} - for kind in hits: - self.fault_log.append((kind, trajectory_id, ordinal)) - return hits - - async def acquire(self, task_id: str, trajectory_id: str) -> Tb2Lease: - ordinal = self.acquired - self.acquired += 1 - task = resolve_task(task_id, self.tasks_dir) - faults = self._faults_for(ordinal, trajectory_id) - if "create_fail" in faults: - raise InjectedCreateFailure(f"injected sandbox creation failure for {trajectory_id!r} (acquire #{ordinal})") - handle = await asyncio.to_thread(self.backend.create, task, trajectory_id) - token = secrets.token_urlsafe(24) - binding = EpisodeBinding() - relay = RelayTerminalEnvironment(handle, binding) - try: - runner, url = await serve_environment(relay, token, host=self.host) - except BaseException: - await asyncio.to_thread(handle.terminate) - raise - deadline = self.lease_seconds if self.lease_seconds is not None else task.agent_timeout_s - # 2026-10-07: TB2 episodes run under the signed TB2 system prompt. - worker_env: dict[str, str] | None = tb2_instructions_env() - if "deadline" in faults: - deadline = float(faults["deadline"] or 1.0) - if "max_turns" in faults: - worker_env = {**(worker_env or {}), TURN_BUDGET_ENV: str(int(faults["max_turns"] or 2))} - state = {"destroyed": False} - - async def destroy() -> None: - if state["destroyed"]: - return - state["destroyed"] = True - try: - await runner.cleanup() - finally: - self.live.pop(trajectory_id, None) - self.destroyed += 1 - await asyncio.to_thread(handle.terminate) - - async def describe() -> str: - if not state["destroyed"]: - return "live" - # Modal reports the exit asynchronously; a short bounded wait keeps - # ``describe`` honest without stalling the rollout. - for _ in range(20): - if not await asyncio.to_thread(handle.alive): - return "gone" - await asyncio.sleep(0.25) - return "live" - - lease = Tb2Lease( - env_url=url, - env_token=token, - verifier=Tb2Verifier(handle, task, binding), - destroy=destroy, - describe=describe, - deadline_seconds=float(deadline), - worker_env=worker_env, - task=task, - environment=relay, - handle=handle, - ) - self.live[trajectory_id] = lease - return lease - - -# ----------------------------------------------------------------------------- factories (module:callable) - - -def _env_float(name: str, default: float | None) -> float | None: - raw = os.environ.get(name) - if raw is None or not raw.strip(): - return default - value = float(raw) - if value <= 0: - raise ValueError(f"{name} must be positive") - return value - - -def _tasks_dir_from_env() -> Path: - raw = os.environ.get(TASKS_DIR_ENV) - if not raw: - raise RuntimeError(f"{TASKS_DIR_ENV} is not set (terminal-bench-2 checkout root)") - tasks_dir = Path(raw).expanduser().resolve() - if not tasks_dir.is_dir(): - raise RuntimeError(f"{TASKS_DIR_ENV}={raw}: not a directory") - return tasks_dir - - -def _provider(backend: SandboxBackend) -> Tb2EnvironmentProvider: - return Tb2EnvironmentProvider( - backend, - _tasks_dir_from_env(), - lease_seconds=_env_float(LEASE_SECONDS_ENV, None), - faults=parse_faults(os.environ.get(FAULT_ENV)), - host=os.environ.get(ENV_HOST_ENV, "127.0.0.1"), - ) - - -def require_modal_client() -> None: - """The Modal client must import in *this* interpreter (the island's python). - - The ports image does not ship it; Modal mounts its own ``modal`` package - into Function containers without its dependencies, so a bare ``import - modal`` there fails only at first Sandbox use, inside every rollout - worker (A-T3-5, codex-smoke-20261003-7). The launcher installs the client - in the island setup (``MODAL_CLIENT_SETUP``); this check makes the driver - preflight fail closed before any GPU work when that did not happen. - """ - try: - import modal # noqa: F401 - except Exception as exc: # noqa: BLE001 - ImportError or Modal's own re-raise - raise RuntimeError( - f"the Modal Sandbox backend needs an importable `modal` client in {sys.executable}: {exc}" - ) from exc - - -def modal_provider(miles_args: Any = None) -> Tb2EnvironmentProvider: - """``YETO_HARNESS_ENVIRONMENT_PROVIDER=yeto.rl.harness.codex.tb2_provider:modal_provider``.""" - del miles_args - require_modal_client() - backend = ModalSandboxBackend( - app_name=os.environ.get(MODAL_APP_ENV, DEFAULT_MODAL_APP), - ttl_s=int(_env_float(SANDBOX_TTL_ENV, DEFAULT_SANDBOX_TTL_S) or DEFAULT_SANDBOX_TTL_S), - idle_timeout_s=int(_env_float(IDLE_TIMEOUT_ENV, DEFAULT_IDLE_TIMEOUT_S) or DEFAULT_IDLE_TIMEOUT_S), - network_policy=NetworkPolicy.from_env(), - ) - return _provider(backend) - - -def local_provider(miles_args: Any = None) -> Tb2EnvironmentProvider: - """CPU-only provider (commands run on this host): smoke tests and dev.""" - del miles_args - return _provider(LocalProcessBackend(network_policy=NetworkPolicy.from_env())) diff --git a/yeto/rl/harness/codex/tbench_reward.py b/yeto/rl/harness/codex/tbench_reward.py deleted file mode 100644 index a8df347b5..000000000 --- a/yeto/rl/harness/codex/tbench_reward.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Reward and group filter for Terminal-Bench Codex samples on the ports path. - -- ``reward_func``: only structurally valid ``tbench_outcome`` rewards (first of - the three validation points, design D7). An infrastructure marker - marks the sample ABORTED and returns the numeric 0.0 Miles requires; the - group filter never lets it into training. A verified outcome also sets - ``sample.metadata["success"]`` to the outcome ``passed`` bit (all - Terminal-Bench tests passed); aborted samples get no flag. -- ``check_group``: validates every sample again, drops groups with aborted - members, and enforces R-D5a: samples of one trajectory (siblings) share one - reward and are counted once (``rollout_id``/``trajectory_id``). -""" - -from __future__ import annotations - -from typing import Any - -from yeto.rl.rewards.types import RewardResult, Trajectory -from yeto.rl.tbench_outcome import OUTCOME_KEY, UntrustedTBenchOutcome, verified_outcome - -INFRASTRUCTURE_KEY = "tbench_infrastructure_error" - - -def _metadata(sample: Any) -> dict[str, Any]: - metadata = getattr(sample, "metadata", None) - if not isinstance(metadata, dict): - raise UntrustedTBenchOutcome("sample has no metadata") - return metadata - - -def _mark_aborted(sample: Any) -> None: - try: - from miles.utils.types import Sample - - sample.status = Sample.Status.ABORTED - except ImportError: # CPU tests without the Miles package - sample.status = "ABORTED" - - -def is_infrastructure(sample: Any) -> bool: - metadata = _metadata(sample) - return INFRASTRUCTURE_KEY in metadata and OUTCOME_KEY not in metadata - - -def tbench_reward(trajectory: Trajectory) -> RewardResult: - """Neutral form (decoupling 3.3). - - Writes the full-success flag into ``trajectory.metadata`` in place (the - outcome pass bit, not reward > 0). An infrastructure marker is no - verdict: ``success`` is removed and the result is ``aborted`` so the - backend adapter marks the sample with its abort status. - """ - - metadata = trajectory.metadata - if not isinstance(metadata, dict): - raise UntrustedTBenchOutcome("sample has no metadata") - if INFRASTRUCTURE_KEY in metadata and OUTCOME_KEY not in metadata: - metadata.pop("success", None) # no verdict: never a positive - return RewardResult(value=0.0, aborted=True, reason=INFRASTRUCTURE_KEY) - outcome, value = verified_outcome(metadata) - metadata["success"] = outcome["passed"] is True - return RewardResult(value=value) - - -def _sample_reward(sample: Any) -> float: - result = tbench_reward(Trajectory(metadata=_metadata(sample))) - if result.aborted: - _mark_aborted(sample) - return result.value - - -async def reward_func(args: Any, samples: Any, **_kwargs: Any) -> float | list[float]: - del args - if isinstance(samples, list): - return [_sample_reward(sample) for sample in samples] - return _sample_reward(samples) - - -def iter_samples(group: list[Any]) -> list[Any]: - flat: list[Any] = [] - for item in group: - flat.extend(item if isinstance(item, list) else [item]) - return flat - - -def sibling_key(sample: Any) -> Any: - rollout_id = getattr(sample, "rollout_id", None) - if rollout_id is not None: - return ("rollout", rollout_id) - trajectory_id = _metadata(sample).get("trajectory_id") - if trajectory_id is not None: - return ("trajectory", trajectory_id) - return ("row", id(sample)) - - -def check_siblings(samples: list[Any]) -> dict[Any, float]: - """R-D5a: siblings share one reward; returns one reward per trajectory.""" - per_trajectory: dict[Any, float] = {} - outcomes: dict[Any, dict[str, Any]] = {} - for sample in samples: - outcome, value = verified_outcome(_metadata(sample)) - key = sibling_key(sample) - if key in per_trajectory: - if per_trajectory[key] != value or outcomes[key] != outcome: - raise UntrustedTBenchOutcome( - f"sibling samples of {key!r} carry different outcomes" - ) - else: - per_trajectory[key] = value - outcomes[key] = outcome - return per_trajectory - - -def check_group(args: Any, samples: list[Any], **_kwargs: Any): - """Dynamic group filter: validate, drop aborted groups, enforce sibling sharing.""" - try: - from miles.rollout.filter_hub.base_types import DynamicFilterOutput - except ImportError: # CPU tests without the Miles package - from types import SimpleNamespace as DynamicFilterOutput # type: ignore[assignment] - flat = iter_samples(samples) - if any(is_infrastructure(sample) for sample in flat): - for sample in flat: - if is_infrastructure(sample): - _mark_aborted(sample) - return DynamicFilterOutput(keep=False, reason="group_has_aborted") - rewards = check_siblings(flat) # raises on tampering / mismatch - distinct = set(rewards.values()) - if len(distinct) < 2: - return DynamicFilterOutput(keep=False, reason=f"zero_std_{next(iter(distinct)) if distinct else 'empty'}") - return DynamicFilterOutput(keep=True, reason=None) diff --git a/yeto/rl/harness/gateway/proxy.py b/yeto/rl/harness/gateway/proxy.py index ac6e7a342..d6a086040 100644 --- a/yeto/rl/harness/gateway/proxy.py +++ b/yeto/rl/harness/gateway/proxy.py @@ -17,8 +17,10 @@ from __future__ import annotations +import copy import json import secrets +import time from contextvars import ContextVar from typing import Any @@ -144,13 +146,66 @@ def _error(status: int, exc: Exception) -> Any: def _responses_sse(request: dict[str, Any], response: dict[str, Any], number: int) -> bytes: - # Reuse the Codex bridge's event sequence (created, item added/done, completed). - from yeto.rl.harness.codex.codex_harness_agent import _sse_response - - return _sse_response(request={"model": request.get("model"), "tools": request.get("tools") or [], - "prompt_cache_key": request.get("prompt_cache_key")}, - output=response["output"], usage=response["usage"], - response_number=number) + """Responses API event stream: created, per-item added/(arguments)/done, completed. + + One event per output item, no token streaming. ``instructions`` echoes the + request (the client owns its system prompt). + """ + output = response["output"] + completed = { + "id": f"resp_yeto_{number}", + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "error": None, + "incomplete_details": None, + "instructions": request.get("instructions"), + "max_output_tokens": None, + "max_tool_calls": None, + "model": request.get("model"), + "output": output, + "parallel_tool_calls": False, + "previous_response_id": None, + "prompt_cache_key": request.get("prompt_cache_key"), + "reasoning": {"effort": "xhigh", "summary": "none"}, + "safety_identifier": None, + "service_tier": "default", + "store": False, + "temperature": None, + "text": {"format": {"type": "text"}}, + "tool_choice": "auto", + "tools": request.get("tools") or [], + "top_logprobs": 0, + "top_p": None, + "truncation": "disabled", + "usage": response["usage"], + "user": None, + "metadata": {}, + } + events: list[dict[str, Any]] = [ + {"type": "response.created", "response": {**completed, "status": "in_progress", "output": []}}] + for index, item in enumerate(output): + added = copy.deepcopy(item) + added["status"] = "in_progress" + if item["type"] == "function_call": + added["arguments"] = "" + events.append({"type": "response.output_item.added", "output_index": index, "item": added}) + if item["type"] == "function_call": + events.append({"type": "response.function_call_arguments.delta", "item_id": item["id"], + "output_index": index, "delta": item["arguments"]}) + events.append({"type": "response.function_call_arguments.done", "item_id": item["id"], + "output_index": index, "arguments": item["arguments"]}) + events.append({"type": "response.output_item.done", "output_index": index, "item": item}) + events.append({"type": "response.completed", "response": completed}) + chunks = [f"event: {event['type']}\n".encode() + b"data: " + _canonical_json(event) + b"\n\n" + for event in events] + chunks.append(b"data: [DONE]\n\n") + return b"".join(chunks) + + +def _canonical_json(value: Any) -> bytes: + return json.dumps(value, ensure_ascii=False, allow_nan=False, sort_keys=True, + separators=(",", ":")).encode("utf-8") def _chat_sse(completion: dict[str, Any]) -> bytes: diff --git a/yeto/rl/harness/harbor/README.md b/yeto/rl/harness/harbor/README.md index 8c7195e97..5dc706b7e 100644 --- a/yeto/rl/harness/harbor/README.md +++ b/yeto/rl/harness/harbor/README.md @@ -18,8 +18,8 @@ backend: not verified. | verl backend: `HarborAgentLoop(AgentLoopBase)`, registered as `yeto_harbor` | yeto (verl unchanged) | `yeto/rl/adapters/verl/harbor_agent_loop.py` | | Token recorder in front of the sampler (verl) | yeto, N10.5 (not written) | port: `yeto/rl/agentic/recorder.py` | -The self-written Codex harness (`yeto/rl/harness/codex/`) is not deleted (user -decision pending). New code does not import it. +The self-written Codex harness (`yeto/rl/harness/codex/`) was removed in S26 +E0.5; the last commit with it is tagged `s26-pre-harness-removal`. ## Outcome rule diff --git a/yeto/rl/harness/harbor/__init__.py b/yeto/rl/harness/harbor/__init__.py index f77cb09a2..3310a2265 100644 --- a/yeto/rl/harness/harbor/__init__.py +++ b/yeto/rl/harness/harbor/__init__.py @@ -1,8 +1,8 @@ """Harbor (harbor-framework/harbor) as the agent harness of yeto agentic RL (S25 N10.1). -Replaces layers L1/L2/L4/L8/L11 of the self-written Codex harness -(``yeto.rl.harness.codex``): task loading, sandbox, tool exec, verifier, -agent build. New code depends on this package, not on ``harness.codex``. +Harbor owns task loading, sandbox, tool exec, verifier and agent build. The +self-written Codex harness that did this before was removed in S26 E0.5 (git +tag ``s26-pre-harness-removal``). """ from .pins import CODEX_VERSION, HARBOR_COMMIT, HARBOR_VERSION, install_spec, installed_problem diff --git a/yeto/rl/harness/reward_env/__init__.py b/yeto/rl/harness/reward_env/__init__.py index debae3ff5..07f9bfb54 100644 --- a/yeto/rl/harness/reward_env/__init__.py +++ b/yeto/rl/harness/reward_env/__init__.py @@ -1,5 +1,6 @@ """Benchmark-neutral reward environments for agentic RL (change rl-agentic-reward-env). -``benchmark``: adapter protocol, registry, judge runner. ``tb2``: Terminal-Bench 2 -adapter, prebaked Modal sandboxes, held-out split helpers. +``benchmark``: adapter protocol, registry, judge runner, held-out split helpers. +``swebench_verified``: SWE-bench Verified adapter. The Terminal-Bench 2 adapter +left with the self-written Codex harness (S26 E0.5). """ diff --git a/yeto/rl/harness/reward_env/benchmark.py b/yeto/rl/harness/reward_env/benchmark.py index 7044ddf25..2af974dd3 100644 --- a/yeto/rl/harness/reward_env/benchmark.py +++ b/yeto/rl/harness/reward_env/benchmark.py @@ -96,7 +96,7 @@ def to_dict(self) -> dict[str, Any]: class JudgeHandle(Protocol): - """What the judge runner needs from a sandbox (``tb2_provider.SandboxHandle`` fits).""" + """What the judge runner needs from a sandbox.""" def exec(self, command: str, *, timeout_s: float, workdir: str | None = None) -> Any: ... @@ -124,7 +124,7 @@ def parse_judge(self, task_id: str, output: str, exit_code: int, timed_out: bool _REGISTRY: dict[str, Callable[..., BenchmarkAdapter]] = {} # benchmark name -> module of this package that registers it (lazy, no import cost) -_BUILTIN_MODULES = {"tb2": "tb2", "swebench-verified": "swebench_verified"} +_BUILTIN_MODULES = {"swebench-verified": "swebench_verified"} def register(name: str, factory: Callable[..., BenchmarkAdapter]) -> None: diff --git a/yeto/rl/harness/reward_env/tb2.py b/yeto/rl/harness/reward_env/tb2.py deleted file mode 100644 index 5fe2d2c73..000000000 --- a/yeto/rl/harness/reward_env/tb2.py +++ /dev/null @@ -1,252 +0,0 @@ -"""Terminal-Bench 2 adapter for the reward environment (change rl-agentic-reward-env). - -Reuses the task resolution and the judge command of the existing codex TB2 -provider (``yeto.rl.harness.codex.tb2_provider``; imported, never modified) and -adds what that provider lacks: - -- ``prebake_plan``: the dependency setup that each task's ``tests/test.sh`` - performs on every judge run (``apt-get install curl``, the uv installer, - ``uvx -p 3.13 -w pytest==… …`` or ``pip install pytest…``), extracted so it can - run once at image build time. ``test.sh`` itself is never changed: at judge - time it runs the same steps again but they hit the baked caches. -- the Modal app names/env knobs for the prebaked sandboxes; the Modal backend - itself (``PrebakedModalSandboxBackend`` + ``modal_provider``) lives in the - cloud layer ``yeto/cloud/modal_reward_env.py`` and is injected through the - neutral ``tb2_provider.SandboxBackend`` / ``YETO_HARNESS_ENVIRONMENT_PROVIDER``. -- held-out evaluation split helpers (30 tasks by default). -""" - -from __future__ import annotations - -import json -import os -import re -import shlex -from pathlib import Path -from typing import Any - -from yeto.rl.harness.codex import tb2_provider as tb2 - -from .benchmark import PrebakePlan, TaskSpec, register - -NAME = "tb2" -DEFAULT_APP = "yeto-reward-env" -PROTECTED_APPS = frozenset({tb2.DEFAULT_MODAL_APP}) # "yeto-tbench2": production, never built into -APP_ENV = "YETO_REWARD_ENV_MODAL_APP" -PREBAKE_ENV = "YETO_REWARD_ENV_PREBAKE" # "0" disables prebaking (A/B timing) - -_INSTALL_PREFIXES = ("apt-get ", "apt ", "curl ", "wget ", "source ", ". ", "export ", "pip install ", - "pip3 install ", "python -m pip install ", "python3 -m pip install ", "uv ") -_CONTROL = re.compile(r"^(if|then|else|elif|fi|echo|exit|set)\b") - - -def _logical_lines(script: str) -> list[str]: - lines: list[str] = [] - buffer = "" - for raw in script.splitlines(): - stripped = raw.strip() - if not buffer and (not stripped or stripped.startswith("#")): - continue - if stripped.endswith("\\"): - buffer += stripped[:-1] + " " - continue - lines.append((buffer + stripped).strip()) - buffer = "" - if buffer.strip(): - lines.append(buffer.strip()) - return lines - - -_UVX_VALUE_FLAGS = frozenset({"-p", "--python", "-w", "--with", "--from", "--index", "--default-index", - "--index-url", "--extra-index-url", "-i", "--with-requirements", - "--with-editable", "--index-strategy", "--prerelease"}) - - -def _uvx_warmup(line: str) -> str | None: - """``uvx -p 3.13 -w a -w b pytest --ctrf …`` -> ``uvx -p 3.13 -w a -w b pytest --version``.""" - try: - tokens = shlex.split(line) - except ValueError: - return None - if not tokens or tokens[0] != "uvx": - return None - i = 1 - while i < len(tokens): - token = tokens[i] - if token in _UVX_VALUE_FLAGS: - i += 2 - continue - if token.startswith("-"): - i += 1 - continue - return shlex.join(tokens[: i + 1] + ["--version"]) - return None - - -def _is_test_run(line: str) -> bool: - return line.startswith(("uvx ", "uv run ")) or bool(re.match(r"^(python3? -m )?pytest\b", line)) - - -def prebake_from_test_sh(script: str, base_image: str) -> PrebakePlan: - """Setup lines before the test run, verbatim, plus a cache warm-up of the uvx tool env.""" - commands: list[str] = [] - skipped: list[str] = [] - for line in _logical_lines(script): - if _is_test_run(line): - # shlex re-quoting would freeze shell variables (e.g. ${COMMIT_HASH} - # set by a task-specific line): no warm-up then, judge installs it. - warm = None if "$" in line else _uvx_warmup(line) - if warm: - commands.append(warm) - elif line.startswith("uvx "): - skipped.append(line) # tool env not warmed - break - if line.startswith(("curl ", "wget ")) and not re.search(r"\|\s*(ba)?sh\b", line): - skipped.append(line) # a data download into the workdir, not a dependency - elif line.startswith(_INSTALL_PREFIXES): - commands.append(line) - elif _CONTROL.match(line) or line.startswith("[") or line.startswith("cd "): - continue # PWD guards etc.: not dependencies - else: - skipped.append(line) - else: - # no recognisable test-run line: do not guess - return PrebakePlan(base_image=base_image, commands=(), complete=False, skipped=tuple(skipped) or ("",)) - return PrebakePlan(base_image=base_image, commands=tuple(commands), complete=not skipped, skipped=tuple(skipped)) - - -def prebake_script(plan: PrebakePlan) -> str: - """One bash script (env from ``source`` must survive into the warm-up).""" - return "set -e\ncd /tmp\n" + "\n".join(plan.commands) + "\n" - - -def _checkout_version(path: Path) -> str: - import subprocess - - try: - rev = subprocess.run(["git", "-C", str(path), "rev-parse", "--short=7", "HEAD"], capture_output=True, - text=True, timeout=10, check=True).stdout.strip() - except (OSError, subprocess.SubprocessError): - return "tb2@unknown" - return f"tb2@{rev}" if rev else "tb2@unknown" - - -def _official_difficulty(task_dir: Path) -> str: - try: - meta = tb2.tomllib.loads((task_dir / "task.toml").read_text()).get("metadata", {}) - except (OSError, ValueError): - return "unknown" - value = meta.get("difficulty") - return str(value) if value else "unknown" - - -class Tb2Benchmark: - name = NAME - - def __init__(self, tasks_dir: str | os.PathLike[str] | None = None, version: str | None = None) -> None: - raw = tasks_dir or os.environ.get(tb2.TASKS_DIR_ENV) - if not raw: - raise RuntimeError(f"{tb2.TASKS_DIR_ENV} is not set (terminal-bench-2 checkout root)") - self.tasks_dir = Path(raw).expanduser().resolve() - self.version = version or _checkout_version(self.tasks_dir) - - def task_ids(self) -> list[str]: - return sorted(p.parent.name for p in self.tasks_dir.glob("*/task.toml")) - - def _task(self, task_id: str) -> tb2.Tb2Task: - return tb2.resolve_task(task_id, self.tasks_dir) - - def task_spec(self, task_id: str) -> TaskSpec: - t = self._task(task_id) - difficulty = _official_difficulty(t.task_dir) - return TaskSpec(benchmark=NAME, task_id=t.task_id, base_image=t.docker_image, cpus=t.cpus, - memory_mb=t.memory_mb, workdir=t.workdir, agent_timeout_s=t.agent_timeout_s, - judge_timeout_s=t.verifier_timeout_s, instruction=t.instruction, - benchmark_version=self.version, difficulty=difficulty, - difficulty_source="tb2-task.toml", - eval_bucket=f"tb2-{difficulty}" if difficulty in ("easy", "medium", "hard") else None) - - def prebake_plan(self, task_id: str) -> PrebakePlan: - t = self._task(task_id) - return prebake_from_test_sh((t.tests_dir / "test.sh").read_text(), t.docker_image) - - def judge_setup_commands(self, task_id: str) -> list[str]: - """Commands run before ``judge_command`` (stage large tests/, S17 G3).""" - return tb2.verifier_stage_commands(self._task(task_id)) - - def judge_command(self, task_id: str, submission: str | None = None) -> str: - del submission # TB2 judges the sandbox the agent worked in (official stage-at-verify) - return tb2.verifier_command(self._task(task_id)) - - def parse_judge(self, task_id: str, output: str, exit_code: int, timed_out: bool) -> dict[str, Any]: - reward = tb2._REWARD_RE.search(output) - rc = tb2._RC_RE.search(output) - passed = bool(reward) and reward.group(1) == "1" - return {"reward": 1.0 if passed else 0.0, "passed": passed, - "testsh_rc": int(rc.group(1)) if rc else exit_code, - "reward_txt": reward.group(1) if reward else None} - - -register(NAME, Tb2Benchmark) - - -# ----------------------------------------------------------------------------- Modal backend -# The Modal sandbox backend for this adapter lives in the cloud layer -# (``yeto/cloud/modal_reward_env.py``; import boundary: neutral core never imports -# a cloud library). It is injected at launch through the neutral provider -# interface ``YETO_HARNESS_ENVIRONMENT_PROVIDER=yeto.cloud.modal_reward_env:modal_provider`` -# (a ``tb2_provider.SandboxBackend`` behind a ``Tb2EnvironmentProvider``). - - -def check_build_app(app_name: str) -> str: - if app_name in PROTECTED_APPS: - raise ValueError(f"refusing to build into production Modal app {app_name!r}") - return app_name - - -# ----------------------------------------------------------------------------- hold-out (WP3 D2/D6b) - -HOLDOUT_SEED = 20261008 -HOLDOUT_QUOTAS = {"tb2-easy": 2, "tb2-medium": 18, "tb2-hard": 10} -HOLDOUT_RULE_UNUSABLE = ("exclude smoke6 (S15 trained) and data/eval/tb2-unusable.json; " - "stratified by difficulty: easy 2 / medium 18 / hard 10") -HOLDOUT_RULE = "exclude smoke6 (S15 trained); stratified by difficulty: easy 2 / medium 18 / hard 10" -# WP3 rl-eval-difficulty-buckets D2 (#123 c8979bd2): the S15 smoke tasks -# (codex-bundle/data/tbench2_smoke6.jsonl) were trained on in both S15 GPU runs, -# so they leave the eval pool before the stratified draw (they may stay in training). -SMOKE6_TASK_IDS = ("fix-git", "regex-log", "sqlite-db-truncate", "log-summary-date-ranges", - "openssl-selfsigned-cert", "git-multibranch") -SMOKE6_REASON = "S15 smoke training" - - -def smoke6_exclusions() -> dict[str, str]: - return {t: SMOKE6_REASON for t in SMOKE6_TASK_IDS} - - -UNUSABLE_SCHEMA = "yeto.tb2-unusable/v1" -# data/eval/tb2-unusable.json: tasks whose official solution does not score 1 in -# our judge env (S17 G3 positive control). They leave training *and* the eval pool. -UNUSABLE_DEFAULT_PATH = Path(__file__).resolve().parents[4] / "data" / "eval" / "tb2-unusable.json" - - -def load_unusable(path: str | os.PathLike[str] | None = None) -> dict[str, str]: - """``task_id -> "unusable/: "`` from the unusable table (fail closed).""" - data = json.loads(Path(path or UNUSABLE_DEFAULT_PATH).read_text()) - if data.get("schema") != UNUSABLE_SCHEMA or data.get("benchmark") != "tb2": - raise ValueError(f"{path or UNUSABLE_DEFAULT_PATH}: not a {UNUSABLE_SCHEMA} tb2 file") - out: dict[str, str] = {} - for item in data["items"]: - tid, cat, why = item["task_id"], item["category"], item["reason"] - if not (tid and cat and why) or tid in out: - raise ValueError(f"bad or duplicate unusable entry: {item!r}") - out[tid] = f"unusable/{cat}: {why}" - return out - - -def build_holdout(adapter: "Tb2Benchmark", *, exclude: Any = None, **kwargs: Any) -> dict[str, Any]: - """TB2 hold-out; ``exclude`` defaults to the smoke-6 tasks (pass ``{}`` for none).""" - from .benchmark import build_holdout as _build - - kwargs.setdefault("rule", HOLDOUT_RULE) - return _build(adapter, HOLDOUT_QUOTAS, seed=HOLDOUT_SEED, - exclude=smoke6_exclusions() if exclude is None else exclude, **kwargs) diff --git a/yeto/rl/harness/session_protocol.md b/yeto/rl/harness/session_protocol.md index fc417f850..17f3e9cda 100644 --- a/yeto/rl/harness/session_protocol.md +++ b/yeto/rl/harness/session_protocol.md @@ -1,5 +1,7 @@ # 会话服务协议（decoupling 任务 3.8，spec `rl-neutral-rewards-harness`） +> S26 E0.5: the Codex harness this file describes was removed (tag `s26-pre-harness-removal`); the protocol stays as the contract of the test double. + codex harness 对"会话服务"的全部依赖。Miles 会话服务是这个协议的一种实现；以后换框架（如 verl）时，只要提供同样的五条路径，harness 不用改。测试替身：`tests/session_service_double.py`。 本协议只写文档与测试替身，**没有**实现 verl 侧。 diff --git a/yeto/rl/island_tasks.py b/yeto/rl/island_tasks.py index 89e485749..524979fe1 100644 --- a/yeto/rl/island_tasks.py +++ b/yeto/rl/island_tasks.py @@ -14,7 +14,7 @@ A token ``-`` names the tasks ```` ... ```` unless the token itself is a task id. A row's task id is -``metadata.task_id`` (what ``yeto.rl.rlf.catalog.prompt_rows`` writes) or a +``metadata.task_id`` or a top-level ``task_id``. This module is backend-neutral: it knows JSONL rows and task ids only. diff --git a/yeto/rl/rlf/__init__.py b/yeto/rl/rlf/__init__.py deleted file mode 100644 index 680abf802..000000000 --- a/yeto/rl/rlf/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -"""Small, trusted-side bridge for rlenvfactory CTF tasks.""" -from .contracts import ( - AdmissionError, TaskAdmission, admit_task_package, TrustedGrade, - EpisodeLease, RewardEvidence, finish_episode, rl_gate, -) -from .catalog import TaskCatalog, import_catalog, prompt_rows, write_prompt_data, task_statement, rlf_identity -from .preflight import RolloutRejected, validate_rollout, trusted_grade_from_evaluation -from .trajectory import CodexTrajectory, from_codex_result -__all__ = ["prompt_rows", "write_prompt_data", "task_statement", "rlf_identity", "trusted_grade_from_evaluation","AdmissionError", "TaskAdmission", "admit_task_package", "TrustedGrade", "EpisodeLease", "RewardEvidence", "finish_episode", "rl_gate", "TaskCatalog", "import_catalog", "RolloutRejected", "validate_rollout", "CodexTrajectory", "from_codex_result"] -from .modal import RlfModalVmConfig, RlfVmLease, acquire_vm diff --git a/yeto/rl/rlf/catalog.py b/yeto/rl/rlf/catalog.py deleted file mode 100644 index 705f9aec4..000000000 --- a/yeto/rl/rlf/catalog.py +++ /dev/null @@ -1,116 +0,0 @@ -from __future__ import annotations -import json, uuid -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any -from .contracts import AdmissionError, EpisodeLease, TaskAdmission, admit_task_package - -@dataclass -class TaskCatalog: - tasks: dict[str, TaskAdmission] - roots: dict[str, Path] = field(default_factory=dict) - def add(self, package: str|Path, **kwargs: Any) -> TaskAdmission: - item = admit_task_package(package, **kwargs) - if item.task_id in self.tasks: raise ValueError(f'duplicate RLF task: {item.task_id}') - self.tasks[item.task_id] = item - self.roots[item.task_id] = Path(package) - return item - def get(self, task_id: str) -> TaskAdmission: - return self.tasks[task_id] - def episode(self, task_id: str, *, trajectory_id: str|None=None) -> EpisodeLease: - """Start one rollout episode for an admitted task (RLF identity producer).""" - admission = self.get(task_id) - return EpisodeLease(task_id=admission.task_id, - trajectory_id=trajectory_id or f'rlf-traj-{uuid.uuid4().hex}') - -def import_catalog(packages: list[str|Path], **kwargs: Any) -> TaskCatalog: - catalog=TaskCatalog({}) - for package in packages: catalog.add(package, **kwargs) - return catalog - -RLF_IDENTITY_KEYS = ('rlf_task_id', 'rlf_episode_id', 'rlf_trajectory_id', 'policy_id') - -def rlf_identity(lease: EpisodeLease, *, policy_id: str|None=None) -> dict[str, str]: - """The RLF identity quadruple for one episode lease (prompt/chain metadata).""" - identity = {'rlf_task_id': lease.task_id, 'rlf_episode_id': lease.episode_id, - 'rlf_trajectory_id': lease.trajectory_id} - if policy_id is not None: identity['policy_id'] = str(policy_id) - return identity - -def task_statement(catalog: TaskCatalog, task_id: str) -> str: - """The task's own statement (public README); never fabricated.""" - root = catalog.roots.get(task_id) - if root is None: raise AdmissionError(f'no package root recorded for task {task_id!r}') - for name in ('README.md', 'README'): - candidate = root/'public'/name - if candidate.is_file(): - text = candidate.read_text().strip() - if text: return text - raise AdmissionError(f'task {task_id!r} has no public README task statement') - -def target_endpoint(catalog: TaskCatalog, task_id: str) -> str | None: - """The agent-visible target address from the package's task.yaml. - - RLF CTF targets run on the episode VM's host network, so the address is - always 127.0.0.1:. Returns None when the package declares no - port (non-network task). - """ - root = catalog.roots.get(task_id) - if root is None: - return None - text = (root/'task.yaml').read_text() - for line in text.splitlines(): - if line.startswith('target_port:'): - port = line.split(':', 1)[1].strip() - return f'127.0.0.1:{port}' if port else None - return None - - -def prompt_rows(catalog: TaskCatalog, task_ids: list[str]|None=None, *, repeat: int=1) -> list[dict[str, Any]]: - """Catalog -> Miles ``--prompt-data`` rows (the catalog consumer for rollout). - - Each row carries ``metadata.task_id`` (read by the harness ``task_identity``) - and ``rlf_task_id``; the per-rollout episode/trajectory ids are produced - later by ``codex_openenv_generate.produce_rlf_identity`` so episodes are - never shared across rollouts. - """ - rows: list[dict[str, Any]] = [] - for task_id in (task_ids or list(catalog.tasks)): - admission = catalog.get(task_id) - statement = task_statement(catalog, task_id) - endpoint = target_endpoint(catalog, task_id) - if endpoint: - statement += ('\n\nThe target service listens on ' - + endpoint + ' (reachable from your shell).') - - metadata: dict[str, Any] = {'task_id': admission.task_id, 'rlf_task_id': admission.task_id, - 'benchmark': 'rlf', 'grader_mode': admission.grader_mode, - 'package_digest': admission.package_digest, 'prompt': statement} - if admission.image_digest: metadata['image_digest'] = admission.image_digest - for _ in range(max(1, repeat)): - rows.append({'metadata': dict(metadata), 'prompt': statement}) - return rows - -def write_prompt_data(catalog: TaskCatalog, path: str|Path, task_ids: list[str]|None=None, *, repeat: int=1) -> Path: - """Write ``prompt_rows`` as the JSONL file Miles ``--prompt-data`` expects.""" - path = Path(path) - rows = prompt_rows(catalog, task_ids, repeat=repeat) - if not rows: raise AdmissionError('catalog produced no prompt rows') - path.write_text(''.join(json.dumps(row, ensure_ascii=False)+'\n' for row in rows)) - return path - -def main(argv: list[str]|None=None) -> None: - """``python -m yeto.rl.rlf.catalog PKG... --out prompt.jsonl [--repeat N]``.""" - import argparse - parser = argparse.ArgumentParser(description='RLF catalog -> Miles prompt data') - parser.add_argument('packages', nargs='+') - parser.add_argument('--out', required=True) - parser.add_argument('--repeat', type=int, default=1) - parser.add_argument('--evidence', default=None, help='JSON admission evidence (must include runtime_quality_gate)') - args = parser.parse_args(argv) - evidence = json.loads(args.evidence) if args.evidence else None - catalog = import_catalog(args.packages, evidence=evidence) - write_prompt_data(catalog, args.out, repeat=args.repeat) - -if __name__ == '__main__': - main() diff --git a/yeto/rl/rlf/contracts.py b/yeto/rl/rlf/contracts.py deleted file mode 100644 index 00e0ae695..000000000 --- a/yeto/rl/rlf/contracts.py +++ /dev/null @@ -1,98 +0,0 @@ -from __future__ import annotations -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Mapping -import hashlib, json, time, uuid - -class AdmissionError(ValueError): pass - -def _sha(path: Path) -> str: - h=hashlib.sha256() - if path.is_file(): h.update(path.read_bytes()) - else: - for p in sorted(path.rglob('*')): - if p.is_file(): h.update(str(p.relative_to(path)).encode()); h.update(p.read_bytes()) - return h.hexdigest() - -@dataclass(frozen=True) -class TaskAdmission: - task_id: str; grader_mode: str; package_digest: str; image_digest: str|None - source_commit: str|None; evidence: Mapping[str,Any]; manifest: Mapping[str,Any] - -def admit_task_package(root: str|Path, *, evidence: Mapping[str,Any]|None=None, image_digest: str|None=None, source_commit: str|None=None) -> TaskAdmission: - root=Path(root); public=root/'public'; yaml=root/'task.yaml'; manifest=root/'MANIFEST.json' - if not public.is_dir() or not yaml.is_file() or not manifest.is_file(): raise AdmissionError('package must contain public/, task.yaml and MANIFEST.json') - try: m=json.loads(manifest.read_text()); text=yaml.read_text() - except Exception as e: raise AdmissionError(f'invalid package metadata: {e}') from e - task_id=m.get('task_id') or next((x.split(':',1)[1].strip() for x in text.splitlines() if x.startswith('task_id:')),None) - mode=m.get('grader_mode') or next((x.split(':',1)[1].strip() for x in text.splitlines() if x.startswith('grader_mode:')),None) - if not task_id or mode not in {'flag','hacked'}: raise AdmissionError('task_id and grader_mode flag/hacked are required') - ev=dict(evidence or {}) - if not ev.get('runtime_quality_gate'): raise AdmissionError('runtime_quality_gate evidence is required') - if ev.get('agentflow_node_output'): raise AdmissionError('AgentFlow node output is not admission evidence') - return TaskAdmission(task_id, mode, _sha(root), image_digest, source_commit, ev, m) - -@dataclass -class EpisodeLease: - task_id: str; trajectory_id: str; episode_id: str=field(default_factory=lambda: str(uuid.uuid4())); started_at: float=field(default_factory=time.time); closed: bool=False - secret: str|None=None - def teardown(self): self.closed=True; self.secret=None - -@dataclass(frozen=True) -class TrustedGrade: - status: str; reward: float; passed: bool; grader_evidence: Mapping[str,Any]; error_class: str|None=None -@dataclass(frozen=True) -class RewardEvidence: - value: float; metadata: Mapping[str,Any] - -def finish_episode(lease: EpisodeLease, grade: TrustedGrade, *, timed_out: bool=False) -> RewardEvidence: - if lease.closed: raise RuntimeError('episode already torn down') - status='timeout' if timed_out else grade.status - md={'success': bool(grade.passed), 'reward': float(grade.reward), 'rlf_episode_id':lease.episode_id, 'rlf_task_id':lease.task_id, 'rlf_trajectory_id':lease.trajectory_id, 'rlf_status':status, 'grader_evidence':dict(grade.grader_evidence)} - if grade.error_class: md['rejected_data_reason']=grade.error_class - lease.teardown(); return RewardEvidence(float(grade.reward), md) - -def rl_gate(rows: list[Mapping[str,Any]]) -> dict[str,Any]: - rewards=[float(r.get('reward',0)) for r in rows]; adv=[float(r.get('advantage',0)) for r in rows] - total=sum(rewards); mean=total/len(rewards) if rewards else 0.0 - variance=sum((x-mean)**2 for x in rewards)/len(rewards) if rewards else 0.0 - result={'reward_sum':total,'reward_std':variance**0.5,'success_count':sum(bool(r.get('metadata',{}).get('success')) for r in rows),'advantage_abs_sum':sum(abs(x) for x in adv),'policy_update':any(bool(r.get('policy_update')) for r in rows)} - criteria={'non_empty':bool(rows),'reward_std_positive':result['reward_std']>0,'success_count_positive':result['success_count']>0,'advantage_abs_sum_positive':result['advantage_abs_sum']>0,'policy_update':result['policy_update']} - result['passed']=all(criteria.values()) - if not result['passed']: - failed=', '.join(k for k,v in criteria.items() if not v) - raise AdmissionError(f'RL gate failed ({failed}): {result}') - return result - -def group_advantages(rewards: list[float]) -> list[float]: - """GRPO-style group baseline: advantage_i = reward_i - mean(group). - - This is the gate-runner/evidence-entry computation of record for - ``advantage_abs_sum`` on paths where the learner does not export - per-sample advantages; callers must not hand-fill the value. - """ - values=[float(r) for r in rewards] - if not values: return [] - mean=sum(values)/len(values) - return [r-mean for r in values] - -def _params_digest(params: Mapping[str,float]) -> str: - blob=json.dumps({k:float(v) for k,v in sorted(params.items())},sort_keys=True,separators=(',',':')) - return hashlib.sha256(blob.encode()).hexdigest() - -@dataclass(frozen=True) -class PolicyUpdateEvidence: - """Structured proof that an optimizer step changed policy parameters.""" - step: int; updated: bool; before_digest: str; after_digest: str; changed: tuple[str,...] - -def policy_update_evidence(before: Mapping[str,float], after: Mapping[str,float], *, step: int) -> PolicyUpdateEvidence: - """Record a learner update by digesting parameters before and after the step. - - ``updated`` is true only when the parameter snapshot actually changed; - gate rows must read ``policy_update`` from this evidence, never from a - hand-written literal. - """ - keys=set(before)|set(after) - changed=tuple(sorted(k for k in keys if float(before.get(k,0.0))!=float(after.get(k,0.0)))) - bd=_params_digest(before); ad=_params_digest(after) - return PolicyUpdateEvidence(int(step), bd!=ad, bd, ad, changed) diff --git a/yeto/rl/rlf/modal.py b/yeto/rl/rlf/modal.py deleted file mode 100644 index 07124bb36..000000000 --- a/yeto/rl/rlf/modal.py +++ /dev/null @@ -1,579 +0,0 @@ -"""RLF Modal VM lifecycle + CTF environment orchestration (audit breakpoints 3/5). - -Learner GPUs are deliberately never allocated here. - -Two layers, both CPU-testable with a fake VM: - -* VM lifecycle (S20): ``RlfModalVmConfig`` / ``RlfVmLease`` / ``acquire_vm``. -* CTF orchestration (breakpoint 3): the nested-Docker sequence verified in - ``evidence/rlf/modal-ctf-trusted-recvfix2.log`` packaged as functions — - dockerd readiness, ``docker build --network host``, target run with - ``-e FLAG`` (env injection only), health wait, idempotent cleanup. - -Transport decision (代拍板, rationale): the worker-facing wire is served by -``serve_environment`` **in the trusted process** and commands are relayed into -the VM through the Modal exec channel — the exact pattern ``tb2_provider`` -uses — instead of running a token-protected exec HTTP service inside the VM. -Reasons: (1) the bearer token never enters the VM, so a compromised agent -command cannot read it; (2) nothing has to be baked into the VM image beyond -docker; (3) the worker speaks the unchanged ``HttpTerminalEnvironment`` wire. -The agent reaches the CTF target at ``127.0.0.1:31337`` because relay -commands exec in the VM's host network namespace and the target container -runs ``--network host``. - -Flag hygiene: the FLAG reaches the target container only as ``-e FLAG`` on -``docker run``. It is never written into stdout/stderr captures, logs, or -exception messages: every orchestration error passes through :func:`scrub` -with the flag, and :class:`OrchestrationError` never embeds the command line. - -Cleanup order (each step idempotent): solver/grader containers -> target -container -> images are kept (layer cache for the next episode) -> VM -terminate. -""" -from __future__ import annotations - -import base64 -import io -import json -import os -import shlex -import tarfile -import time -import uuid -from dataclasses import dataclass, field, replace -from pathlib import Path -from typing import Any, Callable, Mapping - -from .contracts import EpisodeLease - -# ----------------------------------------------------------------------------- VM lifecycle (S20) - - -@dataclass(frozen=True) -class RlfModalVmConfig: - app_name: str = 'yeto-rlf-env' - cpu: float = 2.0 - memory_mb: int = 4096 - ttl_s: int = 1800 - # Idle kill must not fire below ttl_s: agent turns can spend minutes in model - # generation with no exec running; a shorter idle timeout reclaims the VM - # mid-episode (Modal NotFoundError) and the fail-closed preflight then drops - # the whole group. Default: follow ttl_s (hard cap still bounds leaks). - idle_timeout_s: int | None = None - image: str = 'debian:bookworm-slim' - runtime: str = 'vm' - image_run_commands: tuple = () - - def validate(self): - if self.runtime != 'vm': - raise ValueError('RLF environment must use Modal VM runtime') - if self.cpu <= 0 or self.memory_mb < 512: - raise ValueError('invalid VM resources') - - -class RlfVmLease: - def __init__(self, sandbox: Any, *, task_id: str, episode_id: str, trajectory_id: str | None = None): - self.sandbox = sandbox - self.task_id = task_id - self.episode_id = episode_id - self.closed = False - self.trajectory_id = trajectory_id - - @property - def object_id(self): - return getattr(self.sandbox, 'object_id', '') - - def exec(self, *args, **kwargs): - if self.closed: - raise RuntimeError('RLF VM lease is closed') - return self.sandbox.exec(*args, **kwargs) - - def alive(self) -> bool: - try: - poll = getattr(self.sandbox, 'poll', None) - return bool(poll) and poll() is None - except Exception: # noqa: BLE001 - a sandbox the API no longer knows is gone - return False - - def teardown(self): - if not self.closed: - self.closed = True - self.sandbox.terminate() - - -def acquire_vm(task_id: str, episode_id: str, *, config: RlfModalVmConfig | None = None, image: Any = None, - trajectory_id: str | None = None) -> RlfVmLease: - cfg = config or RlfModalVmConfig() - cfg.validate() - import modal - app = modal.App.lookup(cfg.app_name, create_if_missing=True) - if image is None: - image = modal.Image.from_registry(cfg.image) - if cfg.image_run_commands: - image = image.run_commands(*cfg.image_run_commands) - idle_timeout_s = cfg.idle_timeout_s if cfg.idle_timeout_s is not None else cfg.ttl_s - sb = modal.Sandbox.create('sleep', 'infinity', app=app, image=image, runtime='vm', cpu=cfg.cpu, - memory=cfg.memory_mb, timeout=cfg.ttl_s, idle_timeout=idle_timeout_s, - tags={'yeto-purpose': 'rlf-environment', 'yeto-task': task_id[:63], - 'yeto-episode': episode_id[:63]}) - return RlfVmLease(sb, task_id=task_id, episode_id=episode_id, trajectory_id=trajectory_id) - - -# ----------------------------------------------------------------------------- exec surface (fakeable) - - -@dataclass(frozen=True) -class VmExecResult: - exit_code: int - output: str - timed_out: bool = False - - -# One synchronous VM command: (shell command, timeout) -> captured result. -VmExec = Callable[..., VmExecResult] - -_TIMEOUT_EXIT_CODES = frozenset({124, 137}) - - -def _decode_output(data: bytes | str | None) -> str: - if not data: - return '' - return data if isinstance(data, str) else data.decode('utf-8', errors='replace') - - -def vm_exec(lease: RlfVmLease) -> VmExec: - """Adapt a Modal VM lease to the synchronous ``VmExec`` surface. - - Commands run as ``bash -lc`` in the VM's host network namespace (this is - what lets the agent reach the ``--network host`` CTF target). - """ - - def run(command: str, *, timeout_s: float) -> VmExecResult: - process = lease.exec('bash', '-lc', command, timeout=int(timeout_s) + 30, text=False) - output = _decode_output(process.stdout.read()) + _decode_output(process.stderr.read()) - code = int(process.wait()) - return VmExecResult(code, output, code in _TIMEOUT_EXIT_CODES) - - return run - - -class OrchestrationError(RuntimeError): - """A CTF orchestration step failed. Messages are always flag-scrubbed.""" - - -def scrub(text: Any, *secrets: str | None) -> str: - """Replace every secret occurrence with ``***`` (defense in depth).""" - out = str(text) - for secret in secrets: - if secret: - out = out.replace(str(secret), '***') - return out - - -# ----------------------------------------------------------------------------- CTF task spec - - -@dataclass(frozen=True) -class RlfCtfTask: - """One admitted CTF task package as the orchestrator needs it (trusted side).""" - - task_id: str - package_dir: Path - flag: str - target_port: int = 31337 - grader_mode: str = 'flag' - - -def _task_yaml_field(text: str, name: str) -> str | None: - for line in text.splitlines(): - if line.startswith(name + ':'): - return line.split(':', 1)[1].strip() - return None - - -def resolve_ctf_task(task_id: str, tasks_dir: str | Path, flags: Mapping[str, str]) -> RlfCtfTask: - """Fail closed: unknown task, missing package files, non-flag mode or a - missing trusted flag all refuse the acquire (infrastructure, never a 0 reward).""" - package_dir = Path(tasks_dir) / task_id - yaml_path = package_dir / 'task.yaml' - if not yaml_path.is_file(): - raise OrchestrationError(f'{task_id}: no task.yaml under {tasks_dir}') - if not (package_dir / 'Dockerfile').is_file(): - raise OrchestrationError(f'{task_id}: package has no Dockerfile') - text = yaml_path.read_text() - mode = _task_yaml_field(text, 'grader_mode') or 'flag' - if mode != 'flag': - raise OrchestrationError(f'{task_id}: grader_mode {mode!r} is not packaged (flag only)') - flag = flags.get(task_id) - if not flag: - raise OrchestrationError(f'{task_id}: no trusted flag configured') - port = int(_task_yaml_field(text, 'target_port') or 31337) - return RlfCtfTask(task_id=task_id, package_dir=package_dir, flag=str(flag), target_port=port) - - -# ----------------------------------------------------------------------------- dockerd / build / run / health - -DOCKER_SOCK = 'unix:///run/docker/docker.sock' -_DOCKER = f'docker -H {DOCKER_SOCK}' -REMOTE_TASK_DIR = '/opt/task' -_STAGE_PATH = '/tmp/rlf-task.tgz.b64' -_STAGE_CHUNK = 48_000 -# Only these package entries form the docker build context; ``private/`` -# (reference solution) and admission metadata never enter the VM. -_BUILD_CONTEXT_ENTRIES = ('Dockerfile', 'entrypoint.sh', 'public') - - -def start_dockerd(exec_fn: VmExec, *, attempts: int = 40, sleep_s: float = 1.0, - sleep: Callable[[float], None] = time.sleep) -> None: - """Reset and start nested dockerd, then wait for ``docker info``. - - Parameters verified on Modal (``modal-dockerd-probe.log`` / - ``modal-ctf-trusted-recvfix2.log``): custom socket under /run/docker, - ``--iptables=false --bridge=none`` (the VM has no bridge network). - """ - result = exec_fn( - 'rm -rf /var/lib/docker /run/docker && mkdir -p /run/docker && ' - f'(dockerd --host={DOCKER_SOCK} --iptables=false --bridge=none >/tmp/dockerd.log 2>&1 &)', - timeout_s=30, - ) - if result.exit_code != 0: - raise OrchestrationError(f'dockerd launch failed (rc={result.exit_code}): {result.output[-200:]}') - for _ in range(max(1, attempts)): - if exec_fn(f'{_DOCKER} info >/dev/null 2>&1', timeout_s=15).exit_code == 0: - return - sleep(sleep_s) - raise OrchestrationError(f'dockerd did not become ready after {attempts} attempts') - - -def package_tar_b64(package_dir: str | Path) -> str: - """tar+base64 of the build context only (never ``private/``).""" - root = Path(package_dir) - buf = io.BytesIO() - with tarfile.open(fileobj=buf, mode='w:gz') as tar: - for entry in _BUILD_CONTEXT_ENTRIES: - path = root / entry - if path.is_file(): - tar.add(path, arcname=entry) - elif path.is_dir(): - for item in sorted(path.rglob('*')): - if item.is_file() and '__pycache__' not in item.parts: - tar.add(item, arcname=str(item.relative_to(root))) - return base64.b64encode(buf.getvalue()).decode() - - -def stage_package(exec_fn: VmExec, package_dir: str | Path, *, remote_dir: str = REMOTE_TASK_DIR) -> None: - """Ship the build context into the VM (chunked base64, like the TB2 tests staging).""" - b64 = package_tar_b64(package_dir) - commands = [f'rm -rf {remote_dir} && mkdir -p {remote_dir} && : > {_STAGE_PATH}'] - commands += [f'printf %s {b64[i:i + _STAGE_CHUNK]} >> {_STAGE_PATH}' - for i in range(0, len(b64), _STAGE_CHUNK)] - commands.append(f'base64 -d {_STAGE_PATH} | tar xzf - -C {remote_dir} && rm -f {_STAGE_PATH}') - for command in commands: - result = exec_fn(command, timeout_s=120) - if result.exit_code != 0: - raise OrchestrationError(f'staging the task package failed (rc={result.exit_code}): {result.output[-200:]}') - - -def build_target(exec_fn: VmExec, task: RlfCtfTask, *, image: str, remote_dir: str = REMOTE_TASK_DIR) -> None: - """``docker build --network host`` — the default bridge hangs and - ``--network=none`` fails in nested Docker (verified in recvfix2).""" - result = exec_fn(f'{_DOCKER} build --network host -t {shlex.quote(image)} {shlex.quote(remote_dir)}', - timeout_s=600) - if result.exit_code != 0: - raise OrchestrationError(scrub(f'target image build failed (rc={result.exit_code}): ' - f'{result.output[-200:]}', task.flag)) - - -def run_target(exec_fn: VmExec, task: RlfCtfTask, *, image: str, container: str) -> str: - """Start the target with the FLAG as a container env var (the only place - the flag crosses onto the VM). Output/errors are scrubbed; the flag is - never logged or raised.""" - command = (f'{_DOCKER} run -d --name {shlex.quote(container)} --network host ' - f'-e FLAG={shlex.quote(task.flag)} {shlex.quote(image)}') - result = exec_fn(command, timeout_s=60) - if result.exit_code != 0: - raise OrchestrationError(scrub(f'target container failed to start (rc={result.exit_code}): ' - f'{result.output[-200:]}', task.flag)) - cid = result.output.strip().splitlines()[-1] if result.output.strip() else '' - if not cid: - raise OrchestrationError('target container start returned no container id') - return cid - - -def wait_target_healthy(exec_fn: VmExec, task: RlfCtfTask, *, container: str, attempts: int = 30, - sleep_s: float = 1.0, sleep: Callable[[float], None] = time.sleep) -> None: - """Block until the target's HEALTHCHECK reports healthy; fail fast on exit.""" - probe = (f"{_DOCKER} inspect -f '{{{{.State.Status}}}} " - f"{{{{if .State.Health}}}}{{{{.State.Health.Status}}}}{{{{end}}}}' {shlex.quote(container)}") - for _ in range(max(1, attempts)): - result = exec_fn(probe, timeout_s=15) - fields = result.output.split() - status = fields[0] if fields else '' - health = fields[1] if len(fields) > 1 else '' - if status in ('exited', 'dead'): - raise OrchestrationError(f'target container {container} is {status} before becoming healthy') - if result.exit_code == 0 and (health == 'healthy' or (status == 'running' and not health)): - return - sleep(sleep_s) - raise OrchestrationError(f'target container did not become healthy after {attempts} attempts') - - -def remove_containers(exec_fn: VmExec, containers: list[str] | tuple[str, ...]) -> None: - """``docker rm -f`` in the given order; every step tolerates failure - (idempotent teardown — a missing container is already clean).""" - for name in containers: - try: - exec_fn(f'{_DOCKER} rm -f {shlex.quote(name)} >/dev/null 2>&1', timeout_s=30) - except Exception: # noqa: BLE001 - cleanup never masks the original outcome - pass - - -# ----------------------------------------------------------------------------- orchestration state machine - - -@dataclass -class CtfOrchestration: - """Verified nested-Docker sequence as an idempotent up/down state machine.""" - - exec_fn: VmExec - task: RlfCtfTask - target_container: str = 'rlf-ctf-target' - target_image: str = 'rlf-ctf-target' - # Solver/grader containers (cleaned before the target); empty in the online - # path (the agent relays into the VM host namespace and grading is - # trusted-side), but the ordering contract is kept for reference-check runs. - auxiliary_containers: tuple[str, ...] = () - dockerd_attempts: int = 40 - health_attempts: int = 30 - sleep: Callable[[float], None] = time.sleep - _up_done: bool = False - _down_done: bool = False - - def up(self) -> str: - """dockerd -> stage -> build -> run(FLAG) -> health. Returns the target cid.""" - if self._up_done: - raise OrchestrationError('orchestration already brought up') - try: - start_dockerd(self.exec_fn, attempts=self.dockerd_attempts, sleep=self.sleep) - stage_package(self.exec_fn, self.task.package_dir) - build_target(self.exec_fn, self.task, image=self.target_image) - cid = run_target(self.exec_fn, self.task, image=self.target_image, container=self.target_container) - wait_target_healthy(self.exec_fn, self.task, container=self.target_container, - attempts=self.health_attempts, sleep=self.sleep) - except BaseException: - self.down() # a half-up environment is always cleaned - raise - self._up_done = True - return cid - - def down(self) -> None: - """solver/grader containers -> target container; images stay; idempotent.""" - if self._down_done: - return - self._down_done = True - remove_containers(self.exec_fn, [*self.auxiliary_containers, self.target_container]) - - -# ----------------------------------------------------------------------------- relay environment + verifier - - -class CtfRelayEnvironment: - """``TerminalEnvironment`` served in the trusted process; commands exec on the VM. - - One relay == one episode: the first episode id on the wire binds it - (same contract as ``tb2_provider.EpisodeBinding``). ``submit`` records - the agent's submission for the trusted verifier; the flag comparison - never happens here. - """ - - def __init__(self, exec_fn: VmExec) -> None: - self._exec_fn = exec_fn - self.episode_id: str | None = None - self.commands: list[str] = [] - self.submissions: list[dict[str, Any]] = [] - self.submitted = False - - def _bind(self, episode_id: str) -> None: - if self.episode_id is None: - self.episode_id = episode_id - elif episode_id != self.episode_id: - from yeto.rl.harness.codex.client import EpisodeAPIError - raise EpisodeAPIError(404, 'unknown_episode', 'unknown episode') - - async def execute(self, episode_id: str, command: str, timeout_seconds: float, output_bytes: int): - import asyncio - self._bind(episode_id) - if self.submitted: - from yeto.rl.harness.codex.client import EpisodeAPIError - raise EpisodeAPIError(409, 'episode_terminal', 'episode already submitted') - self.commands.append(command) - seconds = max(1, int(timeout_seconds)) - capped = (f'timeout -k 2 {seconds} bash -c {shlex.quote(command)} 2>&1 | ' - f'head -c {int(output_bytes) + 1}; exit ${{PIPESTATUS[0]}}') - result = await asyncio.to_thread(self._exec_fn, capped, timeout_s=timeout_seconds) - truncated = len(result.output.encode('utf-8')) > output_bytes - return {'exit_code': result.exit_code, 'output': result.output[:output_bytes], - 'timed_out': result.timed_out, 'truncated': truncated} - - async def submit(self, episode_id: str, submission: dict[str, Any]): - self._bind(episode_id) - self.submissions.append(dict(submission)) - self.submitted = True - return {'accepted': True} - - -class RlfFlagVerifier: - """Trusted-side flag grader (``grader_mode: flag``). - - Compares the agent's submitted flag against the trusted flag in constant - time. The evaluation payload carries the RLF grade under ``rlf_grade`` - (consumed by ``finish_trusted`` via ``trusted_grade_from_evaluation``) - with the RLF identity quadruple in ``grader_evidence`` — the verifier is - built per-lease, so that identity is trusted-side. The flag itself never - appears in the payload, the evidence, or the log. - """ - - def __init__(self, relay: CtfRelayEnvironment, *, flag: str, task_id: str, episode_id: str, - trajectory_id: str) -> None: - self._relay = relay - self._flag = flag - self._identity = {'rlf_task_id': task_id, 'rlf_episode_id': episode_id, - 'rlf_trajectory_id': trajectory_id} - self.evaluations = 0 - - async def evaluate(self, episode_id: str) -> dict[str, Any]: - if self._relay.episode_id is not None and episode_id != self._relay.episode_id: - raise ValueError(f'verifier bound to {self._relay.episode_id!r}, asked for {episode_id!r}') - self.evaluations += 1 - import hmac as _hmac - submission = self._relay.submissions[-1] if self._relay.submissions else None - candidate = submission.get('flag') if isinstance(submission, Mapping) else None - submitted = isinstance(candidate, str) and bool(candidate) - matched = submitted and _hmac.compare_digest(candidate, self._flag) - evidence = {'grader': 'rlf-flag', 'grader_mode': 'flag', 'submitted': submitted, **self._identity} - grade: dict[str, Any] = {'status': 'completed', 'reward': 1.0 if matched else 0.0, - 'passed': bool(matched), 'grader_evidence': evidence} - if not submitted: - grade['error_class'] = 'missing_submission' - return {'passed': bool(matched), 'rlf_grade': grade, - 'log': 'flag submission matched' if matched else 'flag submission missing or did not match'} - - -# ----------------------------------------------------------------------------- full environment lease - - -@dataclass(frozen=True) -class RlfCtfConfig: - """CTF packaging for :class:`RlfModalProvider` (breakpoint 3).""" - - tasks_dir: str # directory of admitted task packages (task_id -> package dir) - flags: Mapping[str, str] # task_id -> trusted flag (never leaves the trusted side) - vm_image: str = 'docker:27-dind' # nested Docker runtime (verified on Modal) - lease_seconds: float = 1800.0 - dockerd_attempts: int = 40 - health_attempts: int = 30 - - -@dataclass -class RlfCtfLease: - """Full ``EnvironmentLease``: field names match the subprocess agent protocol.""" - - env_url: str - env_token: str - verifier: RlfFlagVerifier - destroy: Callable[[], Any] - describe: Callable[[], Any] - deadline_seconds: float - worker_env: dict[str, str] | None = None - task_id: str = '' - episode_id: str = '' - trajectory_id: str = '' - vm_lease: RlfVmLease | None = None - orchestration: CtfOrchestration | None = None - environment: CtfRelayEnvironment | None = None - - -class RlfModalProvider: - """Preflight-compatible RLF environment provider factory product. - - ``resolve_environment_provider`` calls ``factory(miles_args)`` and requires - the result to expose ``acquire(task_id, trajectory_id)``; this class is that - product. Each acquire opens one ``EpisodeLease`` (RLF identity producer) - and one isolated Modal VM bound to it. - - With ``ctf=None`` the returned lease is the S20 VM lease (exec/teardown + - RLF identity; no HTTP endpoint/verifier). With an :class:`RlfCtfConfig` - each acquire additionally runs the packaged CTF orchestration - (dockerd/target/flag/health) and returns a full :class:`RlfCtfLease` - (env_url/env_token/verifier/destroy/describe/deadline) — audit - breakpoint 3. - """ - - def __init__(self, *, config: RlfModalVmConfig | None = None, ctf: RlfCtfConfig | None = None): - self.config = config or RlfModalVmConfig() - self.ctf = ctf - - async def acquire(self, task_id: str, trajectory_id: str = '') -> RlfVmLease | RlfCtfLease: - episode = EpisodeLease(task_id=task_id, - trajectory_id=trajectory_id or f'rlf-traj-{uuid.uuid4().hex}') - if self.ctf is None: - return acquire_vm(task_id, episode.episode_id, config=self.config, - trajectory_id=episode.trajectory_id) - return await self._acquire_ctf(episode) - - async def _acquire_ctf(self, episode: EpisodeLease) -> RlfCtfLease: - import asyncio - import secrets - - from yeto.rl.harness.codex.environment import serve_environment - - ctf = self.ctf - assert ctf is not None - task = resolve_ctf_task(episode.task_id, ctf.tasks_dir, ctf.flags) - # vm_exec/relay use ``bash -lc``; the default docker:27-dind image is - # alpine-based and has no bash, so add it at image build time. - vm_config = replace(self.config, image=ctf.vm_image, - image_run_commands=('apk add --no-cache bash',)) - vm = await asyncio.to_thread(acquire_vm, episode.task_id, episode.episode_id, - config=vm_config, trajectory_id=episode.trajectory_id) - orchestration = CtfOrchestration(vm_exec(vm), task, - dockerd_attempts=ctf.dockerd_attempts, - health_attempts=ctf.health_attempts) - try: - await asyncio.to_thread(orchestration.up) - relay = CtfRelayEnvironment(vm_exec(vm)) - token = secrets.token_urlsafe(24) - runner, url = await serve_environment(relay, token, host='127.0.0.1') - except BaseException: - orchestration.down() - await asyncio.to_thread(vm.teardown) - raise - deadline = float(ctf.lease_seconds or self.config.ttl_s) - state = {'destroyed': False} - - async def destroy() -> None: - if state['destroyed']: - return - state['destroyed'] = True - # solver/grader containers -> target container -> images stay -> VM terminate - await asyncio.to_thread(orchestration.down) - await runner.cleanup() - await asyncio.to_thread(vm.teardown) - - async def describe() -> str: - if not state['destroyed']: - return 'live' - for _ in range(20): - if not await asyncio.to_thread(vm.alive): - return 'gone' - await asyncio.sleep(0.25) - return 'live' - - return RlfCtfLease(env_url=url, env_token=token, - verifier=RlfFlagVerifier(relay, flag=task.flag, task_id=episode.task_id, - episode_id=episode.episode_id, - trajectory_id=episode.trajectory_id), - destroy=destroy, describe=describe, deadline_seconds=deadline, - task_id=episode.task_id, episode_id=episode.episode_id, - trajectory_id=episode.trajectory_id, vm_lease=vm, - orchestration=orchestration, environment=relay) diff --git a/yeto/rl/rlf/preflight.py b/yeto/rl/rlf/preflight.py deleted file mode 100644 index e99db9437..000000000 --- a/yeto/rl/rlf/preflight.py +++ /dev/null @@ -1,41 +0,0 @@ -from __future__ import annotations -from typing import Any, Mapping -from .contracts import TrustedGrade -_REQUIRED=('rlf_episode_id','rlf_task_id','rlf_trajectory_id','grader_evidence','success') -class RolloutRejected(ValueError): pass -def validate_rollout(row: Mapping[str,Any]) -> None: - md=row.get('metadata') - if not isinstance(md, Mapping): raise RolloutRejected('missing metadata') - missing=[x for x in _REQUIRED if x not in md] - if missing: raise RolloutRejected('missing trusted fields: '+', '.join(missing)) - if not md.get('rlf_episode_id') or not md.get('rlf_task_id'): raise RolloutRejected('missing RLF identity') - if not isinstance(md['grader_evidence'], Mapping): raise RolloutRejected('invalid grader evidence') - if row.get('cleanup_ok') is False or md.get('cleanup_error'): raise RolloutRejected('cleanup evidence failed') - -RLF_GRADE_KEY = 'rlf_grade' -_GRADE_REQUIRED = ('status', 'reward', 'passed', 'grader_evidence') - -def trusted_grade_from_evaluation(evaluation: Mapping[str, Any]) -> TrustedGrade | None: - """Interface for the harness trusted finish (RLF grade first, TB2 HMAC fallback). - - A ``TrustedVerifier.evaluate()`` payload that carries ``rlf_grade`` (or is - itself a full grade mapping) is parsed into a ``TrustedGrade`` so the - trusted finish can close the episode via ``finish_episode``; anything else - returns ``None`` and the caller falls back to the TB2 HMAC-signed outcome. - A present but malformed grade fails closed (``RolloutRejected``) instead of - silently downgrading to HMAC. - - Not yet wired into ``codex_openenv_agent_function.finish_trusted``; that - call site is outside this change's write scope. - """ - if not isinstance(evaluation, Mapping): return None - payload = evaluation.get(RLF_GRADE_KEY) - if payload is None and all(k in evaluation for k in _GRADE_REQUIRED): payload = evaluation - if payload is None: return None - if not isinstance(payload, Mapping) or not all(k in payload for k in _GRADE_REQUIRED): - raise RolloutRejected('malformed RLF grade in trusted evaluation') - evidence = payload['grader_evidence'] - if not isinstance(evidence, Mapping): raise RolloutRejected('RLF grade evidence must be a mapping') - error = payload.get('error_class') - return TrustedGrade(str(payload['status']), float(payload['reward']), bool(payload['passed']), - dict(evidence), str(error) if error else None) diff --git a/yeto/rl/rlf/trajectory.py b/yeto/rl/rlf/trajectory.py deleted file mode 100644 index d00a34bce..000000000 --- a/yeto/rl/rlf/trajectory.py +++ /dev/null @@ -1,14 +0,0 @@ -from __future__ import annotations -from dataclasses import dataclass, field -from typing import Any, Mapping -@dataclass -class CodexTrajectory: - task_id: str; episode_id: str; trajectory_id: str - tokens: list[int]=field(default_factory=list); logprobs: list[float]=field(default_factory=list) - loss_mask: list[int]=field(default_factory=list); segments: list[Mapping[str,Any]]=field(default_factory=list) - policy_id: str|None=None; metadata: dict[str,Any]=field(default_factory=dict) - def as_metadata(self): - return {'rlf_task_id':self.task_id,'rlf_episode_id':self.episode_id,'rlf_trajectory_id':self.trajectory_id,'policy_id':self.policy_id,'segments':self.segments} - -def from_codex_result(result: Mapping[str,Any], *, task_id: str, episode_id: str, trajectory_id: str, policy_id: str|None=None) -> CodexTrajectory: - return CodexTrajectory(task_id,episode_id,trajectory_id,list(result.get('tokens',[])),list(result.get('logprobs',result.get('rollout_logprobs',[]))),list(result.get('loss_mask',[])),list(result.get('segments',[])),policy_id,dict(result.get('metadata',{}))) diff --git a/yeto/rl/secrlenv_outcome.py b/yeto/rl/secrlenv_outcome.py new file mode 100644 index 000000000..c713d14d4 --- /dev/null +++ b/yeto/rl/secrlenv_outcome.py @@ -0,0 +1,111 @@ +"""Structural check of a SecRLEnv grader verdict in trajectory metadata. + +Kept from the SecRLEnv reward of the removed self-written Codex harness (S26 +E0.5, tag s26-pre-harness-removal) so trajectory evidence of the ``secrlenv`` +kind still verifies; the same role ``tbench_outcome`` has for Terminal-Bench. +Rewards are not signed (openspec ``rl-drop-reward-tamper-hash``): this checks +shape, status and range only. +""" + +from __future__ import annotations + +import json +import math +import re +from typing import Any + +OUTCOME_KEY = "secrlenv_trusted_outcome" +INFRASTRUCTURE_STATUS = "infrastructure_error" +CLEANUP_ERROR_STATUS = "cleanup_error" +INFRASTRUCTURE_503_RETRIES_KEY = "secrlenv_infrastructure_503_retries" +ADMISSION_SCHEMA = 2 +ADMISSION_PHASE_KEY = "admission_phase" +ADMISSION_ERROR_CODE_KEY = "admission_error_code" +ADMISSION_NONCE_KEY = "admission_nonce" +ADMISSION_PHASE = "pre_create" +ADMISSION_ERROR_CODE = "infrastructure_error" +_ADMISSION_NONCE = re.compile(r"[0-9a-f]{32}") +_OUTCOME_STATUSES = { + "completed", + "timeout", + "max_turns", + "max_seq_len", + INFRASTRUCTURE_STATUS, + CLEANUP_ERROR_STATUS, +} + + +class UntrustedOutcome(RuntimeError): + pass + + +def canonical_outcome(value: dict[str, Any]) -> bytes: + return json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8") + + +def verified_outcome(metadata: Any) -> tuple[dict[str, Any], float]: + if not isinstance(metadata, dict): + raise UntrustedOutcome("sample metadata is not an object") + outcome = metadata.get(OUTCOME_KEY) + if not isinstance(outcome, dict): + raise UntrustedOutcome("sample has no secrlenv outcome") + schema = outcome.get("schema") + status = outcome.get("status") + if schema not in {1, ADMISSION_SCHEMA} or status not in _OUTCOME_STATUSES: + raise UntrustedOutcome("secrlenv outcome has an invalid schema or status") + infrastructure_503_retries = outcome.get(INFRASTRUCTURE_503_RETRIES_KEY) + if ( + type(infrastructure_503_retries) is not int + or infrastructure_503_retries not in {0, 1} + ): + raise UntrustedOutcome("secrlenv outcome has an invalid retry ledger") + reward = outcome.get("reward") + if isinstance(reward, bool) or not isinstance(reward, (int, float)): + raise UntrustedOutcome("secrlenv outcome reward is not numeric") + value = float(reward) + if not math.isfinite(value) or not 0.0 <= value <= 1.0: + raise UntrustedOutcome("secrlenv outcome reward is outside [0, 1]") + if not isinstance(outcome.get("task_id"), str): + raise UntrustedOutcome("secrlenv outcome has no episode/task identity") + if schema == ADMISSION_SCHEMA: + expected_keys = { + "schema", + "status", + "episode_id", + "task_id", + "reward", + "passed", + "class", + INFRASTRUCTURE_503_RETRIES_KEY, + ADMISSION_PHASE_KEY, + ADMISSION_ERROR_CODE_KEY, + ADMISSION_NONCE_KEY, + } + if ( + set(outcome) != expected_keys + or status != INFRASTRUCTURE_STATUS + or outcome.get("episode_id") is not None + or not outcome["task_id"] + or outcome.get(ADMISSION_PHASE_KEY) != ADMISSION_PHASE + or outcome.get(ADMISSION_ERROR_CODE_KEY) != ADMISSION_ERROR_CODE + or not isinstance(outcome.get(ADMISSION_NONCE_KEY), str) + or _ADMISSION_NONCE.fullmatch(outcome[ADMISSION_NONCE_KEY]) is None + or infrastructure_503_retries != 1 + or value != 0.0 + or outcome.get("passed") is not False + or outcome.get("class") is not None + ): + raise UntrustedOutcome("secrlenv admission outcome is invalid") + elif ( + not isinstance(outcome.get("episode_id"), str) + or any( + key in outcome + for key in ( + ADMISSION_PHASE_KEY, + ADMISSION_ERROR_CODE_KEY, + ADMISSION_NONCE_KEY, + ) + ) + ): + raise UntrustedOutcome("secrlenv outcome has no episode/task identity") + return outcome, value diff --git a/yeto/rl/ssh_harness.py b/yeto/rl/ssh_harness.py index 8c6975e98..bfa968524 100644 --- a/yeto/rl/ssh_harness.py +++ b/yeto/rl/ssh_harness.py @@ -22,37 +22,11 @@ import torch from . import ( - CODEX_APP_SERVER_PROTOCOL_REVISION, - CODEX_APP_SERVER_SCHEMA_SHA256, - CODEX_BASE_INSTRUCTIONS_SHA256, - CODEX_CLI_VERSION, - CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH, - CODEX_CONTAINER_BINARY_PATH, - CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256, - CODEX_HARNESS_AGENT, - CODEX_HARNESS_AGENT_SHA256, - CODEX_LINUX_BINARY_SHA256, - CODEX_LINUX_BINARY_SIZE_BYTES, - CODEX_LINUX_TARGET, - CODEX_NPM_PACKAGE, - CODEX_NPM_TARBALL_SHA256, - CODEX_OPENENV_AGENT, - CODEX_OPENENV_IDENTITY_ENV, - CODEX_PACKAGE_MANIFEST_SHA256, - CODEX_SUBMIT_TOOL_SCHEMA_SHA256, - CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256, MILES_NEXT_COMMIT, MILES_NEXT_REPOSITORY, MILES_PEFT_VERSION, - SECRLENV_AGENTS, - SECRLENV_GENERATE, - SECRLENV_GROUP_FILTER, - SECRLENV_INFRASTRUCTURE_REPLACEMENTS, - SECRLENV_REWARD, - SECRLENV_ZERO_VARIANCE_REPLACEMENTS, SGLANG_NEXT_COMMIT, SGLANG_NEXT_REPOSITORY, - SIGNED_CODEX_AGENTS, TMS_PRELOAD_BASE_BINARY_SHA256, TMS_PRELOAD_PATCH_BASE_COMMIT, TMS_PRELOAD_PATCH_COMMIT, @@ -112,9 +86,6 @@ _REVISION = re.compile(r"[0-9a-f]{40}\Z") _SAFE_IMAGE = re.compile(r"[a-zA-Z0-9][a-zA-Z0-9._/:@-]+\Z") _EVAL_DATASET_NAME = re.compile(r"[a-zA-Z0-9][a-zA-Z0-9._-]{0,63}\Z") -CODEX_REMOTE_BUNDLE_PATH = "codex/codex-x86_64-unknown-linux-musl" -CODEX_REMOTE_MANIFEST_PATH = "codex/codex-package.json" -CODEX_REMOTE_SCHEMA_PATH = "codex/codex_app_server_protocol.v2.schemas.json" _CONTAINER_STATE_FORMAT = ( '{"Status":{{json .State.Status}},' '"ExitCode":{{.State.ExitCode}},' @@ -502,276 +473,6 @@ def _syncer_source_sha256() -> str: return digest.hexdigest() -def _pinned_codex_adapter_identity() -> dict[str, str]: - return { - "base_instructions_sha256": CODEX_BASE_INSTRUCTIONS_SHA256, - "terminal_exec_tool_schema_sha256": ( - CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256 - ), - "submit_tool_schema_sha256": CODEX_SUBMIT_TOOL_SCHEMA_SHA256, - "dynamic_tools_schema_sha256": CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256, - } - - -def _codex_adapter_source_path() -> Path: - """Resolve the optional legacy adapter from its installed package.""" - - try: - from yeto.rl.harness.codex import codex_harness_agent - except ImportError as exc: - raise HarnessError("cannot attest the Yeto Codex harness adapter") from exc - value = getattr(codex_harness_agent, "__file__", None) - if not isinstance(value, str): - raise HarnessError("Yeto Codex harness adapter has no source identity") - source = Path(value) - if source.is_symlink() or not source.is_file(): - raise HarnessError("Yeto Codex harness adapter source is unavailable") - return source.resolve() - - -def _codex_adapter_identity() -> dict[str, str]: - """Read hashes derived by the live optional Codex adapter.""" - - try: - from yeto.rl.harness.codex import codex_harness_agent - - identity = codex_harness_agent.codex_harness_identity() - except (ImportError, AttributeError, TypeError, ValueError, RuntimeError) as exc: - raise HarnessError("cannot attest the Yeto Codex harness adapter") from exc - required = { - "base_instructions_sha256", - "terminal_exec_tool_schema_sha256", - "submit_tool_schema_sha256", - "dynamic_tools_schema_sha256", - } - if not isinstance(identity, dict) or set(identity) != required: - raise HarnessError("Yeto Codex adapter returned an invalid identity contract") - normalized = {name: str(value).lower() for name, value in identity.items()} - if any(not _SHA256.fullmatch(value) for value in normalized.values()): - raise HarnessError("Yeto Codex adapter identity contains an invalid SHA256") - expected = _pinned_codex_adapter_identity() - if normalized != expected: - raise HarnessError("Yeto Codex adapter identity does not match its Yeto pins") - return normalized - - -def _codex_harness_contract(namespace, args) -> dict[str, Any]: - """Attest the official Linux Codex artifact and the Yeto adapter surface.""" - - from ..provenance import canonical_json_sha256, file_sha256 - - binary_value = getattr(namespace, "codex_harness_binary", None) - manifest_value = getattr(namespace, "codex_package_manifest", None) - schema_value = getattr(namespace, "codex_app_server_schema", None) - if not binary_value or not manifest_value or not schema_value: - raise HarnessError( - "the stock Codex agent requires --codex-harness-binary and " - "--codex-package-manifest and --codex-app-server-schema" - ) - binary = Path(binary_value).expanduser() - manifest = Path(manifest_value).expanduser() - schema = Path(schema_value).expanduser() - for path, flag in ( - (binary, "--codex-harness-binary"), - (manifest, "--codex-package-manifest"), - (schema, "--codex-app-server-schema"), - ): - if path.is_symlink() or not path.is_file(): - raise HarnessError(f"{flag} must be a regular, non-symlink file") - binary = binary.resolve() - manifest = manifest.resolve() - schema = schema.resolve() - if not binary.stat().st_mode & 0o111: - raise HarnessError("--codex-harness-binary must be executable") - if binary.stat().st_size != CODEX_LINUX_BINARY_SIZE_BYTES: - raise HarnessError("Codex Linux binary size does not match its Yeto pin") - if file_sha256(binary) != CODEX_LINUX_BINARY_SHA256: - raise HarnessError("Codex Linux binary SHA256 does not match its Yeto pin") - if file_sha256(manifest) != CODEX_PACKAGE_MANIFEST_SHA256: - raise HarnessError("Codex package manifest SHA256 does not match its Yeto pin") - try: - schema_sha = canonical_json_sha256(schema) - except ValueError: - schema_sha = None - if schema_sha != CODEX_APP_SERVER_SCHEMA_SHA256: - raise HarnessError("Codex app-server schema SHA256 does not match its Yeto pin") - openenv = ( - getattr(args, "custom_agent_function_path", None) == CODEX_OPENENV_AGENT - ) - if openenv: - # The OpenEnv wrapper is shipped by the pinned Miles bundle and is - # attested inside the learner before model allocation. - identity = _pinned_codex_adapter_identity() - else: - identity = _codex_adapter_identity() - adapter_path = _codex_adapter_source_path() - if file_sha256(adapter_path) != CODEX_HARNESS_AGENT_SHA256: - raise HarnessError( - "Yeto Codex harness adapter source does not match its pin" - ) - from .codex_backend import stock_codex_backend_contract_with_reasoning_policy - - profile_name = getattr(args, "codex_backend_profile", None) or args.tito_model - backend = stock_codex_backend_contract_with_reasoning_policy( - profile_name, - args.rollout_max_response_len, - not getattr(args, "codex_exclude_reasoning_from_loss", False), - ) - contract = { - "agent_function_path": CODEX_HARNESS_AGENT, - "agent_source_sha256": CODEX_HARNESS_AGENT_SHA256, - "controller_binary_path": str(binary), - "controller_package_manifest_path": str(manifest), - "controller_app_server_schema_path": str(schema), - "bundle_binary_path": CODEX_REMOTE_BUNDLE_PATH, - "bundle_package_manifest_path": CODEX_REMOTE_MANIFEST_PATH, - "bundle_app_server_schema_path": CODEX_REMOTE_SCHEMA_PATH, - "container_binary_path": CODEX_CONTAINER_BINARY_PATH, - "container_app_server_schema_path": ( - CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH - ), - "binary_sha256": CODEX_LINUX_BINARY_SHA256, - "binary_size_bytes": CODEX_LINUX_BINARY_SIZE_BYTES, - "cli_version": CODEX_CLI_VERSION, - "npm_package": CODEX_NPM_PACKAGE, - "target": CODEX_LINUX_TARGET, - "npm_tarball_sha256": CODEX_NPM_TARBALL_SHA256, - "package_manifest_sha256": CODEX_PACKAGE_MANIFEST_SHA256, - "app_server_protocol_revision": CODEX_APP_SERVER_PROTOCOL_REVISION, - "app_server_schema_sha256": CODEX_APP_SERVER_SCHEMA_SHA256, - **identity, - "reasoning_effort": args.codex_reasoning_effort, - "backend": backend, - } - if openenv: - contract["openenv_identity_env"] = dict(CODEX_OPENENV_IDENTITY_ENV) - return contract - - -def _validate_codex_harness(value: Any, learner: dict[str, Any]) -> None: - required = { - "agent_function_path", - "agent_source_sha256", - "controller_binary_path", - "controller_package_manifest_path", - "controller_app_server_schema_path", - "bundle_binary_path", - "bundle_package_manifest_path", - "bundle_app_server_schema_path", - "container_binary_path", - "container_app_server_schema_path", - "binary_sha256", - "binary_size_bytes", - "cli_version", - "npm_package", - "target", - "npm_tarball_sha256", - "package_manifest_sha256", - "app_server_protocol_revision", - "app_server_schema_sha256", - "base_instructions_sha256", - "terminal_exec_tool_schema_sha256", - "submit_tool_schema_sha256", - "dynamic_tools_schema_sha256", - "reasoning_effort", - "backend", - } - custom_agent = learner.get("custom_agent_function_path") - if custom_agent == CODEX_OPENENV_AGENT: - required.add("openenv_identity_env") - if not isinstance(value, dict) or set(value) != required: - raise HarnessError("plan has an invalid stock Codex harness contract") - pinned = { - "agent_function_path": CODEX_HARNESS_AGENT, - "agent_source_sha256": CODEX_HARNESS_AGENT_SHA256, - "bundle_binary_path": CODEX_REMOTE_BUNDLE_PATH, - "bundle_package_manifest_path": CODEX_REMOTE_MANIFEST_PATH, - "bundle_app_server_schema_path": CODEX_REMOTE_SCHEMA_PATH, - "container_binary_path": CODEX_CONTAINER_BINARY_PATH, - "container_app_server_schema_path": ( - CODEX_CONTAINER_APP_SERVER_SCHEMA_PATH - ), - "binary_sha256": CODEX_LINUX_BINARY_SHA256, - "binary_size_bytes": CODEX_LINUX_BINARY_SIZE_BYTES, - "cli_version": CODEX_CLI_VERSION, - "npm_package": CODEX_NPM_PACKAGE, - "target": CODEX_LINUX_TARGET, - "npm_tarball_sha256": CODEX_NPM_TARBALL_SHA256, - "package_manifest_sha256": CODEX_PACKAGE_MANIFEST_SHA256, - "app_server_protocol_revision": CODEX_APP_SERVER_PROTOCOL_REVISION, - "app_server_schema_sha256": CODEX_APP_SERVER_SCHEMA_SHA256, - "base_instructions_sha256": CODEX_BASE_INSTRUCTIONS_SHA256, - "terminal_exec_tool_schema_sha256": ( - CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256 - ), - "submit_tool_schema_sha256": CODEX_SUBMIT_TOOL_SCHEMA_SHA256, - "dynamic_tools_schema_sha256": CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256, - "reasoning_effort": "xhigh", - } - if any(value.get(name) != expected for name, expected in pinned.items()): - raise HarnessError("plan does not use the pinned stock Codex runtime") - if ( - custom_agent == CODEX_OPENENV_AGENT - and value.get("openenv_identity_env") != CODEX_OPENENV_IDENTITY_ENV - ): - raise HarnessError("plan does not use the pinned Codex OpenEnv surface") - for name in ( - "controller_binary_path", - "controller_package_manifest_path", - "controller_app_server_schema_path", - ): - path = Path(str(value.get(name, ""))) - if not path.is_absolute() or ".." in path.parts: - raise HarnessError(f"stock Codex contract has an invalid {name}") - from .codex_backend import ( - stock_codex_backend_contract_with_reasoning_policy, - validate_stock_codex_fields, - ) - - backend = value.get("backend") - profile_name = str( - learner.get("codex_backend_profile") or learner.get("tito_model", "") - ) - expected_backend = stock_codex_backend_contract_with_reasoning_policy( - profile_name, - learner.get("rollout_max_response_len"), - not learner.get("codex_exclude_reasoning_from_loss", False), - ) - if backend != expected_backend: - raise HarnessError("stock Codex backend/TITO contract drifted") - if custom_agent not in SIGNED_CODEX_AGENTS: - raise HarnessError("stock Codex contract requires a signed agent path") - if ( - custom_agent == CODEX_OPENENV_AGENT - and profile_name != "qwen35_08b" - ): - raise HarnessError( - "the Codex OpenEnv adapter requires backend profile qwen35_08b" - ) - try: - validate_stock_codex_fields( - tito_model=str(learner.get("tito_model", "")), - codex_backend_profile=profile_name, - rl_model_recipe=str(learner.get("rl_model_recipe", "generic")), - model=str(learner.get("model", "")), - model_revision=str(learner.get("model_revision", "")), - rollout_model=learner.get("rollout_model"), - rollout_model_revision=learner.get("rollout_model_revision"), - apply_chat_template_kwargs=learner.get( - "apply_chat_template_kwargs" - ), - tito_allowed_append_roles=learner.get( - "tito_allowed_append_roles" - ), - codex_reasoning_effort=learner.get("codex_reasoning_effort"), - lora_targets=str(learner.get("lora_targets", "")), - expert_full_count=int(learner.get("expert_full_count", 0)), - lora_expert_rank=int(learner.get("rl_lora_expert_rank") or 0), - ) - except (TypeError, ValueError) as exc: - raise HarnessError("stock Codex model profile drifted") from exc - - def _tms_preload_patch(path: str | Path) -> dict[str, str]: """Attest the narrow upstream fix over the pinned Miles image binary.""" @@ -1153,59 +854,8 @@ def _validate_plan(plan: dict[str, Any]) -> None: if plan.get("rl_engine") != "ports": raise HarnessError("plan rl_engine must be 'ports'") _require_ports_plan(learner) - daemon = plan.get("secrlenv_daemon") - if learner.get("custom_agent_function_path") in SECRLENV_AGENTS: - if not isinstance(daemon, dict): - raise HarnessError("secrlenv agent requires a pinned daemon contract") - for name in ("source_root", "task_pack", "state_root"): - _validate_data_path(str(daemon.get(name, "")), f"secrlenv {name}") - for name in ("source_sha256", "task_pack_sha256"): - if not _SHA256.fullmatch(str(daemon.get(name, ""))): - raise HarnessError(f"secrlenv daemon has an invalid {name}") - enable_dind_debug = daemon.get("enable_dind_debug") - if type(enable_dind_debug) is not bool: - raise HarnessError( - "secrlenv daemon enable_dind_debug must be an explicit boolean" - ) - dind_image = daemon.get("dind_image") - dind_image_id = daemon.get("dind_image_id") - dind_debug_script_sha256 = daemon.get("dind_debug_script_sha256") - if enable_dind_debug: - if not _SAFE_IMAGE.fullmatch(str(dind_image or "")) or not re.fullmatch( - r"sha256:[0-9a-f]{64}", str(dind_image_id or "") - ): - raise HarnessError( - "enabled secrlenv DinD debug requires a pinned image and image ID" - ) - if not _SHA256.fullmatch(str(dind_debug_script_sha256 or "")): - raise HarnessError( - "enabled secrlenv DinD debug requires a pinned debug script" - ) - elif any( - value is not None - for value in (dind_image, dind_image_id, dind_debug_script_sha256) - ): - raise HarnessError( - "disabled secrlenv DinD debug must not carry unused debug identities" - ) - if ( - daemon.get("bind") != "127.0.0.1" - or daemon.get("placement", "all-hosts") - not in {"all-hosts", "island-heads"} - or not isinstance(daemon.get("port"), int) - or not 1024 <= daemon["port"] <= 65535 - or not isinstance(daemon.get("max_active_episodes"), int) - or daemon["max_active_episodes"] <= 0 - or not _SAFE_IMAGE.fullmatch(str(daemon.get("operator_image", ""))) - or not re.fullmatch( - r"sha256:[0-9a-f]{64}", str(daemon.get("operator_image_id", "")) - ) - ): - raise HarnessError("secrlenv daemon contract is invalid") - if not plan.get("remote_env_file"): - raise HarnessError("secrlenv daemon requires --remote-env-file") - elif daemon is not None: - raise HarnessError("secrlenv daemon requires the secrlenv custom agent") + if plan.get("secrlenv_daemon") is not None: + raise HarnessError("secrlenv daemon plans are no longer supported (S26 E0.5)") tms_patch = plan.get("tms_preload_patch") tms_train_disk_backup = plan.get("tms_train_disk_backup") if learner.get("rl_offload_train"): @@ -1274,126 +924,9 @@ def _validate_plan(plan: dict[str, Any]) -> None: sync_preset = learner.get("sync_preset", "strict-avg") if sync_preset not in {"strict-avg", "decoupled"}: raise HarnessError("plan has an invalid learner sync_preset") - evaluation = learner.get("evaluation") - if evaluation is not None: - if not isinstance(evaluation, dict): - raise HarnessError("plan has an invalid learner evaluation") - if evaluation.get("eval_only") is not True: - raise HarnessError("SSH evaluation must use a separate eval-only run") - if learner.get("custom_agent_function_path") not in SECRLENV_AGENTS: - raise HarnessError("SSH evaluation is restricted to the secrlenv agent") - if learner.get("reward_function") != SECRLENV_REWARD: - raise HarnessError( - "SSH evaluation requires the signed secrlenv reward contract" - ) - if sync_preset != "strict-avg": - raise HarnessError( - "final SSH evaluation requires strict-avg synchronization" - ) - if data_local_path is None or not PurePosixPath( - str(learner.get("data", "")) - ).suffix: - raise HarnessError( - "SSH evaluation requires one local training dataset file" - ) - if ( - evaluation.get("data") != learner.get("data") - or evaluation.get("data_sha256") != data_sha256 - ): - raise HarnessError( - "evaluation must reuse the exact training dataset path and SHA256" - ) - if not _EVAL_DATASET_NAME.fullmatch( - str(evaluation.get("dataset_name", "")) - ): - raise HarnessError("plan has an invalid evaluation dataset_name") - interval = evaluation.get("interval") - samples = evaluation.get("samples_per_prompt") - prompt_count = evaluation.get("prompt_count") - if ( - not isinstance(interval, int) - or isinstance(interval, bool) - or interval != 1 - ): - raise HarnessError("eval-only Miles evaluation interval must be 1") - if ( - not isinstance(samples, int) - or isinstance(samples, bool) - or samples <= 0 - ): - raise HarnessError("plan has an invalid evaluation samples_per_prompt") - if ( - not isinstance(prompt_count, int) - or isinstance(prompt_count, bool) - or prompt_count <= 0 - ): - raise HarnessError("plan has an invalid evaluation prompt_count") - if evaluation.get("skip_before_train") is not True: - raise HarnessError( - "post-training evaluation must skip pre-train evaluation" - ) - if final_ack_timeout_s is None or final_ack_timeout_s <= 3600: - raise HarnessError( - "evaluation requires an explicitly extended final_ack_timeout_s" - ) - temperature = evaluation.get("temperature") - if temperature is not None and ( - not isinstance(temperature, (int, float)) - or isinstance(temperature, bool) - or not math.isfinite(temperature) - or temperature < 0 - ): - raise HarnessError("plan has an invalid evaluation temperature") - top_p = evaluation.get("top_p") - if top_p is not None and ( - not isinstance(top_p, (int, float)) - or isinstance(top_p, bool) - or not math.isfinite(top_p) - or not 0 < top_p <= 1 - ): - raise HarnessError("plan has an invalid evaluation top_p") - for name in ("max_prompt_len", "max_response_len", "max_context_len"): - value = evaluation.get(name) - if value is not None and ( - not isinstance(value, int) - or isinstance(value, bool) - or value <= 0 - or value > learner["seq_len"] - ): - raise HarnessError(f"plan has an invalid evaluation {name}") - checkpoint = plan.get("eval_checkpoint") - if not isinstance(checkpoint, dict) or set(checkpoint) != { - "source_host", - "source_path", - "sha256", - "size_bytes", - "global_step", - }: - raise HarnessError("eval-only plan requires an attested syncer checkpoint") - try: - _validate_target(str(checkpoint.get("source_host", ""))) - _validate_remote_path( - str(checkpoint.get("source_path", "")), - "eval checkpoint source_path", - ) - except HarnessError as error: - raise HarnessError("plan has an invalid eval checkpoint identity") from error - if ( - not _SHA256.fullmatch(str(checkpoint.get("sha256", ""))) - or not isinstance(checkpoint.get("size_bytes"), int) - or isinstance(checkpoint.get("size_bytes"), bool) - or checkpoint["size_bytes"] <= 12 - or checkpoint.get("global_step") != learner["global_rounds"] - ): - raise HarnessError("plan has an invalid eval checkpoint identity") - if ( - checkpoint["source_host"] == _syncer_host(plan) - and checkpoint["source_path"] - == f"{plan['remote_run']}/state/state.ckpt" - ): - raise HarnessError("eval checkpoint source and fresh destination coincide") - elif plan.get("eval_checkpoint") is not None: - raise HarnessError("eval checkpoint requires an eval-only learner") + if learner.get("evaluation") is not None or plan.get("eval_checkpoint") is not None: + # SSH evaluation ran only the SecRLEnv agent, removed in S26 E0.5. + raise HarnessError("SSH evaluation is restricted to the removed secrlenv agent") fragments = learner.get("fragments", 1) pipeline = learner.get("pipeline", 1) local_horizon = learner.get("local_horizon", 1) @@ -1435,38 +968,7 @@ def _validate_plan(plan: dict[str, Any]) -> None: type(max_replacements) is not int or max_replacements < 0 ): raise HarnessError("plan has an invalid dynamic_sampling_max_replacements") - infrastructure_replacements = learner.get( - "secrlenv_max_infrastructure_replacements" - ) - agent_function = learner.get("custom_agent_function_path") - if agent_function in SECRLENV_AGENTS: - if learner.get("custom_generate_function_path") != SECRLENV_GENERATE: - raise HarnessError( - "secrlenv agent requires the signed generate wrapper" - ) - if learner.get("reward_function") != SECRLENV_REWARD: - raise HarnessError("secrlenv agent requires the signed reward function") - if dynamic_filter != SECRLENV_GROUP_FILTER: - raise HarnessError("secrlenv agent requires the signed group filter") - if max_replacements != SECRLENV_ZERO_VARIANCE_REPLACEMENTS or isinstance( - max_replacements, bool - ): - raise HarnessError( - "secrlenv agent requires zero variance replacements" - ) - if ( - infrastructure_replacements - != SECRLENV_INFRASTRUCTURE_REPLACEMENTS - or isinstance(infrastructure_replacements, bool) - ): - raise HarnessError( - "secrlenv agent requires exactly one infrastructure replacement" - ) - if learner["over_sampling_batch_size"] != learner["groups_per_round"] + 1: - raise HarnessError( - "secrlenv agent requires one reserved same-task retry slot" - ) - elif infrastructure_replacements is not None: + if learner.get("secrlenv_max_infrastructure_replacements") is not None: raise HarnessError( "secrlenv infrastructure replacements require a secrlenv agent" ) @@ -1526,10 +1028,7 @@ def _validate_plan(plan: dict[str, Any]) -> None: or not learner.get("tito_model") ): raise HarnessError("plan has an incomplete agentic session contract") - codex_harness = plan.get("codex_harness") - if custom_agent in SIGNED_CODEX_AGENTS: - _validate_codex_harness(codex_harness, learner) - elif codex_harness is not None: + if plan.get("codex_harness") is not None: raise HarnessError("stock Codex contract requires the signed Codex agent") roles = learner.get("tito_allowed_append_roles") if roles is not None and ( @@ -1926,71 +1425,6 @@ def prepare(namespace) -> Path: if jit_cache_root is not None: _require_cuda_jit_cache(plan) plan["jit_cache"] = _jit_cache_contract(jit_cache_root, plan) - daemon_source_root = getattr(namespace, "secrlenv_source_root", None) - if args.custom_agent_function_path in SECRLENV_AGENTS: - required = { - "--secrlenv-source-root": daemon_source_root, - "--secrlenv-source-sha256": getattr( - namespace, "secrlenv_source_sha256", None - ), - "--secrlenv-task-pack": getattr(namespace, "secrlenv_task_pack", None), - "--secrlenv-task-pack-sha256": getattr( - namespace, "secrlenv_task_pack_sha256", None - ), - "--secrlenv-operator-image": getattr( - namespace, "secrlenv_operator_image", None - ), - "--secrlenv-operator-image-id": getattr( - namespace, "secrlenv_operator_image_id", None - ), - } - missing = [flag for flag, value in required.items() if not value] - if missing: - raise HarnessError( - "secrlenv custom agent requires " + ", ".join(missing) - ) - plan["secrlenv_daemon"] = { - "source_root": daemon_source_root, - "source_sha256": namespace.secrlenv_source_sha256, - "task_pack": namespace.secrlenv_task_pack, - "task_pack_sha256": namespace.secrlenv_task_pack_sha256, - "state_root": f"/data/yeto-rl/secrlenv-runs/{run_id}", - "bind": "127.0.0.1", - "placement": "island-heads", - "port": namespace.secrlenv_port, - "operator_image": namespace.secrlenv_operator_image, - "operator_image_id": namespace.secrlenv_operator_image_id, - "max_active_episodes": namespace.secrlenv_max_active_episodes, - "enable_dind_debug": bool(namespace.secrlenv_enable_dind_debug), - "dind_image": ( - namespace.secrlenv_dind_image - if namespace.secrlenv_enable_dind_debug - else None - ), - "dind_image_id": ( - namespace.secrlenv_dind_image_id - if namespace.secrlenv_enable_dind_debug - else None - ), - "dind_debug_script_sha256": ( - namespace.secrlenv_dind_debug_script_sha256 - if namespace.secrlenv_enable_dind_debug - else None - ), - } - elif daemon_source_root is not None: - raise HarnessError( - "--secrlenv-source-root requires the secrlenv custom agent" - ) - codex_binary = getattr(namespace, "codex_harness_binary", None) - codex_manifest = getattr(namespace, "codex_package_manifest", None) - codex_schema = getattr(namespace, "codex_app_server_schema", None) - if args.custom_agent_function_path in SIGNED_CODEX_AGENTS: - plan["codex_harness"] = _codex_harness_contract(namespace, args) - elif any(value is not None for value in (codex_binary, codex_manifest, codex_schema)): - raise HarnessError( - "Codex artifact options require the signed stock Codex agent" - ) if getattr(args, "expert_full_count", 0): # Recorded only so plan validation rejects it: expert-full needs the # full-parameter sync preset, which this harness does not run. @@ -2044,228 +1478,10 @@ def _remote_vars(plan: dict[str, Any]) -> str: return f'RUN="$HOME/{plan["remote_run"]}"' -def _secrlenv_daemon_script(plan: dict[str, Any]) -> str: - daemon = plan["secrlenv_daemon"] - dind_image_check = "" - dind_server_args = "" - dind_existing_identity = """ - case "$ARGS" in - *--enable-dind-debug*) - echo 'existing secrlenv daemon unexpectedly enables DinD debug' >&2 - exit 1 - ;; - esac -""" - if daemon["enable_dind_debug"]: - dind_image_check = f""" -DIND_IMAGE_ID="$(docker image inspect --format '{{{{.Id}}}}' {shlex.quote(daemon['dind_image'])})" -test "$DIND_IMAGE_ID" = {shlex.quote(daemon['dind_image_id'])} -DIND_DEBUG_SCRIPT="$SOURCE/pipeline/debug_instance.sh" -test -f "$DIND_DEBUG_SCRIPT" && test ! -L "$DIND_DEBUG_SCRIPT" -printf '%s %s\n' {shlex.quote(daemon['dind_debug_script_sha256'])} "$DIND_DEBUG_SCRIPT" | sha256sum --check - -""" - dind_server_args = ( - f" --dind-image {shlex.quote(daemon['dind_image_id'])} \\\n" - " --enable-dind-debug \\\n" - ) - dind_existing_identity = f""" - case "$ARGS" in - *--dind-image*{shlex.quote(daemon['dind_image_id'])}*--enable-dind-debug*) ;; - *) echo 'existing secrlenv daemon has the wrong DinD debug identity' >&2; exit 1 ;; - esac -""" - return f"""set -euo pipefail -{_remote_vars(plan)} -SOURCE={shlex.quote(daemon['source_root'])} -TASK_PACK={shlex.quote(daemon['task_pack'])} -STATE_ROOT={shlex.quote(daemon['state_root'])} -ENV_FILE="$HOME/{plan['remote_env_file']}" -test -d "$SOURCE/secrlenv_rl" -test -d "$TASK_PACK" -test -f "$ENV_FILE" && test ! -L "$ENV_FILE" -test "$(stat -c '%a' "$ENV_FILE")" = 600 -SOURCE_SHA="$(cd "$SOURCE" && find secrlenv_rl -type f -name '*.py' -print0 | sort -z | xargs -0 sha256sum | sha256sum | awk '{{print $1}}')" -test "$SOURCE_SHA" = {shlex.quote(daemon['source_sha256'])} -PYTHONPATH="$RUN/source:$SOURCE${{PYTHONPATH:+:$PYTHONPATH}}" \ - python3 -m yeto.rl.secrlenv_task_images \ - --task-pack "$TASK_PACK" \ - --expected-task-pack-sha256 {shlex.quote(daemon['task_pack_sha256'])} \ - --data-root /data -IMAGE_ID="$(docker image inspect --format '{{{{.Id}}}}' {shlex.quote(daemon['operator_image'])})" -test "$IMAGE_ID" = {shlex.quote(daemon['operator_image_id'])} -{dind_image_check}if [ -s "$STATE_ROOT/daemon.pid" ]; then - PID="$(cat "$STATE_ROOT/daemon.pid")" - ARGS="$(ps -p "$PID" -o args= 2>/dev/null || true)" - case "$ARGS" in - *secrlenv_rl.server*"$STATE_ROOT/state"*) ;; - *) echo 'existing secrlenv daemon PID has the wrong identity' >&2; exit 1 ;; - esac -{dind_existing_identity}else - if python3 - <<'PY' -import socket -s = socket.socket() -try: - s.bind(({daemon['bind']!r}, {daemon['port']})) -finally: - s.close() -PY - then :; else - echo 'secrlenv daemon port is occupied by an unattested process' >&2 - exit 1 - fi - mkdir -p "$STATE_ROOT/state" - TOKEN_FILE="$STATE_ROOT/daemon.token" - if [ ! -e "$TOKEN_FILE" ]; then - umask 077 - python3 - <<'PY' >"$TOKEN_FILE" -import secrets - -print(secrets.token_hex(32)) -PY - fi - test -f "$TOKEN_FILE" && test ! -L "$TOKEN_FILE" - test "$(stat -c '%a' "$TOKEN_FILE")" = 600 - cd "$SOURCE" - nohup setsid env PYTHONPATH="$SOURCE${{PYTHONPATH:+:$PYTHONPATH}}" \ - python3 -m secrlenv_rl.server \ - --task-pack "$TASK_PACK" \ - --state-dir "$STATE_ROOT/state" \ - --repository-root "$SOURCE" \ - --token-file "$TOKEN_FILE" \ - --bind {shlex.quote(daemon['bind'])} \ - --port {daemon['port']} \ - --operator-image {shlex.quote(daemon['operator_image'])} \ -{dind_server_args} --max-active-episodes {daemon['max_active_episodes']} \ - >"$STATE_ROOT/daemon.log" 2>&1 < /dev/null & - echo "$!" > "$STATE_ROOT/daemon.pid" -fi -python3 - <<'PY' -import json -import time -from urllib.request import urlopen - -deadline = time.monotonic() + 60 -last_error = None -while time.monotonic() < deadline: - try: - with urlopen("http://{daemon['bind']}:{daemon['port']}/healthz", timeout=2) as response: - value = json.load(response) - if ( - response.status == 200 - and value.get("ok") is True - and value.get("task_pack_sha256") == {daemon['task_pack_sha256']!r} - and value.get("max_active_episodes") == {daemon['max_active_episodes']} - ): - print("secrlenv_daemon=ready task_pack=" + value["task_pack_sha256"]) - raise SystemExit(0) - last_error = "health identity mismatch" - except Exception as exc: - last_error = str(exc) - time.sleep(1) -raise SystemExit("secrlenv daemon did not become ready: " + str(last_error)) -PY -""" - - -def _start_secrlenv_daemons(plan: dict[str, Any]) -> None: - if plan.get("secrlenv_daemon") is None: - return - for target in _secrlenv_daemon_hosts(plan): - _ssh(plan, target, _secrlenv_daemon_script(plan)) - - -def _stop_secrlenv_daemons(plan: dict[str, Any]) -> None: - daemon = plan.get("secrlenv_daemon") - if daemon is None: - return - for target in _secrlenv_daemon_hosts(plan): - _ssh( - plan, - target, - f"""set -euo pipefail -STATE_ROOT={shlex.quote(daemon['state_root'])} -if [ ! -s "$STATE_ROOT/daemon.pid" ]; then exit 0; fi -PID="$(cat "$STATE_ROOT/daemon.pid")" -ARGS="$(ps -p "$PID" -o args= 2>/dev/null || true)" -case "$ARGS" in - *secrlenv_rl.server*"$STATE_ROOT/state"*) ;; - '') exit 0 ;; - *) echo 'refusing to stop a process with the wrong identity' >&2; exit 1 ;; -esac -kill -TERM -- -"$PID" -for _ in {{1..600}}; do - kill -0 "$PID" 2>/dev/null || exit 0 - sleep 0.1 -done -echo 'secrlenv daemon did not exit after SIGTERM' >&2 -exit 1 -""", - ) - - -def _secrlenv_daemon_status_script(plan: dict[str, Any]) -> str: - daemon = plan.get("secrlenv_daemon") - if daemon is None: - return "" - return f"""STATE_ROOT={shlex.quote(daemon['state_root'])} -if [ ! -s "$STATE_ROOT/daemon.pid" ]; then - echo 'secrlenv_daemon=stopped' -else - PID="$(cat "$STATE_ROOT/daemon.pid")" - case "$PID" in - ''|*[!0-9]*) echo 'secrlenv_daemon=identity-drifted' >&2; exit 1 ;; - esac - ARGS="$(ps -p "$PID" -o args= 2>/dev/null || true)" - case "$ARGS" in - *secrlenv_rl.server*"$STATE_ROOT/state"*) ;; - '') echo 'secrlenv_daemon=stopped' ;; - *) echo 'secrlenv_daemon=identity-drifted' >&2; exit 1 ;; - esac - if [ -n "$ARGS" ]; then - python3 - "$PID" <<'PY' || exit 1 -import json -import sys -from urllib.request import urlopen - -try: - with urlopen("http://{daemon['bind']}:{daemon['port']}/healthz", timeout=2) as response: - value = json.load(response) -except Exception: - print("secrlenv_daemon=unhealthy") - raise SystemExit(1) -if not ( - response.status == 200 - and value.get("ok") is True - and value.get("task_pack_sha256") == {daemon['task_pack_sha256']!r} - and value.get("max_active_episodes") == {daemon['max_active_episodes']} -): - print("secrlenv_daemon=identity-drifted") - raise SystemExit(1) -print("secrlenv_daemon=running pid=" + sys.argv[1]) -PY - fi -fi -""" - - def _all_hosts(plan: dict[str, Any]) -> list[str]: return [host for island in plan["islands"] for host in island["hosts"]] -def _secrlenv_daemon_hosts(plan: dict[str, Any]) -> list[str]: - """Return the pinned daemon roster, retaining old-plan cleanup semantics.""" - - daemon = plan.get("secrlenv_daemon") - if daemon is None: - return [] - placement = daemon.get("placement", "all-hosts") - if placement == "all-hosts": - return _all_hosts(plan) - if placement == "island-heads": - return [island["hosts"][0] for island in plan["islands"]] - raise HarnessError("secrlenv daemon has an invalid placement") - - def _deployment_hosts(plan: dict[str, Any]) -> list[str]: hosts = _all_hosts(plan) syncer_host = _syncer_host(plan) @@ -2292,54 +1508,6 @@ def _attest_local(plan: dict[str, Any]) -> None: != plan["reward_sha256"] ): raise HarnessError("reward source changed after the plan was prepared") - learner = plan["learner"] - codex_harness = plan.get("codex_harness") - if codex_harness is not None: - _validate_codex_harness(codex_harness, learner) - binary = Path(codex_harness["controller_binary_path"]) - manifest = Path(codex_harness["controller_package_manifest_path"]) - schema = Path(codex_harness["controller_app_server_schema_path"]) - for path, expected_size, expected_sha256, executable in ( - ( - binary, - codex_harness["binary_size_bytes"], - codex_harness["binary_sha256"], - True, - ), - (manifest, None, codex_harness["package_manifest_sha256"], False), - (schema, None, codex_harness["app_server_schema_sha256"], False), - ): - if ( - path.is_symlink() - or not path.is_file() - or (executable and not path.stat().st_mode & 0o111) - or (expected_size is not None and path.stat().st_size != expected_size) - or file_sha256(path) != expected_sha256 - ): - raise HarnessError("Codex controller artifact changed after prepare") - if learner.get("custom_agent_function_path") != CODEX_OPENENV_AGENT: - try: - adapter_path = _codex_adapter_source_path() - adapter_identity = _codex_adapter_identity() - except HarnessError: - raise HarnessError( - "Yeto Codex adapter changed after prepare" - ) from None - if ( - file_sha256(adapter_path) - != codex_harness["agent_source_sha256"] - or adapter_identity - != { - name: codex_harness[name] - for name in ( - "base_instructions_sha256", - "terminal_exec_tool_schema_sha256", - "submit_tool_schema_sha256", - "dynamic_tools_schema_sha256", - ) - } - ): - raise HarnessError("Yeto Codex adapter changed after prepare") data_local_path = plan["learner"].get("data_local_path") if ( data_local_path is not None @@ -2405,10 +1573,8 @@ def deploy(plan_path: str | Path) -> None: for target in _deployment_hosts(plan): deploy_data = local_data is not None and target in learner_hosts deploy_checkpoint = eval_checkpoint is not None and target == _syncer_host(plan) - deploy_codex = plan.get("codex_harness") is not None and target in learner_hosts data_directory = '\nmkdir -p "$RUN/data"' if deploy_data else "" state_directory = '\nmkdir -p "$RUN/state"' if deploy_checkpoint else "" - codex_directory = '\nmkdir -p "$RUN/codex"' if deploy_codex else "" initialize = f"""set -euo pipefail {_remote_vars(plan)} if [ -f "$RUN/control/plan.sha256" ]; then @@ -2422,7 +1588,7 @@ def deploy(plan_path: str | Path) -> None: printf '%s\n' {shlex.quote(digest)} > "$RUN/control/deploying.sha256" fi fi -{data_directory}{state_directory}{codex_directory} +{data_directory}{state_directory} """ _ssh(plan, target, initialize) _run( @@ -2462,46 +1628,6 @@ def deploy(plan_path: str | Path) -> None: + ("/" if is_directory else ""), ] ) - if deploy_codex: - codex = plan["codex_harness"] - for source_path, bundle_path in ( - (codex["controller_binary_path"], codex["bundle_binary_path"]), - ( - codex["controller_package_manifest_path"], - codex["bundle_package_manifest_path"], - ), - ( - codex["controller_app_server_schema_path"], - codex["bundle_app_server_schema_path"], - ), - ): - _run( - [ - "rsync", - "-az", - "-e", - _rsync_shell(plan), - source_path, - f"{target}:{plan['remote_run']}/{bundle_path}", - ] - ) - _ssh( - plan, - target, - f"""set -euo pipefail -{_remote_vars(plan)} -CODEX="$RUN/{codex['bundle_binary_path']}" -MANIFEST="$RUN/{codex['bundle_package_manifest_path']}" -SCHEMA="$RUN/{codex['bundle_app_server_schema_path']}" -test -f "$CODEX" && test ! -L "$CODEX" && test -x "$CODEX" -test "$(stat -c %s "$CODEX")" = {codex['binary_size_bytes']} -printf '%s %s\n' {shlex.quote(codex['binary_sha256'])} "$CODEX" | sha256sum --check - -test -f "$MANIFEST" && test ! -L "$MANIFEST" -printf '%s %s\n' {shlex.quote(codex['package_manifest_sha256'])} "$MANIFEST" | sha256sum --check - -test -f "$SCHEMA" && test ! -L "$SCHEMA" -printf '%s %s\n' {shlex.quote(codex['app_server_schema_sha256'])} "$SCHEMA" | sha256sum --check - -""", - ) if deploy_checkpoint: if eval_checkpoint["source_host"] == target: _ssh( @@ -3274,13 +2400,6 @@ def _learner_argv(plan: dict[str, Any], learner_id: int) -> list[str]: str(learner["dynamic_sampling_max_replacements"]), ) ) - if learner.get("secrlenv_max_infrastructure_replacements") is not None: - values.extend( - ( - "--secrlenv-max-infrastructure-replacements", - str(learner["secrlenv_max_infrastructure_replacements"]), - ) - ) if learner.get("dynamic_sampling_filter_path"): values.extend( ( @@ -3326,13 +2445,6 @@ def _learner_argv(plan: dict[str, Any], learner_id: int) -> list[str]: learner["custom_agent_function_path"], ) ) - if plan.get("codex_harness") is not None: - values.extend( - ( - "--codex-harness-contract", - _canonical_json(plan["codex_harness"]), - ) - ) if learner.get("agent_max_seq_len") is not None: values.extend(("--agent-max-seq-len", str(learner["agent_max_seq_len"]))) if learner.get("use_session_server"): @@ -3366,7 +2478,6 @@ def _node_start_script(plan: dict[str, Any], learner_id: int, node_id: int) -> s island = plan["islands"][learner_id] if node_id not in range(len(island["hosts"])): raise HarnessError("node ID is outside the island topology") - daemon_on_node = island["hosts"][node_id] in _secrlenv_daemon_hosts(plan) name = _container_name(plan, learner_id, node_id) head = _target_host(island["hosts"][0]) gpus = island["gpus_per_node"] @@ -3389,22 +2500,6 @@ def _node_start_script(plan: dict[str, Any], learner_id: int, node_id: int) -> s if learner.get("wandb") else "" ) - daemon_env = ( - " --env SECRLENV_DAEMON_URL=" - f"http://127.0.0.1:{plan['secrlenv_daemon']['port']} " - "--env SECRLENV_TASK_PACK_SHA256=" - f"{plan['secrlenv_daemon']['task_pack_sha256']} " - "--env SECRLENV_BEARER_TOKEN_FILE=/run/secrlenv/daemon.token \\\n" - if daemon_on_node - else "" - ) - daemon_volume = ( - " --volume " - f"{shlex.quote(plan['secrlenv_daemon']['state_root'])}/daemon.token:" - "/run/secrlenv/daemon.token:ro \\\n" - if daemon_on_node - else "" - ) reward_scheme = shlex.quote(learner.get("cybergym_reward_scheme", "binary")) reward_view = shlex.quote(learner.get("cybergym_reward_view", "train")) data_volume = ( @@ -3416,69 +2511,6 @@ def _node_start_script(plan: dict[str, Any], learner_id: int, node_id: int) -> s f" --volume {shlex.quote(f'{mount}:{mount}:ro')} \\\n" for mount in learner.get("model_mounts", []) ) - codex = plan.get("codex_harness") - if codex is not None: - codex_contract_sha256 = _plan_digest(codex) - codex_setup = f"""CODEX_HOST="$RUN/{codex['bundle_binary_path']}" -CODEX_MANIFEST_HOST="$RUN/{codex['bundle_package_manifest_path']}" -CODEX_SCHEMA_HOST="$RUN/{codex['bundle_app_server_schema_path']}" -test -f "$CODEX_HOST" && test ! -L "$CODEX_HOST" && test -x "$CODEX_HOST" -test "$(stat -c %s "$CODEX_HOST")" = {codex['binary_size_bytes']} -printf '%s %s\n' {shlex.quote(codex['binary_sha256'])} "$CODEX_HOST" | sha256sum --check - -test -f "$CODEX_MANIFEST_HOST" && test ! -L "$CODEX_MANIFEST_HOST" -printf '%s %s\n' {shlex.quote(codex['package_manifest_sha256'])} "$CODEX_MANIFEST_HOST" | sha256sum --check - -test -f "$CODEX_SCHEMA_HOST" && test ! -L "$CODEX_SCHEMA_HOST" -printf '%s %s\n' {shlex.quote(codex['app_server_schema_sha256'])} "$CODEX_SCHEMA_HOST" | sha256sum --check - -""" - openenv_env = "".join( - f"--env {name}={shlex.quote(value)} " - for name, value in codex.get("openenv_identity_env", {}).items() - ) - codex_env = ( - " --env YETO_CODEX_BINARY_PATH=" - f"{shlex.quote(codex['container_binary_path'])} " - "--env YETO_CODEX_BINARY_SHA256=" - f"{codex['binary_sha256']} " - f"--env YETO_CODEX_BINARY_SIZE_BYTES={codex['binary_size_bytes']} " - f"--env YETO_CODEX_VERSION={shlex.quote(codex['cli_version'])} " - "--env YETO_CODEX_APP_SERVER_PROTOCOL_REVISION=" - f"{codex['app_server_protocol_revision']} " - "--env YETO_CODEX_APP_SERVER_SCHEMA_SHA256=" - f"{codex['app_server_schema_sha256']} " - "--env YETO_CODEX_BASE_INSTRUCTIONS_SHA256=" - f"{codex['base_instructions_sha256']} " - "--env YETO_CODEX_TERMINAL_EXEC_TOOL_SCHEMA_SHA256=" - f"{codex['terminal_exec_tool_schema_sha256']} " - "--env YETO_CODEX_SUBMIT_TOOL_SCHEMA_SHA256=" - f"{codex['submit_tool_schema_sha256']} " - "--env YETO_CODEX_DYNAMIC_TOOLS_SCHEMA_SHA256=" - f"{codex['dynamic_tools_schema_sha256']} " - f"--env YETO_CODEX_REASONING_EFFORT={codex['reasoning_effort']} " - "--env YETO_CODEX_BACKEND_MAX_TOKENS=" - f"{codex['backend']['max_tokens']} " - "--env YETO_CODEX_BACKEND_REASONING_EFFORT=" - f"{codex['backend']['reasoning_effort']} " - "--env YETO_CODEX_BACKEND_THINKING=" - f"{codex['backend']['thinking']['type']} " - f"--env YETO_CODEX_CHAT_TEMPLATE={codex['backend']['chat_template']} " - "--env YETO_CODEX_CHAT_TEMPLATE_KWARGS=" - f"{shlex.quote(_canonical_json(codex['backend']['chat_template_kwargs']))} " - "--env YETO_CODEX_TITO_ALLOWED_APPEND_ROLES=tool,user " - f"{openenv_env}" - f"--env YETO_CODEX_HARNESS_CONTRACT_SHA256={codex_contract_sha256} \\\n" - ) - codex_volumes = ( - " --volume \"$CODEX_HOST:" - f"{codex['container_binary_path']}:ro\" \\\n" - " --volume \"$CODEX_MANIFEST_HOST:/opt/yeto/codex/" - "codex-package.json:ro\" \\\n" - " --volume \"$CODEX_SCHEMA_HOST:" - f"{codex['container_app_server_schema_path']}:ro\" \\\n" - ) - else: - codex_setup = "" - codex_env = "" - codex_volumes = "" tms_patch = plan.get("tms_preload_patch") tms_patch_volume = ( " --volume " @@ -3617,7 +2649,7 @@ def _node_start_script(plan: dict[str, Any], learner_id: int, node_id: int) -> s echo 'refusing to reuse a same-name learner container; use status or explicit restart-learner' >&2 exit 1 fi -{codex_setup}{tms_disk_setup}{jit_cache_setup}mkdir -p "$RUN/island-{learner_id}"/{{state,output,audit,cores}} +{tms_disk_setup}{jit_cache_setup}mkdir -p "$RUN/island-{learner_id}"/{{state,output,audit,cores}} docker run --detach \ --name {shlex.quote(name)} \ --gpus {shlex.quote(gpu_request)} \ @@ -3630,12 +2662,12 @@ def _node_start_script(plan: dict[str, Any], learner_id: int, node_id: int) -> s --env NCCL_IB_DISABLE=1 --env NCCL_SOCKET_IFNAME={network_interface} \ --env GLOO_SOCKET_IFNAME={network_interface} --env TORCH_NCCL_ASYNC_ERROR_HANDLING=1 \ --env CUDA_DEVICE_MAX_CONNECTIONS=1 \ -{attention_env}{recipe_env}{diagnostics_env}{tms_disk_env}{agent_env}{wandb_env}{daemon_env}{codex_env} --env CYBERGYM_URL={shlex.quote(learner['cybergym_url'])} \ +{attention_env}{recipe_env}{diagnostics_env}{tms_disk_env}{agent_env}{wandb_env} --env CYBERGYM_URL={shlex.quote(learner['cybergym_url'])} \ --env CYBERGYM_AGENT_ID={shlex.quote(learner['cybergym_agent_id'])} \ --env CYBERGYM_TIMEOUT={shlex.quote(str(learner['cybergym_timeout']))} \ --env CYBERGYM_REWARD_SCHEME={reward_scheme} \ --env CYBERGYM_REWARD_VIEW={reward_view} \ -{data_volume}{model_volumes}{tms_patch_volume}{tms_disk_volume}{jit_cache_volumes}{daemon_volume}{codex_volumes} --volume "$RUN/source:/workspace/yeto:ro" \ +{data_volume}{model_volumes}{tms_patch_volume}{tms_disk_volume}{jit_cache_volumes} --volume "$RUN/source:/workspace/yeto:ro" \ --volume "$RUN/sglang:/workspace/sglang:ro" \ --volume "$RUN/miles:/workspace/miles" \ --volume "$RUN/island-{learner_id}/state:/workspace/state" \ @@ -3661,7 +2693,6 @@ def start(plan_path: str | Path) -> None: _, plan = load_plan(plan_path) deploy(plan_path) _check_island_cards(plan) - _start_secrlenv_daemons(plan) for island in plan["islands"]: for target in island["hosts"]: _ssh(plan, target, _host_setup_script(plan, island["gpus_per_node"])) @@ -3729,7 +2760,6 @@ def _syncer_status_script(plan: dict[str, Any]) -> str: def status(plan_path: str | Path) -> None: _, plan = load_plan(plan_path) - daemon_hosts = set(_secrlenv_daemon_hosts(plan)) if _syncer_host(plan) not in _all_hosts(plan): result = _ssh( plan, @@ -3751,18 +2781,13 @@ def status(plan_path: str | Path) -> None: syncer = "" if _syncer_host(plan) == target: syncer = _syncer_status_script(plan) - daemon = ( - _secrlenv_daemon_status_script(plan) - if target in daemon_hosts - else "" - ) result = _ssh( plan, target, f"""set -u {_remote_vars(plan)} echo host={shlex.quote(target)} island={learner_id} node={node_id} -{syncer}{daemon}if docker inspect {shlex.quote(name)} >/dev/null 2>&1; then +{syncer}if docker inspect {shlex.quote(name)} >/dev/null 2>&1; then printf 'container=' docker inspect --format {shlex.quote(_CONTAINER_STATE_FORMAT)} {shlex.quote(name)} else echo container=missing; fi @@ -3956,7 +2981,6 @@ def stop(plan_path: str | Path) -> None: _syncer_host(plan), _syncer_stop_script(plan), ) - _stop_secrlenv_daemons(plan) def _collect_event_evidence( @@ -4854,41 +3878,6 @@ def build_parser() -> argparse.ArgumentParser: prepare_parser.add_argument("--eval-max-response-len", type=int, default=None) prepare_parser.add_argument("--eval-max-context-len", type=int, default=None) prepare_parser.add_argument("--model-manifest-sha256", default=None) - prepare_parser.add_argument("--secrlenv-source-root", default=None) - prepare_parser.add_argument("--secrlenv-source-sha256", default=None) - prepare_parser.add_argument("--secrlenv-task-pack", default=None) - prepare_parser.add_argument("--secrlenv-task-pack-sha256", default=None) - prepare_parser.add_argument("--secrlenv-operator-image", default=None) - prepare_parser.add_argument("--secrlenv-operator-image-id", default=None) - prepare_parser.add_argument("--secrlenv-port", type=int, default=28765) - prepare_parser.add_argument( - "--secrlenv-max-active-episodes", type=int, default=16 - ) - prepare_parser.add_argument( - "--secrlenv-enable-dind-debug", - action="store_true", - help="give each solver a private flagless DinD challenge copy", - ) - prepare_parser.add_argument("--secrlenv-dind-image", default=None) - prepare_parser.add_argument("--secrlenv-dind-image-id", default=None) - prepare_parser.add_argument( - "--secrlenv-dind-debug-script-sha256", default=None - ) - prepare_parser.add_argument( - "--codex-harness-binary", - default=None, - help="controller path to the pinned official Codex Linux binary", - ) - prepare_parser.add_argument( - "--codex-package-manifest", - default=None, - help="controller path to the pinned codex-package.json identity", - ) - prepare_parser.add_argument( - "--codex-app-server-schema", - default=None, - help="controller path to the pinned experimental v2 app-server schema", - ) prepare_parser.add_argument("--tms-preload-patch", default=None) prepare_parser.add_argument( "--tms-train-disk-backup-root", diff --git a/yeto/rl/tbench_direct_preflight.py b/yeto/rl/tbench_direct_preflight.py deleted file mode 100644 index f8609c393..000000000 --- a/yeto/rl/tbench_direct_preflight.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Pre-allocation attestation for direct Miles Terminal-Bench SAO launches.""" - -from __future__ import annotations - -import argparse -import importlib -import os -import sys -from pathlib import Path -from typing import Any - -from . import ( - CODEX_OPENENV_AGENT, - CODEX_OPENENV_AGENT_MODULES, - CODEX_OPENENV_IDENTITY_ENV, -) -from .codex_backend import QWEN35_08B_MODEL, QWEN35_08B_REVISION - -_REWARD_FUNCTION = "openenv_generate.reward_func" -_FILTER_FUNCTION = "openenv_generate.check_terminal_bench_episode" -_GENERATE_FUNCTION = "miles.rollout.generate_hub.agentic_tool_call.generate" -_COMPACTION_ENV = { - "YETO_CODEX_COMPACTION_ENABLED": "1", - "YETO_CODEX_COMPACTION_TRIGGER_TOKENS": "6144", - "YETO_CODEX_COMPACTION_SUMMARY_MAX_TOKENS": "1024", - "YETO_CODEX_MAX_COMPACTIONS": "3", -} -_TIMEOUT_ENV = { - "OPENENV_MAX_ROLLOUT_TIME_SECONDS": "1800", - "SECRLENV_MAX_ROLLOUT_TIME_SECONDS": "1800", -} - - -def _direct_options(argv: list[str]) -> argparse.Namespace: - parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False) - parser.add_argument("--custom-generate-function-path") - parser.add_argument("--custom-agent-function-path") - parser.add_argument("--custom-rm-path") - parser.add_argument("--dynamic-sampling-filter-path") - parser.add_argument("--input-key") - parser.add_argument("--tito-model") - parser.add_argument("--sao-online-recipe") - parser.add_argument("--max-seq-len", type=int) - parser.add_argument("--num-gpus-per-node", type=int) - parser.add_argument("--sglang-mem-fraction-static", type=float) - parser.add_argument("--sglang-max-total-tokens", type=int) - parser.add_argument("--sglang-max-mamba-cache-size", type=int) - parser.add_argument("--sao-compaction", action="store_true") - parser.add_argument("--sao-one-gpu-island", action="store_true") - parser.add_argument("--colocate", action="store_true") - values, _ = parser.parse_known_args(argv) - return values - - -def _attest_adapter(miles_root: Path) -> None: - from yeto.rl.harness.codex.preflight import forward_legacy_openenv_preflight - - if miles_root.is_symlink() or not miles_root.is_dir(): - raise ValueError("the pinned Miles root is not a real directory") - forward_legacy_openenv_preflight(None, "qwen35_08b") - - -def preflight_tbench_codex_streaming( - argv: list[str], - *, - sao_context: dict[str, Any], - streaming_runtime: Any, - miles_root: str | Path, -) -> None: - """Fail before Miles parses arguments or creates any model/Ray actor.""" - - if ( - getattr(streaming_runtime, "trajectory_evidence_kind", None) - != "terminal-bench-2.1" - or getattr(streaming_runtime, "trajectory_evidence_schema_version", None) - != 2 - ): - raise ValueError( - "Terminal-Bench direct launch requires trajectory evidence v2" - ) - if ( - sao_context.get("benchmark") != "terminal-bench-2.1" - or sao_context.get("model") != QWEN35_08B_MODEL - or sao_context.get("base_model_revision") != QWEN35_08B_REVISION - or sao_context.get("rollout_model_revision") != QWEN35_08B_REVISION - ): - raise ValueError("Terminal-Bench direct launch model identity drifted") - - options = _direct_options(argv) - expected = { - "custom_generate_function_path": _GENERATE_FUNCTION, - "custom_agent_function_path": CODEX_OPENENV_AGENT, - "custom_rm_path": _REWARD_FUNCTION, - "dynamic_sampling_filter_path": _FILTER_FUNCTION, - "input_key": "messages", - "tito_model": "qwen35", - "sao_online_recipe": "coding", - "max_seq_len": 8192, - "num_gpus_per_node": 1, - "sglang_mem_fraction_static": 0.15, - "sglang_max_total_tokens": 393216, - "sglang_max_mamba_cache_size": 256, - "sao_compaction": True, - "sao_one_gpu_island": True, - "colocate": True, - } - mismatched = [ - name for name, value in expected.items() if getattr(options, name) != value - ] - if mismatched: - raise ValueError( - "Terminal-Bench direct Miles arguments drifted: " - + ", ".join(mismatched) - ) - - environment = { - **CODEX_OPENENV_IDENTITY_ENV, - **_COMPACTION_ENV, - **_TIMEOUT_ENV, - } - mismatched_environment = [ - name for name, value in environment.items() if os.getenv(name) != value - ] - if mismatched_environment: - raise ValueError( - "Terminal-Bench trusted container environment drifted: " - + ", ".join(mismatched_environment) - ) - _attest_adapter(Path(miles_root).expanduser().resolve()) diff --git a/yeto/rl/trajectory_evidence.py b/yeto/rl/trajectory_evidence.py index 7ee7c67da..3087c7e71 100644 --- a/yeto/rl/trajectory_evidence.py +++ b/yeto/rl/trajectory_evidence.py @@ -498,14 +498,12 @@ def _verified_outcome( if "tbench_trusted_outcome" in metadata: raise ValueError("trajectory metadata mixes benchmark evidence kinds") - from yeto.rl.harness.codex.reward import ( + from .secrlenv_outcome import ( CLEANUP_ERROR_STATUS, INFRASTRUCTURE_STATUS, - _canonical, - ) - from yeto.rl.harness.codex.reward import ( - _verified_outcome as verified_secrlenv_outcome, + canonical_outcome as _canonical, ) + from .secrlenv_outcome import verified_outcome as verified_secrlenv_outcome outcome, reward = verified_secrlenv_outcome(dict(metadata)) if outcome["status"] in {INFRASTRUCTURE_STATUS, CLEANUP_ERROR_STATUS}: From 9ced3bf41429d8a84d3021a9ae18bda5ec5458c4 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:57:15 +0000 Subject: [PATCH 003/233] Drop dead Codex/SecRLEnv plumbing left after the harness removal The removed agents made these paths unreachable: - --codex-backend-profile, --codex-reasoning-effort, --codex-exclude-reasoning-from-loss and --secrlenv-max-infrastructure-replacements become launcher-only flags (still refused); the island no longer parses them or --codex-harness-contract, and no longer copies them into miles_args. - The SSH harness drops its SecRLEnv-only eval path: eval checkpoint contract, prepare options, checkpoint deploy, eval-only argv, syncer lifecycle evidence, eval collect/verify, and the codex fields of the plan. Plans that still carry eval, daemon or Codex fields are refused. - The island checkpoint config keeps its codex/secrlenv keys as None so checkpoints written before the removal still resume. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_ssh_harness.py | 55 --- yeto/rl/adapters/miles/island_entry.py | 25 -- yeto/rl/adapters/miles/run_flags.py | 4 +- yeto/rl/engine/progress.py | 19 +- yeto/rl/engine/run_flags.py | 6 +- yeto/rl/ssh_harness.py | 479 +------------------------ 6 files changed, 23 insertions(+), 565 deletions(-) diff --git a/tests/test_rl_ssh_harness.py b/tests/test_rl_ssh_harness.py index 091bbe386..b5fb662c7 100644 --- a/tests/test_rl_ssh_harness.py +++ b/tests/test_rl_ssh_harness.py @@ -589,61 +589,6 @@ def test_plan_rejects_previous_pipeline_identity_schema(): ssh_harness._validate_plan(plan) -def test_prepare_cli_exposes_final_secrlenv_eval_controls(): - args = ssh_harness.build_parser().parse_args( - [ - "prepare", - "--host", - "root@h200-n1,root@h200-n2", - "--run-id", - "flaky100", - "--final-ack-timeout-s", - "21600", - "--eval-checkpoint-host", - "root@h200-n6", - "--eval-checkpoint-path", - "yeto-rl-ssh-data/train/state/state.ckpt", - "--eval-checkpoint-sha256", - "f" * 64, - "--eval-checkpoint-size-bytes", - "138632000000", - "--eval-checkpoint-global-step", - "8", - "--eval-dataset-name", - "flaky100", - "--eval-interval", - "1", - "--eval-samples-per-prompt", - "2", - "--eval-temperature", - "0.7", - "--eval-top-p", - "0.95", - "--eval-max-prompt-len", - "4096", - "--eval-max-response-len", - "32768", - "--eval-max-context-len", - "32768", - ] - ) - - assert args.final_ack_timeout_s == 21600 - assert args.eval_checkpoint_host == "root@h200-n6" - assert args.eval_checkpoint_path.endswith("train/state/state.ckpt") - assert args.eval_checkpoint_sha256 == "f" * 64 - assert args.eval_checkpoint_size_bytes == 138_632_000_000 - assert args.eval_checkpoint_global_step == 8 - assert args.eval_dataset_name == "flaky100" - assert args.eval_interval == 1 - assert args.eval_samples_per_prompt == 2 - assert args.eval_temperature == 0.7 - assert args.eval_top_p == 0.95 - assert args.eval_max_prompt_len == 4096 - assert args.eval_max_response_len == 32768 - assert args.eval_max_context_len == 32768 - - def test_plan_requires_the_patched_sglang_pin(): plan = _plan() plan.pop("sglang") diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 103d8ca24..0193f0c35 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -1526,18 +1526,6 @@ def run_miles( write_boot_only_marker(args, ports_launch, run_config, ref_probe) return # boot-only: exit 0 before Ray/GPU work; no training - if getattr(args, "codex_exclude_reasoning_from_loss", False): - # S17 WP6: read by rollout_meta_hook.record_trained_groups in the - # rollout process (with or without Yeto policy sync); absent (default) == reasoning counted, no-op. - miles_args.yeto_rl_codex_reasoning_in_loss = False - from yeto.rl.codex_backend import stock_codex_reasoning_markers - - miles_args.yeto_rl_codex_reasoning_markers = list( - stock_codex_reasoning_markers( - getattr(args, "codex_backend_profile", None) or args.tito_model - ) - or () - ) # The full-parameter sync contracts bind the spec hash (4.8). The island # contract v3 (strict, decoupled, elastic) binds the semantic hash and the # run code pack instead (rl-identity-tiering L9.7); placement capabilities @@ -1586,22 +1574,9 @@ def run_miles( spec.name for spec in specs ) miles_args.yeto_rl_reward_sha256 = args.reward_sha256 - miles_args.yeto_rl_codex_harness_contract = getattr( - args, - "codex_harness_contract", - None, - ) - miles_args.yeto_rl_codex_backend_profile = getattr( - args, - "codex_backend_profile", - None, - ) miles_args.yeto_rl_dynamic_sampling_max_replacements = getattr( args, "dynamic_sampling_max_replacements", None ) - miles_args.yeto_rl_secrlenv_max_infrastructure_replacements = getattr( - args, "secrlenv_max_infrastructure_replacements", None - ) miles_args.yeto_rl_completed_groups_path = args.completed_groups_path miles_args.yeto_rl_event_tape = args.event_tape miles_args.yeto_rl_learner_id = args.learner_id diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 8a9f52027..3319f5d76 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -9,7 +9,6 @@ from __future__ import annotations import argparse -import json from collections.abc import Mapping from types import MappingProxyType @@ -241,7 +240,7 @@ "on few-core single-GPU islands to avoid the parallel-import storm " "(32 sglang imports can miss the 300 s ready window)", ), - shared_flag( + cli_flag( # refused since S26 E0.5 (removed Codex/SecRLEnv agents) "--codex-backend-profile", default=None, help=( @@ -337,7 +336,6 @@ island_flag("--groups-per-round", type=int, required=True), island_flag("--samples-per-group", type=int, required=True), island_flag("--optimizer-steps", type=int, required=True), - island_flag("--codex-harness-contract", type=json.loads, default=None), island_flag("--completed-groups-path", required=True), island_flag("--event-tape", required=True), island_flag("--audit-dir", default=None), diff --git a/yeto/rl/engine/progress.py b/yeto/rl/engine/progress.py index 71441da40..9c1891f13 100644 --- a/yeto/rl/engine/progress.py +++ b/yeto/rl/engine/progress.py @@ -54,11 +54,9 @@ def _island_checkpoint_config(args) -> dict[str, Any]: "dynamic_sampling_max_replacements": getattr( args, "yeto_rl_dynamic_sampling_max_replacements", None ), - "secrlenv_max_infrastructure_replacements": getattr( - args, - "yeto_rl_secrlenv_max_infrastructure_replacements", - None, - ), + # The SecRLEnv and Codex agents left in S26 E0.5; the keys stay (None) + # so island checkpoints written before that still resume. + "secrlenv_max_infrastructure_replacements": None, "rl_offload_train": bool(getattr(args, "offload_train", False)), "rl_distributed_timeout_minutes": getattr( args, "distributed_timeout_minutes", 10 @@ -69,17 +67,8 @@ def _island_checkpoint_config(args) -> dict[str, Any]: "custom_generate_function_path": args.custom_generate_function_path, "use_session_server": args.use_session_server, "tito_model": args.tito_model, - "codex_backend_profile": getattr( - args, "yeto_rl_codex_backend_profile", None - ), + "codex_backend_profile": None, } - codex_harness_contract = getattr( - args, - "yeto_rl_codex_harness_contract", - None, - ) - if codex_harness_contract is not None: - config["codex_harness_contract"] = codex_harness_contract initial_adapter_sha256 = getattr( args, "yeto_rl_initial_adapter_sha256", diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index b4de91487..7ab563b2f 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -301,7 +301,7 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "fallback after this many rejected groups" ), ), - shared_flag( + cli_flag( # refused since S26 E0.5 (removed Codex/SecRLEnv agents) "--secrlenv-max-infrastructure-replacements", type=int, default=None, @@ -798,7 +798,7 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. _AGENT = ( shared_flag("--apply-chat-template-kwargs", type=json.loads, default=None), shared_flag("--custom-generate-function-path", default=None), - shared_flag( + cli_flag( # refused since S26 E0.5 (removed Codex/SecRLEnv agents) "--codex-reasoning-effort", choices=["xhigh"], default=None, @@ -811,7 +811,7 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. shared_flag("--session-server-ip", default=None), shared_flag("--session-server-port", type=int, nargs="+", default=None), shared_flag("--tito-model", default=None), - shared_flag( + cli_flag( # refused since S26 E0.5 (removed Codex/SecRLEnv agents) "--codex-exclude-reasoning-from-loss", action="store_true", help=( diff --git a/yeto/rl/ssh_harness.py b/yeto/rl/ssh_harness.py index bfa968524..686e73236 100644 --- a/yeto/rl/ssh_harness.py +++ b/yeto/rl/ssh_harness.py @@ -5,12 +5,10 @@ import argparse import hashlib import json -import math import os import re import shlex import shutil -import statistics import subprocess import sys import time @@ -85,7 +83,6 @@ }) _REVISION = re.compile(r"[0-9a-f]{40}\Z") _SAFE_IMAGE = re.compile(r"[a-zA-Z0-9][a-zA-Z0-9._/:@-]+\Z") -_EVAL_DATASET_NAME = re.compile(r"[a-zA-Z0-9][a-zA-Z0-9._-]{0,63}\Z") _CONTAINER_STATE_FORMAT = ( '{"Status":{{json .State.Status}},' '"ExitCode":{{.State.ExitCode}},' @@ -111,15 +108,6 @@ "rl_policy_apply": ("event", "sync/global_policy_hash"), "rl_final_cut": ("event",), }, - "eval-learner": { - "rl_policy_apply": ( - "event", "island_id", "policy_version", "sync/global_policy_hash" - ), - "rl_eval_result": ( - "event", "island_id", "rollout_id", "policy_version", "dataset_name", - "sample_count", "rl/eval/result", "rl/eval/pass_at_1" - ), - }, "syncer": { "syncer_commit": ( "step", "fragment", "launch_base_version", "expected", "responded", @@ -226,26 +214,6 @@ def scalar(name, value): except (OSError, UnicodeError, ValueError, json.JSONDecodeError): raise SystemExit("invalid event evidence source") """ -_SYNCER_LIFECYCLE_EVIDENCE = r""" -import json -import sys -from pathlib import Path - -path = Path(sys.argv[1]) -if path.is_symlink() or not path.is_file() or path.stat().st_size > 64 * 1024 * 1024: - raise SystemExit("invalid syncer lifecycle source") -resume = 0 -final_ack = 0 -try: - with path.open(encoding="utf-8", errors="strict") as source: - for line in source: - resume += line.count("resumed from checkpoint") - final_ack += line.count("learner finalized authoritative parameters") -except (OSError, UnicodeError): - raise SystemExit("invalid syncer lifecycle source") -print(json.dumps({"schema": 1, "resume_count": resume, "final_ack_count": final_ack}, - sort_keys=True, separators=(",", ":"))) -""" _EVENT_EVIDENCE_FIELDS = { "strict-learner": { "rl_local_round": { @@ -271,24 +239,6 @@ def scalar(name, value): }, "rl_final_cut": {"event"}, }, - "eval-learner": { - "rl_policy_apply": { - "event", - "island_id", - "policy_version", - "sync/global_policy_hash", - }, - "rl_eval_result": { - "event", - "island_id", - "rollout_id", - "policy_version", - "dataset_name", - "sample_count", - "rl/eval/result", - "rl/eval/pass_at_1", - }, - }, "syncer": { "syncer_commit": { "step", @@ -709,49 +659,6 @@ def _local_data_sha256(path: str | Path) -> str: return digest.hexdigest() -def _local_jsonl_record_count(path: str | Path) -> int: - source = Path(path) - if source.is_symlink() or not source.is_file(): - raise HarnessError("evaluation dataset must be one regular JSONL file") - try: - with source.open("rb") as handle: - count = sum(1 for line in handle if line.strip()) - except OSError as error: - raise HarnessError("cannot read evaluation dataset") from error - if count <= 0: - raise HarnessError("evaluation dataset is empty") - return count - - -def _eval_checkpoint_contract( - source_host: str, - source_path: str, - sha256: str, - size_bytes: int, - global_step: int, -) -> dict[str, Any]: - """Pin a remote terminal checkpoint without staging it on the controller.""" - - source_host = _validate_target(source_host) - source_path = _validate_remote_path(source_path, "--eval-checkpoint-path") - sha256 = str(sha256).lower() - if not _SHA256.fullmatch(sha256): - raise HarnessError("--eval-checkpoint-sha256 must be a lowercase SHA256") - for flag, value in ( - ("--eval-checkpoint-size-bytes", size_bytes), - ("--eval-checkpoint-global-step", global_step), - ): - if not isinstance(value, int) or isinstance(value, bool) or value <= 0: - raise HarnessError(f"{flag} must be a positive integer") - return { - "source_host": source_host, - "source_path": source_path, - "sha256": sha256, - "size_bytes": size_bytes, - "global_step": global_step, - } - - def _canonical_json(value: Any) -> str: return json.dumps(value, sort_keys=True, separators=(",", ":")) @@ -1176,7 +1083,6 @@ def prepare(namespace) -> Path: data = args.data data_local_path = None data_sha256 = None - data_prompt_count = None if dataset["source"] == "local": source = Path(args.data).expanduser().absolute() data_sha256 = _local_data_sha256(source) @@ -1185,8 +1091,6 @@ def prepare(namespace) -> Path: data = "/workspace/data/dataset" if source.is_file(): data += source.suffix - if getattr(namespace, "eval_checkpoint_host", None) is not None: - data_prompt_count = _local_jsonl_record_count(source) tms_patch_path = getattr(namespace, "tms_preload_patch", None) tms_preload_patch = ( @@ -1292,9 +1196,6 @@ def prepare(namespace) -> Path: "dynamic_sampling_max_replacements": getattr( args, "dynamic_sampling_max_replacements", None ), - "secrlenv_max_infrastructure_replacements": getattr( - args, "secrlenv_max_infrastructure_replacements", None - ), "rl_offload_train": bool(getattr(args, "rl_offload_train", False)), "rl_distributed_timeout_minutes": getattr( args, "rl_distributed_timeout_minutes", 10 @@ -1308,17 +1209,11 @@ def prepare(namespace) -> Path: "custom_agent_function_path": getattr( args, "custom_agent_function_path", None ), - "codex_reasoning_effort": getattr( - args, "codex_reasoning_effort", None - ), "agent_max_seq_len": getattr(args, "agent_max_seq_len", None), "use_session_server": args.use_session_server, "session_server_ip": args.session_server_ip, "session_server_port": args.session_server_port, "tito_model": args.tito_model, - "codex_backend_profile": getattr( - args, "codex_backend_profile", None - ), "tito_allowed_append_roles": getattr( args, "tito_allowed_append_roles", None ), @@ -1352,75 +1247,9 @@ def prepare(namespace) -> Path: "cybergym_reward_view": os.environ.get("CYBERGYM_REWARD_VIEW", "train"), }, } - if getattr(args, "codex_exclude_reasoning_from_loss", False): - # S17 WP6: only present when opted out, so default plans are unchanged. - plan["learner"]["codex_exclude_reasoning_from_loss"] = True final_ack_timeout_s = getattr(namespace, "final_ack_timeout_s", None) if final_ack_timeout_s is not None: plan["final_ack_timeout_s"] = final_ack_timeout_s - eval_checkpoint_host = getattr(namespace, "eval_checkpoint_host", None) - eval_checkpoint_path = getattr(namespace, "eval_checkpoint_path", None) - eval_checkpoint_sha256 = getattr(namespace, "eval_checkpoint_sha256", None) - eval_checkpoint_size_bytes = getattr( - namespace, "eval_checkpoint_size_bytes", None - ) - eval_checkpoint_global_step = getattr( - namespace, "eval_checkpoint_global_step", None - ) - checkpoint_values = ( - eval_checkpoint_host, - eval_checkpoint_path, - eval_checkpoint_sha256, - eval_checkpoint_size_bytes, - eval_checkpoint_global_step, - ) - if any(value is not None for value in checkpoint_values) and any( - value is None for value in checkpoint_values - ): - raise HarnessError( - "all remote eval checkpoint identity options are required together" - ) - eval_interval = getattr(namespace, "eval_interval", None) - eval_dataset_name = getattr(namespace, "eval_dataset_name", None) - eval_options = { - "temperature": getattr(namespace, "eval_temperature", None), - "top_p": getattr(namespace, "eval_top_p", None), - "max_prompt_len": getattr(namespace, "eval_max_prompt_len", None), - "max_response_len": getattr(namespace, "eval_max_response_len", None), - "max_context_len": getattr(namespace, "eval_max_context_len", None), - } - if eval_checkpoint_host is None: - if ( - eval_interval is not None - or eval_dataset_name is not None - or getattr(namespace, "eval_samples_per_prompt", None) is not None - or any(value is not None for value in eval_options.values()) - ): - raise HarnessError("evaluation options require --eval-checkpoint-host") - else: - plan["eval_checkpoint"] = _eval_checkpoint_contract( - eval_checkpoint_host, - eval_checkpoint_path, - eval_checkpoint_sha256, - eval_checkpoint_size_bytes, - eval_checkpoint_global_step, - ) - eval_samples_per_prompt = getattr( - namespace, "eval_samples_per_prompt", None - ) - plan["learner"]["evaluation"] = { - "eval_only": True, - "dataset_name": eval_dataset_name, - "data": data, - "data_sha256": data_sha256, - "interval": 1 if eval_interval is None else eval_interval, - "samples_per_prompt": ( - 1 if eval_samples_per_prompt is None else eval_samples_per_prompt - ), - "prompt_count": data_prompt_count, - "skip_before_train": True, - **eval_options, - } jit_cache_root = getattr(namespace, "jit_cache_root", None) if jit_cache_root is not None: _require_cuda_jit_cache(plan) @@ -1514,11 +1343,6 @@ def _attest_local(plan: dict[str, Any]) -> None: and _local_data_sha256(data_local_path) != plan["learner"]["data_sha256"] ): raise HarnessError("local dataset changed after the plan was prepared") - evaluation = plan["learner"].get("evaluation") - if evaluation is not None and _local_jsonl_record_count( - data_local_path - ) != evaluation.get("prompt_count"): - raise HarnessError("evaluation dataset record count changed after prepare") def deploy(plan_path: str | Path) -> None: @@ -1528,31 +1352,6 @@ def deploy(plan_path: str | Path) -> None: _attest_local(plan) digest = _plan_digest(plan) local_data = plan["learner"].get("data_local_path") - eval_checkpoint = plan.get("eval_checkpoint") - if eval_checkpoint is not None: - _require_program("scp") - source_path = eval_checkpoint["source_path"] - _ssh( - plan, - eval_checkpoint["source_host"], - f"""set -euo pipefail -SOURCE="$HOME/{source_path}" -test -f "$SOURCE" && test ! -L "$SOURCE" -test "$(stat -c %s "$SOURCE")" = {eval_checkpoint['size_bytes']} -printf '%s %s\n' {shlex.quote(eval_checkpoint['sha256'])} "$SOURCE" | sha256sum --check - -python3 - "$SOURCE" {eval_checkpoint['global_step']} <<'PY' -import sys -from pathlib import Path - -header = Path(sys.argv[1]).open("rb").read(12) -expected_step = int(sys.argv[2]) -if len(header) != 12 or int.from_bytes(header[:4], "little") != {SYNCER_CHECKPOINT_MAGIC}: - raise SystemExit("invalid syncer checkpoint header") -if int.from_bytes(header[4:12], "little") != expected_step: - raise SystemExit("syncer checkpoint global step mismatch") -PY -""", - ) excludes = ( ".git/", ".venv/", @@ -1572,9 +1371,7 @@ def deploy(plan_path: str | Path) -> None: learner_hosts = set(_all_hosts(plan)) for target in _deployment_hosts(plan): deploy_data = local_data is not None and target in learner_hosts - deploy_checkpoint = eval_checkpoint is not None and target == _syncer_host(plan) data_directory = '\nmkdir -p "$RUN/data"' if deploy_data else "" - state_directory = '\nmkdir -p "$RUN/state"' if deploy_checkpoint else "" initialize = f"""set -euo pipefail {_remote_vars(plan)} if [ -f "$RUN/control/plan.sha256" ]; then @@ -1588,7 +1385,7 @@ def deploy(plan_path: str | Path) -> None: printf '%s\n' {shlex.quote(digest)} > "$RUN/control/deploying.sha256" fi fi -{data_directory}{state_directory} +{data_directory} """ _ssh(plan, target, initialize) _run( @@ -1628,29 +1425,6 @@ def deploy(plan_path: str | Path) -> None: + ("/" if is_directory else ""), ] ) - if deploy_checkpoint: - if eval_checkpoint["source_host"] == target: - _ssh( - plan, - target, - f"""set -euo pipefail -{_remote_vars(plan)} -cp --reflink=auto -- "$HOME/{eval_checkpoint['source_path']}" "$RUN/state/state.ckpt" -""", - ) - else: - _run( - [ - "scp", - "-3", - *plan.get("ssh_options", []), - ( - f"{eval_checkpoint['source_host']}:" - f"{eval_checkpoint['source_path']}" - ), - f"{target}:{plan['remote_run']}/state/state.ckpt", - ] - ) _ssh( plan, target, @@ -2081,16 +1855,9 @@ def _syncer_runtime_identity_functions() -> str: def _syncer_start_script(plan: dict[str, Any]) -> str: unit = shlex.quote(_syncer_unit_name(plan)) - eval_checkpoint = plan.get("eval_checkpoint") - checkpoint_check = "" - if eval_checkpoint is not None: - checkpoint_check = f"""test -f "$RUN/state/state.ckpt" -test "$(stat -c %s "$RUN/state/state.ckpt")" = {eval_checkpoint['size_bytes']} -printf '%s %s\n' {shlex.quote(eval_checkpoint['sha256'])} "$RUN/state/state.ckpt" | sha256sum --check - -""" return f"""set -euo pipefail {_remote_vars(plan)} -{checkpoint_check}command -v systemctl >/dev/null +command -v systemctl >/dev/null command -v systemd-run >/dev/null test -d /run/systemd/system UNIT={unit} @@ -2350,30 +2117,6 @@ def _learner_argv(plan: dict[str, Any], learner_id: int) -> list[str]: values.extend(("--wandb-mode", str(learner.get("wandb_mode", "online")))) if learner.get("wandb_entity"): values.extend(("--wandb-entity", str(learner["wandb_entity"]))) - evaluation = learner.get("evaluation") - if evaluation is not None: - values.append("--eval-only") - values.extend( - ( - "--eval-dataset-name", - evaluation["dataset_name"], - "--eval-data-sha256", - evaluation["data_sha256"], - "--eval-interval", - str(evaluation["interval"]), - "--eval-samples-per-prompt", - str(evaluation["samples_per_prompt"]), - ) - ) - for flag, name in ( - ("--eval-temperature", "temperature"), - ("--eval-top-p", "top_p"), - ("--eval-max-prompt-len", "max_prompt_len"), - ("--eval-max-response-len", "max_response_len"), - ("--eval-max-context-len", "max_context_len"), - ): - if evaluation.get(name) is not None: - values.extend((flag, str(evaluation[name]))) if learner.get("rollout_model") is not None: values.extend(("--rollout-model", learner["rollout_model"])) values.extend( @@ -2456,15 +2199,9 @@ def _learner_argv(plan: dict[str, Any], learner_id: int) -> list[str]: values.extend(str(port) for port in learner["session_server_port"]) if learner.get("tito_model"): values.extend(("--tito-model", learner["tito_model"])) - if learner.get("codex_backend_profile"): - values.extend( - ("--codex-backend-profile", learner["codex_backend_profile"]) - ) if learner.get("tito_allowed_append_roles"): values.append("--tito-allowed-append-roles") values.extend(learner["tito_allowed_append_roles"]) - if learner.get("codex_exclude_reasoning_from_loss"): - values.append("--codex-exclude-reasoning-from-loss") if learner["trust_remote_code"]: values.append("--trust-remote-code") return values @@ -3030,37 +2767,9 @@ def _collect_event_evidence( _atomic_bytes(destination, payload) -def _parse_syncer_lifecycle_evidence(raw: str) -> dict[str, int]: - if len(raw.encode()) > 4096: - raise HarnessError("syncer lifecycle evidence is unexpectedly large") - try: - evidence = json.loads(raw) - except (TypeError, json.JSONDecodeError) as error: - raise HarnessError("syncer lifecycle evidence is malformed") from error - if not isinstance(evidence, dict) or set(evidence) != { - "schema", - "resume_count", - "final_ack_count", - }: - raise HarnessError("syncer lifecycle evidence has unapproved fields") - if ( - evidence.get("schema") != 1 - or any( - isinstance(evidence.get(name), bool) - or not isinstance(evidence.get(name), int) - or not 0 <= evidence[name] <= 1_000_000 - for name in ("resume_count", "final_ack_count") - ) - ): - raise HarnessError("syncer lifecycle evidence has invalid counters") - return evidence - - def _collect_capacity_preflight(plan: dict[str, Any], artifacts: Path) -> None: """Refuse a large training transfer before it can fill the controller disk.""" - if plan["learner"].get("evaluation") is not None: - return total_bytes = (len(plan["islands"]) + 1) * 64 * 1024 * 1024 if plan["learner"].get("sync_preset", "strict-avg") != "decoupled": for learner_id, island in enumerate(plan["islands"]): @@ -3107,16 +2816,9 @@ def collect(plan_path: str | Path) -> Path: artifacts = plan_file.parent / "artifacts" artifacts.mkdir(parents=True, exist_ok=True) _collect_capacity_preflight(plan, artifacts) - eval_only = plan["learner"].get("evaluation") is not None sync_preset = plan["learner"].get("sync_preset", "strict-avg") learner_event_mode = ( - "eval-learner" - if eval_only - else ( - "decoupled-learner" - if sync_preset == "decoupled" - else "strict-learner" - ) + "decoupled-learner" if sync_preset == "decoupled" else "strict-learner" ) for learner_id, island in enumerate(plan["islands"]): destination = artifacts / f"island-{learner_id}" @@ -3146,7 +2848,7 @@ def collect(plan_path: str | Path) -> Path: learner_event_mode, output / "events.jsonl", ) - if not eval_only and sync_preset != "decoupled": + if sync_preset != "decoupled": audit = destination / "audit" audit.mkdir(parents=True, exist_ok=True) audit_includes = [ @@ -3168,34 +2870,17 @@ def collect(plan_path: str | Path) -> Path: ) syncer = artifacts / "syncer" syncer.mkdir(parents=True, exist_ok=True) - if eval_only: - lifecycle = _ssh( - plan, - _syncer_host(plan), - f"{_remote_vars(plan)}\n" - "python3 -c " - f"{shlex.quote(_SYNCER_LIFECYCLE_EVIDENCE)} " - '"$RUN/state/syncer.log"', - capture=True, - ) - evidence = _parse_syncer_lifecycle_evidence(lifecycle.stdout) - _atomic_bytes( - syncer / "lifecycle-evidence.json", - (_canonical_json(evidence) + "\n").encode(), - ) - syncer_includes = ("--include=/syncer.exit",) - else: - _collect_event_evidence( - plan, - _syncer_host(plan), - "state/events.jsonl", - "syncer", - syncer / "events.jsonl", - ) - syncer_includes = ( - "--include=/state.ckpt", - "--include=/syncer.exit", - ) + _collect_event_evidence( + plan, + _syncer_host(plan), + "state/events.jsonl", + "syncer", + syncer / "events.jsonl", + ) + syncer_includes = ( + "--include=/state.ckpt", + "--include=/syncer.exit", + ) _run( [ "rsync", @@ -3657,118 +3342,6 @@ def _container_succeeded(inspection: Any) -> bool: ) -def _read_successful_syncer_exit(path: Path) -> None: - try: - if path.stat().st_size > 4096: - raise HarnessError("eval syncer exit record is unexpectedly large") - lines = path.read_text(encoding="utf-8").splitlines() - pairs = [line.split("=", 1) for line in lines] - except OSError as error: - raise HarnessError("missing eval syncer exit record") from error - if any(len(pair) != 2 for pair in pairs) or len({pair[0] for pair in pairs}) != len( - pairs - ): - raise HarnessError("eval syncer exit record is malformed") - fields = dict(pairs) - if ( - fields.get("result") != "success" - or fields.get("exit_code") != "exited" - or fields.get("exit_status") != "0" - ): - raise HarnessError("eval syncer did not exit successfully after FINAL_ACK") - - -def _eval_metric(event: dict[str, Any], name: str) -> float: - value = event.get(name) - if ( - not isinstance(value, (int, float)) - or isinstance(value, bool) - or not math.isfinite(value) - or not 0.0 <= value <= 1.0 - ): - raise HarnessError(f"eval event has an invalid {name}") - return float(value) - - -def _verify_eval( - plan: dict[str, Any], artifacts: Path -) -> tuple[str, int, int, float, float]: - learner = plan["learner"] - evaluation = learner["evaluation"] - expected_version = learner["global_rounds"] - expected_samples = ( - evaluation["prompt_count"] * evaluation["samples_per_prompt"] - ) - if (artifacts / "syncer" / "state.ckpt").exists(): - raise HarnessError("eval artifact collection must not contain state.ckpt") - _read_successful_syncer_exit(artifacts / "syncer" / "syncer.exit") - try: - lifecycle = _parse_syncer_lifecycle_evidence( - (artifacts / "syncer" / "lifecycle-evidence.json").read_text( - encoding="utf-8" - ) - ) - except OSError as error: - raise HarnessError("missing eval syncer lifecycle evidence") from error - if lifecycle["resume_count"] != 1: - raise HarnessError("eval syncer lacks one checkpoint resume event") - if lifecycle["final_ack_count"] != _learner_count(plan): - raise HarnessError("eval syncer lacks one FINAL_ACK per learner") - - hashes = set() - results = [] - pass_at_1 = [] - for learner_id in range(_learner_count(plan)): - events = _json_lines(_event_path(artifacts, learner_id)) - apply_indices = [ - index - for index, event in enumerate(events) - if event.get("event") == "rl_policy_apply" - and event.get("policy_version") == expected_version - ] - eval_indices = [ - index - for index, event in enumerate(events) - if event.get("event") == "rl_eval_result" - and event.get("dataset_name") == evaluation["dataset_name"] - ] - if len(apply_indices) != 1 or len(eval_indices) != 1: - raise HarnessError( - f"eval learner {learner_id} lacks one final apply/result event" - ) - if apply_indices[0] >= eval_indices[0]: - raise HarnessError( - f"eval learner {learner_id} evaluated before final policy apply" - ) - apply_event = events[apply_indices[0]] - eval_event = events[eval_indices[0]] - policy_hash = apply_event.get("sync/global_policy_hash") - if not isinstance(policy_hash, str) or not _SHA256.fullmatch(policy_hash): - raise HarnessError(f"eval learner {learner_id} has no final policy hash") - if ( - apply_event.get("island_id") != learner_id - or eval_event.get("island_id") != learner_id - or eval_event.get("rollout_id") != 0 - or eval_event.get("policy_version") != expected_version - or eval_event.get("sample_count") != expected_samples - ): - raise HarnessError( - f"eval learner {learner_id} result has the wrong policy/sample identity" - ) - hashes.add(policy_hash) - results.append(_eval_metric(eval_event, "rl/eval/result")) - pass_at_1.append(_eval_metric(eval_event, "rl/eval/pass_at_1")) - if len(hashes) != 1: - raise HarnessError("eval learners applied different terminal policies") - return ( - evaluation["dataset_name"], - evaluation["prompt_count"], - evaluation["samples_per_prompt"], - statistics.fmean(results), - statistics.fmean(pass_at_1), - ) - - def verify(plan_path: str | Path, export_dir: str | None = None) -> None: plan_file, plan = load_plan(plan_path) artifacts = plan_file.parent / "artifacts" @@ -3790,15 +3363,6 @@ def verify(plan_path: str | Path, export_dir: str | None = None) -> None: raise HarnessError( f"island {learner_id} node {node_id} did not exit successfully" ) - if plan["learner"].get("evaluation") is not None: - if export_dir: - raise HarnessError("eval-only verification cannot export a checkpoint") - dataset, prompts, samples, result, pass_at_1 = _verify_eval(plan, artifacts) - print( - f"verified terminal eval {dataset}: {prompts} prompts x {samples} " - f"sample(s), result={result:.6f}, pass@1={pass_at_1:.6f}" - ) - return from ..export import parse_checkpoint checkpoint_path = artifacts / "syncer" / "state.ckpt" @@ -3864,19 +3428,6 @@ def build_parser() -> argparse.ArgumentParser: prepare_parser.add_argument("--remote-root", default=DEFAULT_REMOTE_ROOT) prepare_parser.add_argument("--remote-env-file", default=None) prepare_parser.add_argument("--final-ack-timeout-s", type=int, default=None) - prepare_parser.add_argument("--eval-checkpoint-host", default=None) - prepare_parser.add_argument("--eval-checkpoint-path", default=None) - prepare_parser.add_argument("--eval-checkpoint-sha256", default=None) - prepare_parser.add_argument("--eval-checkpoint-size-bytes", type=int, default=None) - prepare_parser.add_argument("--eval-checkpoint-global-step", type=int, default=None) - prepare_parser.add_argument("--eval-dataset-name", default=None) - prepare_parser.add_argument("--eval-interval", type=int, default=None) - prepare_parser.add_argument("--eval-samples-per-prompt", type=int, default=None) - prepare_parser.add_argument("--eval-temperature", type=float, default=None) - prepare_parser.add_argument("--eval-top-p", type=float, default=None) - prepare_parser.add_argument("--eval-max-prompt-len", type=int, default=None) - prepare_parser.add_argument("--eval-max-response-len", type=int, default=None) - prepare_parser.add_argument("--eval-max-context-len", type=int, default=None) prepare_parser.add_argument("--model-manifest-sha256", default=None) prepare_parser.add_argument("--tms-preload-patch", default=None) prepare_parser.add_argument( From fbefc71bf22a328446d34bc359ea1b17fd817a15 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:19:40 +0000 Subject: [PATCH 004/233] Thread preflight on by default for eval launches; log and record every loosened value The C1b eval scripts passed --preflight-threads off. They now keep the default (wait) and may loosen it only with explicit flags (cfg preflight_args / PREFLIGHT env); the launcher prints each loosened value with its default and records it under preflight.threads.relaxed. Co-Authored-By: Claude Opus 5.5 --- tests/test_launch_preflight_guards.py | 27 ++++++++++++++++++++++++++- tools/eval/miles.py | 8 +++++--- tools/eval/miles_head.sh | 7 +++++-- yeto/launch_preflight.py | 27 +++++++++++++++++++++++++-- 4 files changed, 61 insertions(+), 8 deletions(-) diff --git a/tests/test_launch_preflight_guards.py b/tests/test_launch_preflight_guards.py index ab9be3986..c8d9ee6fe 100644 --- a/tests/test_launch_preflight_guards.py +++ b/tests/test_launch_preflight_guards.py @@ -143,7 +143,32 @@ def test_threads_off_skips_and_warns(): out = io.StringIO() rec = lp.thread_preflight(_targs(preflight_threads="off"), counter=lambda: pytest.fail("off must not read"), out=out) - assert rec == {"mode": "off", "disabled": True} and "OFF" in out.getvalue() + assert rec == {"mode": "off", "disabled": True, "relaxed": {"mode": ["off", "wait"]}} + assert "OFF" in out.getvalue() and "RELAXED" in out.getvalue() + + +def test_threads_default_is_on_and_not_relaxed(): + out = io.StringIO() + rec = lp.thread_preflight(_targs(), counter=_counter([100]), out=out) + assert rec["mode"] == "wait" and rec["relaxed"] == {} and "RELAXED" not in out.getvalue() + + +def test_threads_loosened_values_are_logged_and_recorded(): + out = io.StringIO() + rec = lp.thread_preflight(_targs(preflight_thread_start=2950, preflight_thread_hard=3200), + counter=_counter([2900]), out=out) + assert rec["relaxed"] == {"start": [2950, 2800], "hard": [3200, 3000]} + assert "start 2950 (default 2800)" in out.getvalue() and "hard 3200 (default 3000)" in out.getvalue() + # a stricter setting is not a relaxation + assert lp.thread_relaxations(lp.ThreadConfig(mode="error", start=100, hard=200)) == {} + + +def test_eval_launch_args_keep_the_thread_preflight_on(): + from pathlib import Path + from tools.eval import miles + root = Path(__file__).resolve().parents[1] + assert "--preflight-threads off" not in (root / "tools/eval/miles_head.sh").read_text() + assert "--preflight-threads off" not in Path(miles.__file__).read_text() def test_thread_message_text(): diff --git a/tools/eval/miles.py b/tools/eval/miles.py index 3f9538e69..401e3fae7 100644 --- a/tools/eval/miles.py +++ b/tools/eval/miles.py @@ -30,9 +30,11 @@ def _gpu(cfg: dict, n: int) -> str: def _common(cfg: dict, run_id: str, prefix: str) -> str: - # --preflight-threads off: Modal islands put no SkyPilot threads on this host; the scripts - # check the host thread count (< 3000) before the launch instead - return (f"launch --training-mode rl --rl-engine ports --preflight-threads off " + # the launcher thread preflight stays on (default wait); a cell may loosen it only with + # explicit flags in cfg["preflight_args"] (e.g. "--preflight-thread-start 2950"), which the + # launcher prints and records in the run manifest + preflight = cfg.get("preflight_args", "") + return (f"launch --training-mode rl --rl-engine ports {preflight + ' ' if preflight else ''}" f"--modal-gpu-exact --modal-retries 0 --modal-cpu 8 " f"--modal-memory 64 --modal-tape-volume {TAPE_VOLUME} --controller local --cluster-prefix {prefix} " f"--tuning lora --lora-r {cfg['lora_r']} --rl-observe-timeline {MODEL} {DATA} " diff --git a/tools/eval/miles_head.sh b/tools/eval/miles_head.sh index dac653735..e4ac5b032 100755 --- a/tools/eval/miles_head.sh +++ b/tools/eval/miles_head.sh @@ -8,6 +8,9 @@ set -u : "${P:?}" "${R:?}" "${REPO:?}" "${PY:?}" "${SKY:?}" HARD=${HARD:-2700}; WD=$(( HARD + 300 )); STEPS=${STEPS:-4}; SROUNDS=$(( 2 * STEPS + 6 )); GPU=${GPU:-a10g} +# launcher thread preflight stays on; loosen only with explicit flags, e.g. +# PREFLIGHT="--preflight-thread-start 2950" (the launcher logs and records the loosened values) +PREFLIGHT=${PREFLIGHT:-} SHA=${SHA:-HEAD}; MODAL="$PY -m modal" APP=yeto-$P; SP=$P-s; SAPP=yeto-$SP; TAPEVOL=yeto-event-tapes; STOREVOL=${STOREVOL:-yeto-eval-store}; HEAD_REGION=nebius/eu-north1 export CARGO_HOME=/home/michael/work/gpu-default-modal/cargo RUSTUP_HOME=/home/michael/work/gpu-default-modal/rustup @@ -30,8 +33,8 @@ LORA="--tuning lora --lora-r 16" MODALX="--modal-gpu-exact --modal-retries 0 --modal-cpu 8 --modal-memory 64 --modal-tape-volume $TAPEVOL" STORE="--rl-checkpoint-store modal-volume://$STOREVOL/$P" COMMON="--training-mode rl --rl-engine ports $MODALX $MODEL $DATA $LORA $STORE --rl-algorithm-spec examples/rl_algorithms/decoupled-ppo-c2.json --rl-max-policy-age 2 --rl-observe-timeline --reward-function gsm8k_reward:score --fragments 1 --pipeline 1 --rollout-batch-size 32 --n-samples-per-prompt 4 --rollout-max-response-len 1024 --seq-len 2048 --inner-lr 1e-5 --seed 17 --apply-chat-template-kwargs '{\"enable_thinking\": false}' --trust-remote-code" -TARGS="launch --controller head --syncer-region $HEAD_REGION --preflight-threads off --rl-stall-timeout 2400 --on-demand --no-island-relaunch --modal-timeout-s $HARD --cluster-prefix $P --gpu modal:1x$GPU,modal:1x$GPU --rl-island-scheduling elastic --rl-sample-role trainer --rl-sample-wait-s 900 $COMMON --total-steps $STEPS" -SARGS="launch --preflight-threads off --controller local --cluster-prefix $SP --gpu modal:1x$GPU --rl-single-island-no-sync --rl-sample-role sampler --rl-sampler-max-backlog-rounds 4 $COMMON --total-steps $SROUNDS" +TARGS="launch --controller head --syncer-region $HEAD_REGION $PREFLIGHT --rl-stall-timeout 2400 --on-demand --no-island-relaunch --modal-timeout-s $HARD --cluster-prefix $P --gpu modal:1x$GPU,modal:1x$GPU --rl-island-scheduling elastic --rl-sample-role trainer --rl-sample-wait-s 900 $COMMON --total-steps $STEPS" +SARGS="launch $PREFLIGHT --controller local --cluster-prefix $SP --gpu modal:1x$GPU --rl-single-island-no-sync --rl-sample-role sampler --rl-sampler-max-backlog-rounds 4 $COMMON --total-steps $SROUNDS" git -C $REPO rev-parse $SHA > $R/yeto_sha.txt; echo "$TARGS" > $R/args.txt; echo "$SARGS" > $R/sargs.txt echo "$APP $SAPP $HEAD_CL" > $R/resources.txt eval "$(/usr/bin/python3 - <<'PYC' diff --git a/yeto/launch_preflight.py b/yeto/launch_preflight.py index 539b07a07..c173c4da9 100644 --- a/yeto/launch_preflight.py +++ b/yeto/launch_preflight.py @@ -185,6 +185,24 @@ def thread_message(count: ThreadCount, cfg: ThreadConfig, headline: str) -> str: ]) +def thread_relaxations(cfg: ThreadConfig) -> dict: + """The thread settings that are looser than the defaults, as ``{name: [value, default]}``. + + The check is on by default; a run may only loosen it with an explicit flag, and the + launcher prints and records every loosened value (empty dict: defaults or stricter).""" + + relaxed = {} + if cfg.mode == "off": + relaxed["mode"] = ["off", DEFAULT_THREAD_MODE] + if cfg.start > DEFAULT_THREAD_START: + relaxed["start"] = [cfg.start, DEFAULT_THREAD_START] + if cfg.hard > DEFAULT_THREAD_HARD: + relaxed["hard"] = [cfg.hard, DEFAULT_THREAD_HARD] + if cfg.mode == "wait" and cfg.wait_s > DEFAULT_THREAD_WAIT_S: + relaxed["wait_s"] = [cfg.wait_s, DEFAULT_THREAD_WAIT_S] + return relaxed + + def thread_preflight(args, *, counter: Callable[[], ThreadCount] | None = None, clock: Callable[[], float] = time.monotonic, sleep: Callable[[float], None] = time.sleep, @@ -197,10 +215,15 @@ def thread_preflight(args, *, counter: Callable[[], ThreadCount] | None = None, out = sys.stderr if out is None else out cfg = ThreadConfig.from_args(args) + relaxed = thread_relaxations(cfg) + if relaxed: + print("[preflight] WARNING: thread preflight RELAXED by explicit flags: " + + ", ".join(f"{k} {v[0]} (default {v[1]})" for k, v in relaxed.items()), + file=out, flush=True) if cfg.mode == "off": print("[preflight] WARNING: thread preflight is OFF (--preflight-threads off); the " "s1run thread guard may tear down this run", file=out, flush=True) - return {"mode": "off", "disabled": True} + return {"mode": "off", "disabled": True, "relaxed": relaxed} counter = counter or count_threads started = clock() count = counter() @@ -233,7 +256,7 @@ def thread_preflight(args, *, counter: Callable[[], ThreadCount] | None = None, f"(SkyPilot API server {count.sky_api if count.sky_api_found else 'not found'})", file=out, flush=True) return {"mode": cfg.mode, "start": cfg.start, "hard": cfg.hard, "wait_s": cfg.wait_s, - "waited": waited, "readings": readings, **count.to_dict()} + "waited": waited, "readings": readings, "relaxed": relaxed, **count.to_dict()} # --------------------------------------------------------------------------- From 59ab2b492d934e204783889282d4bdab521a3e13 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:28:54 +0000 Subject: [PATCH 005/233] Samplers draw disjoint data slots; trainer ends FINISHED on an exhausted buffer; launcher follows rescheduled Modal containers - sample_roles.DataSlot: sampler i of n keeps prompt rows r % n == i. The launcher sets --rl-sampler-data-slot per pair sampler island; the Miles buffer writer splits its RolloutDataSource before the first draw and any resume seek; the engine sampler spec takes prompts.data_slot (stage 1 sets it with several sampler nodes). - BatchReader raises SampleBufferExhausted when the buffer ended with no trainable group; IslandDriver ends the run FINISHED before that round (event rl_buffer_exhausted), except under a strict roster. - ContainerLogFollower: when a new container is expected, the launcher lists the app's containers and streams each new one's log ('modal container logs'), with a warning; the call log stream misses it. Co-Authored-By: Claude Opus 5.5 --- tests/test_s26_e9_data_slots.py | 273 ++++++++++++++++++ yeto/launch/modal_islands.py | 17 ++ yeto/launcher.py | 13 +- yeto/modal_runner.py | 101 +++++++ yeto/rl/adapters/miles/buffer_rollout.py | 44 ++- .../rl/adapters/miles/buffer_rollout_miles.py | 6 + yeto/rl/adapters/miles/island_entry.py | 2 +- yeto/rl/engine/bridges.py | 8 + yeto/rl/engine/buffer_payload.py | 24 ++ yeto/rl/engine/driver.py | 13 +- yeto/rl/engine/engine_sampler_entry.py | 26 +- yeto/rl/engine/run_flags.py | 3 + yeto/rl/engine/sample_roles.py | 51 ++++ yeto/rl/stage1.py | 12 +- 14 files changed, 575 insertions(+), 18 deletions(-) create mode 100644 tests/test_s26_e9_data_slots.py diff --git a/tests/test_s26_e9_data_slots.py b/tests/test_s26_e9_data_slots.py new file mode 100644 index 000000000..8609d4d88 --- /dev/null +++ b/tests/test_s26_e9_data_slots.py @@ -0,0 +1,273 @@ +"""S26 E9.3: disjoint sampler data slots, a normal trainer end on an exhausted +buffer, and the launcher following a rescheduled Modal container's log.""" + +from __future__ import annotations + +import io +import json +import random +import sys +from types import SimpleNamespace + +import pytest +import torch + +from yeto.rl.engine import sample_roles +from yeto.rl.engine.buffer_payload import SampleBufferExhausted, SampleBufferRolloutError, is_buffer_exhausted +from yeto.rl.engine.engine_sampler_entry import JsonlPrompts, SamplerSpecError +from yeto.rl.engine.sample_roles import DataSlot + + +# ------------------------------------------------------------------ DataSlot +def test_data_slot_parse_and_rows(): + assert DataSlot.parse(None) == DataSlot() and DataSlot().whole + slot = DataSlot.parse("1/3") + assert (slot.index, slot.count, slot.env()) == (1, 3, "1/3") + assert slot.rows(list(range(10))) == [1, 4, 7] + for bad in ("3/3", "-1/2", "1/0", "2"): + with pytest.raises(ValueError): + DataSlot.parse(bad) + with pytest.raises(ValueError, match="empty"): + DataSlot(2, 3).rows([0, 1]) + + +def test_slots_of_all_samplers_cover_the_data_once(): + rows = list(range(17)) + drawn = [r for k in range(3) for r in DataSlot(k, 3).rows(rows)] + assert sorted(drawn) == rows + + +def test_pair_island_args_give_each_sampler_its_slot_and_flag(): + pair = SimpleNamespace(rl_island_sample_role=["0:sampler", "1:sampler", "2:trainer"], + rl_sampler_max_backlog_rounds=None, rl_sample_wait_s=None, total_steps=8, + rl_sampler_spot=False) + a0, a1, t = (sample_roles.island_args(pair, m, 3) for m in range(3)) + assert (a0.rl_sampler_data_slot, a1.rl_sampler_data_slot) == ("0/2", "1/2") + assert getattr(t, "rl_sampler_data_slot", None) is None + assert " --rl-sampler-data-slot 1/2" in sample_roles.island_flags(a1) + one = SimpleNamespace(rl_island_sample_role=["0:sampler", "1:trainer"], rl_sampler_max_backlog_rounds=None, + rl_sample_wait_s=None, total_steps=8, rl_sampler_spot=False) + assert sample_roles.island_args(one, 0, 2).rl_sampler_data_slot is None # one sampler: the whole data + + +def test_data_slot_flag_on_both_layers(): + import argparse + + from yeto.rl.engine.run_flags import CLI, ISLAND, flags_for + + for layer in (CLI, ISLAND): + parser = argparse.ArgumentParser() + for flag in flags_for(layer): + if flag.options[0] == "--rl-sampler-data-slot": + parser.add_argument(*flag.options, **dict(flag.kwargs)) + assert parser.parse_args(["--rl-sampler-data-slot", "0/2"]).rl_sampler_data_slot == "0/2" + + +# ------------------------------------------------------------ engine sampler +def _prompts(tmp_path, n=10): + path = tmp_path / "p.jsonl" + path.write_text("".join(json.dumps({"q": f"q{i}", "a": str(i)}) + "\n" for i in range(n))) + return path + + +def test_engine_sampler_prompts_of_two_slots_are_disjoint(tmp_path): + path = _prompts(tmp_path) + ids = [] + for slot in ("0/2", "1/2"): + prompts = JsonlPrompts(path, input_key="q", label_key="a", groups_per_round=2, data_slot=slot) + ids.append([item.prompt_id for r in range(2) for item in prompts(r)]) + assert ids == [["row0", "row2", "row4", "row6"], ["row1", "row3", "row5", "row7"]] + whole = JsonlPrompts(path, input_key="q", label_key="a", groups_per_round=2) + assert [i.prompt_id for i in whole(1)] == ["row2", "row3"] # unchanged without a slot + with pytest.raises(SamplerSpecError, match="data_slot"): + JsonlPrompts(path, input_key="q", label_key="a", groups_per_round=2, data_slot="2/2") + + +def test_stage1_sampler_spec_carries_the_slot_only_with_several_samplers(): + from yeto.rl import stage1 + + src = open(stage1.__file__).read() + assert "slot=DataSlot(samplers.index(node)" in src and '"data_slot": slot.env()' in src + + +# --------------------------------------------------------------- Miles source +class _FakeDataset: + """The fields and shuffle of Miles' ``miles.utils.data.Dataset``.""" + + def __init__(self, rows, seed=3): + self.origin_samples = list(rows) + self.samples = list(rows) + self.seed, self.epoch_id = seed, -1 + + def shuffle(self, new_epoch_id): + if self.epoch_id == new_epoch_id: + return + random.seed(self.seed + new_epoch_id) + permutation = list(range(len(self.samples))) + random.shuffle(permutation) + self.samples = [self.origin_samples[i] for i in permutation] + self.epoch_id = new_epoch_id + + def __len__(self): + return len(self.samples) + + +def _source(rows, shuffle): + source = SimpleNamespace(dataset=_FakeDataset(rows), epoch_id=0, sample_offset=0, sample_group_index=0, + sample_index=0) + if shuffle: + source.dataset.shuffle(0) + return source + + +@pytest.mark.parametrize("shuffle", [False, True]) +def test_miles_data_source_slots_are_disjoint(shuffle): + from yeto.rl.adapters.miles.buffer_rollout import apply_data_slot + + rows = list(range(20)) + drawn = [] + for k in range(2): + source = _source(rows, shuffle) + assert apply_data_slot(source, DataSlot(k, 2), shuffle=shuffle) == 10 + assert sorted(source.dataset.samples) == rows[k::2] + drawn.append(set(source.dataset.samples)) + assert not drawn[0] & drawn[1] and drawn[0] | drawn[1] == set(rows) + + +def test_miles_data_slot_refuses_after_the_first_draw(): + from yeto.rl.adapters.miles.buffer_rollout import apply_data_slot + + source = _source(list(range(4)), False) + source.sample_offset = 2 + with pytest.raises(SampleBufferRolloutError, match="before the first draw"): + apply_data_slot(source, DataSlot(0, 2), shuffle=False) + whole = _source(list(range(4)), False) + assert apply_data_slot(whole, DataSlot(), shuffle=False) == 4 + + +def test_island_env_exports_the_slot_for_a_sampler_only(): + from yeto.rl.adapters.miles import buffer_rollout as br + + env = br.island_env("sampler", store_root="/s", sampler_id="island0", shards=["island0", "island1"], + data_slot="0/2") + assert env[br.DATA_SLOT_ENV] == "0/2" + with pytest.raises(ValueError, match="only a sampler"): + br.island_env("trainer", store_root="/s", max_policy_age=2, data_slot="0/2") + + +# ------------------------------------------------------- exhausted buffer end +def test_is_buffer_exhausted_follows_wrappers(): + inner = SampleBufferExhausted("ended") + assert is_buffer_exhausted(inner) + try: + try: + raise inner + except SampleBufferExhausted as e: + raise RuntimeError("ray task failed") from e + except RuntimeError as outer: + assert is_buffer_exhausted(outer) + ray_like = RuntimeError("RayTaskError") + ray_like.cause = inner + assert is_buffer_exhausted(ray_like) + dual = type("SampleBufferExhausted", (RuntimeError,), {})("dual class of another process") + assert is_buffer_exhausted(dual) + assert not is_buffer_exhausted(SampleBufferRolloutError("other")) + assert not is_buffer_exhausted(None) + + +def test_batch_reader_raises_exhausted_on_an_ended_empty_buffer(): + from yeto.rl.adapters.miles.buffer_rollout import BatchReader + + buffer = SimpleNamespace(read_cursor=5, data_ended=True, select=lambda consumer, n: []) + reader = BatchReader(buffer, max_policy_age=2, wait_s=0, poll_s=0) + with pytest.raises(SampleBufferExhausted, match="no trainable group"): + reader.take(3, 4, current_version=3) + + +def _fake_driver(tmp_path, sync, exhaust_at): + from yeto.rl.engine.algorithm import AlgorithmSpec + from yeto.rl.engine.driver import EventTape, IslandDriver + from yeto.rl.engine.fake import FakeEngine, fake_capabilities + + engine = FakeEngine(tensors={"base_model.model.layer.lora_A.weight": torch.zeros(1, 2)}, step_delta=1.0, capability_profile="miles-like") + generate = engine.rollout.generate + + def gen(rollout_id, **kw): + if rollout_id == exhaust_at: + raise RuntimeError("ray task failed") from SampleBufferExhausted("ended") + return generate(rollout_id, **kw) + + engine.rollout.generate = gen + return IslandDriver(learner_id=0, rollout=engine.rollout, trainer=engine.trainer, + policy_state=engine.policy_state, publisher=engine.publisher, + placement=engine.placement, algorithm=AlgorithmSpec(), sync=sync, + events=EventTape(tmp_path / "e.jsonl", 0), capabilities=fake_capabilities("miles-like")) + + +def test_driver_ends_finished_when_the_buffer_is_exhausted(tmp_path): + from yeto.rl.engine.bridges import LocalOnlySync + + final = _fake_driver(tmp_path, LocalOnlySync(5), exhaust_at=3).run() + assert final.policy_version == 3 + events = [json.loads(line) for line in (tmp_path / "e.jsonl").read_text().splitlines()] + names = [e["event"] for e in events] + assert "rl_buffer_exhausted" in names and "island_leave" not in names + assert next(e for e in events if e["event"] == "rl_buffer_exhausted")["rollout_id"] == 3 + + +def test_driver_still_fails_on_other_errors_and_on_a_strict_roster(tmp_path): + from yeto.rl.engine.bridges import LocalOnlySync, StrictAvgSync + + assert StrictAvgSync.ends_on_data_end is False + sync = LocalOnlySync(5) + sync.ends_on_data_end = False + with pytest.raises(RuntimeError, match="ray task failed"): + _fake_driver(tmp_path, sync, exhaust_at=2).run() + + +# ------------------------------------------------- rescheduled container log +def test_container_log_follower_streams_new_containers_only(): + from yeto.modal_runner import ContainerLogFollower + + lists = iter([{"ta-a"}, {"ta-a"}, {"ta-a", "ta-b"}, {"ta-b"}]) + streamed = [] + + def stream(cid): + streamed.append(cid) + yield "[modal-island 1] rank 0 container ta-b\n" + yield "round 2 done\nround 3 done\n" + + class Collector: + def __init__(self): + self.lines = [] + + def feed(self, line): + self.lines.append(line) + + collector = Collector() + out = io.StringIO() + follower = ContainerLogFollower(lambda: next(lists), stream, expected=1, + route=lambda learner: ("island1", collector) if learner == 1 else None, + out=out) + assert follower.poll_once() == [] and follower.baseline == {"ta-a"} + assert follower.poll_once() == [] + assert follower.poll_once() == ["ta-b"] + follower.threads[0].join(5) + assert follower.poll_once() == [] # followed once + text = out.getvalue() + assert "new Modal container ta-b" in text + assert "[island1] round 2 done" in text and "[island1] round 3 done" in text + assert streamed == ["ta-b"] and collector.lines[-1] == "round 2 done\nround 3 done" + + +def test_container_log_follower_only_when_new_containers_are_expected(): + from yeto.launch.modal_islands import container_log_follower_for + + ops = SimpleNamespace(list_container_ids=lambda: set(), stream_container_logs=lambda cid: iter(())) + cfgs = {"island0": SimpleNamespace(num_nodes=1, learner_id=0)} + assert container_log_follower_for(None, ops, cfgs, None, poller=object(), start=False) is None + assert container_log_follower_for(None, None, cfgs, None, poller=None, start=False) is None + follower = container_log_follower_for(None, ops, cfgs, None, poller=None, start=False) + assert follower is not None and follower.expected_count == 1 + assert container_log_follower_for(None, SimpleNamespace(list_container_ids=lambda: set()), cfgs, None, + poller=None, start=False) is None diff --git a/yeto/launch/modal_islands.py b/yeto/launch/modal_islands.py index a7f12af6f..74649a6dc 100644 --- a/yeto/launch/modal_islands.py +++ b/yeto/launch/modal_islands.py @@ -301,6 +301,23 @@ def container_set_poller_for(args, modal_ops, modal_cfgs, on_change, *, start: b return poller.start() if start else poller +def container_log_follower_for(args, modal_ops, modal_cfgs, route, *, poller, start: bool = True): + """The follower of rescheduled containers' logs (modal_runner.ContainerLogFollower), + only when a new container is expected (no failing container guard: ``poller`` + is None) and the ops can stream one container. None otherwise.""" + if poller is not None or modal_ops is None or not modal_cfgs: + return None + if not (callable(getattr(modal_ops, "list_container_ids", None)) + and callable(getattr(modal_ops, "stream_container_logs", None))): + return None + from ..modal_runner import ContainerLogFollower + + expected = sum(max(1, int(getattr(cfg, "num_nodes", 1) or 1)) for cfg in modal_cfgs.values()) + follower = ContainerLogFollower(modal_ops.list_container_ids, modal_ops.stream_container_logs, + expected, route=route) + return follower.start() if start else follower + + def effective_recover_timeout(args) -> float: """The fleet controller's learner relaunch budget. diff --git a/yeto/launcher.py b/yeto/launcher.py index 0abd2f4b5..6ee357ca9 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -60,7 +60,7 @@ ) from .launch.modal_islands import ( _modal_tape_run_dir, _tail_modal, _verify_modal_app_stopped, build_modal_island_config, - container_set_poller_for, effective_recover_timeout, pull_modal_tapes, + container_log_follower_for, container_set_poller_for, effective_recover_timeout, pull_modal_tapes, require_modal_for_gpu_exact, ) from .launch.fleet import ( @@ -618,6 +618,17 @@ def container_changed() -> bool: if container_poller is not None: container_guards.append(container_poller) + def route_container(learner_id: int): + # the island of a rescheduled container's lines (unique learner id only: + # the islands of a pair all run as learner 0) + names = [n for n, cfg in modal_cfgs.items() if int(cfg.learner_id) == int(learner_id)] + if len(names) != 1: + return None + return names[0], (event_collectors.get(names[0]) if names[0] in echo_names else None) + + # daemon threads; each container's log ends when the app stops + container_log_follower_for(args, modal_ops, modal_cfgs, route_container, poller=container_poller) + from . import runs from .modal_runner import RoutingOps diff --git a/yeto/modal_runner.py b/yeto/modal_runner.py index b25be3c4f..19ba5353c 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -1194,6 +1194,19 @@ def stream_logs(self, call_id: str): for entry in modal.FunctionCall.from_id(call_id).logs.stream(): yield getattr(entry, "message", str(entry)) + def stream_container_logs(self, container_id: str): + """The log lines of one container (``modal container logs``; follows while it + runs). The function-call stream of :meth:`stream_logs` does not deliver the + lines of a container that Modal started after a reschedule (S25 C1).""" + proc = subprocess.Popen([sys.executable, "-m", "modal", "container", "logs", str(container_id)], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + try: + assert proc.stdout is not None + yield from proc.stdout + finally: + if proc.poll() is None: + proc.terminate() + CONTAINER_LINE_RE = re.compile(r"\[modal-island (\d+)\] rank (\d+) container (\S+)") @@ -1291,6 +1304,94 @@ def stop(self) -> None: self._stop.set() +class ContainerLogFollower: + """Follows the logs of the Modal containers the call log streams miss (S26 E9.3). + + After ``modal container stop`` (or a real reclaim) Modal runs the island's + input again in a new container, and the function-call log stream never + delivers that container's lines (S19, S25 C1), so the launcher showed only + the first container. Every ``interval_s`` this lists the app's running + containers; the first ``expected`` ids are the baseline (the call streams + show them); for every later id it prints one launcher warning and streams + that container's own log (``stream_container(id)``) on a daemon thread. + + Each line is printed with a ``[ 对 ，容差未标定` on that connection only. The syncer does not exit and the admitted islands keep running. H100 and H200 are different groups. - Phase 1 rule: groups must be identical. `CARD_PAIR_TOLERANCE` in the catalog is @@ -342,7 +340,7 @@ only those (other fragments are stored with numel 0 at version 0). for every shard, takes each version from its owner, and ACKs each shard's manifest separately. The terminal cut is the union. - `N = 1` (default) is the unsharded syncer, unchanged. Strict policy - sweeps, elastic scheduling and learner-budget cutoff refuse `N > 1`. + sweeps, dynamic membership and learner-budget cutoff refuse `N > 1`. ## bf16 pushes over an f32 session (dtype 4) @@ -526,15 +524,75 @@ Every merge appends event-tape fields for protocol/delta semantics, attempt, launch base version, frozen members and generations, responders, staleness, weights, counters, quorum/grace timings, and sync latency. -## Elastic island departure (rl-island-graceful-leave) +## Island pool (dynamic membership) -Under `--rl-island-scheduling elastic` an island that stops sends one LEAVE -frame with a one-byte reason code. Only the class of the departure goes on +The syncer has one round loop. Without membership options the roster is +fixed: every island answers every round (strict RL). Any of the options below +makes the islands a dynamic pool; unset options keep the fixed-roster value, +and the syncer and island command lines stay unchanged. + +The launcher takes the `--rl-sync-*` membership options for the strict-avg +(LoRA), decoupled and full-parameter presets; the full-parameter preset needs +`--rl-sync-full-parameter-rounds async`. The launcher refuses them for the +full-parameter sweep, `sao-streaming-full`, `--experimental-rl-sync`, +`--syncer-shards` > 1 and algorithms with a critic. The rounds must be able +to merge without some island (`--rl-sync-quorum-theta`, `--rl-sync-quorum` +below the island count, `--rl-sync-grace-ms` or `--rl-sync-max-base-lag` >= 1). + +| Syncer flag | Launcher option | Effect | +|---|---|---| +| `--quorum-theta F`, `--q-min N` | `--rl-sync-quorum-theta`, `--rl-sync-q-min` | merge once the arrived islands hold F of the capacity the round expects (at least N islands); replaces `--quorum` | +| `--soft-deadline-s S` | `--rl-sync-soft-deadline-s` | merge what arrived after S seconds even below quorum; also caps the grace window | +| `--carry-gamma G` | `--rl-sync-carry-gamma` | weight `G^lag` for a push whose base lags the round; the largest lag is `--max-base-lag` (`--rl-sync-max-base-lag`) | +| `--island-lease-s L` | `--rl-sync-island-lease-s` | an island silent for L seconds leaves the pool and must re-JOIN | +| `--syncer-epoch E` | `--rl-syncer-epoch` | a HELLO with a lower epoch is refused (fencing) | +| `--final-grace-s S` | `--rl-sync-final-grace-s` | after the last step, wait S seconds for the connected islands to ACK the final cut, then finish without the missing ones | +| `--max-learners N` | `--rl-sync-max-learners` | admit island ids 0..N-1 (islands may join late); not for full-parameter, whose island session binds the launched roster | + +Every island in a pool heartbeats every lease/3 (default lease 150 s, +`yeto.rl.pool_client.LeaseHeartbeat`). A lease expiry gets +`MSG_ERROR "island lease expired: re-JOIN required"`; the client redials with +a new connection generation (outside its reconnect budget) and is caught up +to the current versions. Per preset: + +- strict-avg (`PoolIslandClient`, session `PoolAvgSync`) syncs its LoRA vector + as one f32 fragment, answers the newest PULL it has not answered with + `local - base` against the base it holds, and applies the newest broadcast. +- decoupled (`yeto.rl.decoupled.DecoupledRlBridge`) applies the newest + broadcast of each fragment, drops a PULL whose round closed without it (a + newer PULL of the fragment supersedes it), and answers the PULLs the syncer + replays on a new connection. A lost link redials up to 5 times. +- full-parameter, asynchronous rounds (`MilesFullParameterRoleStream`, + `FullParameterSync`) overwrites a fragment with a broadcast that skips + versions (catch-up) and releases its retained pushes of the skipped rounds. + A fresh full-parameter island cannot join a pool after its initial cut yet. At the end the island ACKs the final cut like a strict +island; a push that races the end ends the island normally (`elastic_finished` +on its tape). + +With a pool the syncer also writes island-pool records on its event tape +(`{"kind": ..., "syncer_epoch": ..., "time_unix": ...}`): + +| kind | When | +|---|---| +| `pool_join` | an island connected (`catch_up` true after the run started) | +| `pool_leave` | an island left: `reason` from its LEAVE, `lease_expired`, or `disconnected` | +| `delta_accepted` / `delta_carried_over` / `delta_rejected` | a push was merged on time, merged with `lag` and `discount`, or refused (`carry_lag_exceeded`, `fenced`, `out_of_round`, ...) | +| `outer_step` | a round merged: `arrived`, capacities, `timed_out`, normalized `weights` `[island, base_version, weight]`, `absent` | +| `final_grace_expired` | the final grace window ended with islands missing | + +It keeps `status.json` next to the tape (schema +`yeto.syncer.elastic-status/v1`: state, syncer epoch, outer version, live +islands with capacity and remaining lease). + +### Island departure (LEAVE) + +An island that leaves a pool sends LEAVE (type 27: `u32 learner_id | u8 +reason`) and closes its connection. Only the class of the departure goes on the wire; for a failed departure the exception type and text go to the -island's event tape (`island_leave`; a finished run adds no record), and a -failed send is recorded as `elastic_leave_failed`. The codes -are defined in `yeto/rl/elastic_client.py` (`LeaveReason`) and -`syncer/src/elastic.rs` (`LEAVE_REASON_*`); a test pins both. +island's event tape (`island_leave`), and a failed send is recorded as +`elastic_leave_failed`. The codes are defined in `yeto/rl/pool_client.py` +(`LeaveReason`) and `syncer/src/protocol.rs` (`LEAVE_REASON_NAMES`); a test +pins both. | Code | Name | Sent when | |---|---|---| @@ -544,27 +602,19 @@ are defined in `yeto/rl/elastic_client.py` (`LeaveReason`) and | 3 | PREEMPTED | a droppable island received a cloud reclaim notice | | 4 | INFRA_FAILED | node or trainer lost, reconfiguration or storage failure, any unmapped error | | 5 | TRAINING_FAILED | a training check failed (publication, policy identity, optimizer step, non-finite values) | -| 6 | CONFIG_MISMATCH | incompatible configuration; a refused first JOIN is recorded on the tape only | +| 6 | CONFIG_MISMATCH | incompatible configuration | The exception table is `yeto/rl/engine/leave_reason.py`. An island sends at most one LEAVE: a reclaim listener's PREEMPTED is not overwritten by the -driver's close. A SIGTERM that kills the process (`yeto down`, platform stop) -sends nothing; the syncer removes the island after the lease. The syncer logs -an unknown code as `unknown()` and removes the island as usual. - -The lease defaults to 150 s (`syncer --island-lease-s`, the launcher, `yeto ---rl-island-lease-s` and the island client). The launcher passes the same -value to the syncer and to every elastic island (`--rl-island-lease-s`), and -islands heartbeat every lease/3 (50 s at the default) on their own thread, so -two lost heartbeats are tolerated. +driver's close. A SIGTERM that kills the process sends nothing; the syncer +records `disconnected` (or `lease_expired` after the lease). At the end of the run the launcher waives the missing `rl_learner_finalized` record only for islands whose last syncer `pool_leave` reason is `finished`, `stopped`, `preempted` or `lease_expired`; any other reason keeps the run failed. -To add an island to a running elastic fleet there is no CLI: start another -island process against the same syncer address with the same -`--rl-syncer-epoch`, the same HMAC key and a compatible backend identity and -compat group. It JOINs with `catch_up` and receives the current base; its -weight is zero in the round it joins. +To add an island to a running pool there is no CLI: start another island +process against the same syncer address with the same `--rl-syncer-epoch`, +the same HMAC key and a compatible backend identity and compat group (and an +id below `--max-learners`). It is caught up to the current base. diff --git a/docs/RL_ELASTIC_BENCHMARK.md b/docs/RL_ELASTIC_BENCHMARK.md index 559b272ff..8e35b733e 100644 --- a/docs/RL_ELASTIC_BENCHMARK.md +++ b/docs/RL_ELASTIC_BENCHMARK.md @@ -117,7 +117,7 @@ can get another value of a whitelisted parameter. This is for negative tests only. ``` -yeto launch ... --training-mode rl --rl-island-scheduling elastic \ +yeto launch ... --training-mode rl --rl-sync-quorum-theta 0.75 --rl-sync-max-base-lag 2 \ --gpu modal:1xh100,modal:1xh100 \ --rl-negative-test-run \ --rl-island-override 1:identity_test_salt=n12 @@ -129,7 +129,7 @@ yeto launch ... --training-mode rl --rl-island-scheduling elastic \ | `rl_max_policy_age` | another policy-age limit (落后上限); checked against the backend like the global flag | | `identity_test_salt` | TEST ONLY: mixed into `island_contract_sha256`, so the island declares a different identity | -`--rl-max-carry-lag` is a syncer parameter and cannot be changed per island +`--rl-sync-max-base-lag` is a syncer parameter and cannot be changed per island (10-09 ruling of the main agent for the user). Any other key, an island number out of range, or `--rl-island-override` without `--rl-negative-test-run` stops the launch before any cloud resource. diff --git a/docs/TESTING.md b/docs/TESTING.md index 5987e27a6..fde15d686 100644 --- a/docs/TESTING.md +++ b/docs/TESTING.md @@ -59,7 +59,7 @@ fix the side that is wrong. |---|---|---|---|---|---|---| | 2026-10-09 | s18-fkrt (base fc82f175) | 5055 | 45 | 23 | 0 | 413 s | -The 45 skips name the missing item: codex-cli 0.145.0, an elastic syncer -binary (`YETO_TEST_ELASTIC_SYNCER`), external adapters (yeto_miles_secrlenv, +The 45 skips name the missing item: codex-cli 0.145.0, a syncer +binary (`YETO_TEST_SYNCER`), external adapters (yeto_miles_secrlenv, SecRLEnv, CyberGym, swebench), upstream Miles/Megatron/mlx, CUDA, ray (the guard self-tests; they pass in a venv with ray), and a trimmed evidence log. diff --git a/openspec/changes/rl-code-quality-cleanup/specs/rl-codebase-consistency/spec.md b/openspec/changes/rl-code-quality-cleanup/specs/rl-codebase-consistency/spec.md index b02b6a945..32cefd869 100644 --- a/openspec/changes/rl-code-quality-cleanup/specs/rl-codebase-consistency/spec.md +++ b/openspec/changes/rl-code-quality-cleanup/specs/rl-codebase-consistency/spec.md @@ -55,7 +55,7 @@ The rule table of the three sync modes (identical line by line to the module doc |---|---|---|---| | strict (``StrictAvgSync``, ``DualStrictAvgSync`` for the critic) | syncer ``--learner-weight equal`` (launcher RL default); ``tokens2-over-steps`` = ``merge::learner_weight`` when chosen | ``--max-base-lag 0``: every island trains from the exact global version v and pushes against v; quorum = all islands | stop when the committed version reaches ``global_rounds``; ``finish`` checks the final policy equals the committed one | | decoupled (``DecoupledSync``) | same syncer flags as strict (``equal``, ``--max-base-lag 0``), per fragment | fragments rotate with ``--pipeline`` fragments in flight; each fragment is merged against its own fragment version | stop when the syncer is finalizing or a submission reaches ``total_fragment_steps``; ``learner_budget_steps`` consolidates into one final cut | -| elastic (``ElasticAvgSync``) | ``merge_weight`` = ``merge::learner_weight`` (tokens^2 / steps); a catch-up island has weight 0 on its join version; a late delta is carried with ``carry_gamma^lag`` | a delta with lag > ``DEFAULT_MAX_CARRY_LAG`` (2) is rejected (``carry_lag_exceeded``); samples at most ``MAX_SAMPLE_OUTER_LAG`` (2) outer steps stale; step at quorum ``theta`` (0.75) of capacity, at least ``q_min`` (1) | the syncer answers FINISHED after ``total_steps``, and keeps answering for ``final_grace_s`` (default: the soft deadline) | +| island pool (``PoolAvgSync``) | syncer ``--learner-weight`` (``tokens2-over-steps`` by default here); a late delta is carried with ``carry_gamma^lag`` | a delta whose base lags the round by more than ``--rl-sync-max-base-lag`` is rejected (``carry_lag_exceeded``); a round merges at quorum ``theta`` of capacity (at least ``q_min``) or at the soft deadline; islands hold leases and re-JOIN after an expiry | the syncer sends the final cut after ``total_steps`` and waits ``final_grace_s`` for the islands still connected | #### Scenario: No second implementation - **WHEN** the repo is searched for `StrictRlBridge`, `ElasticRlBridge`, `TwoRoleStrictAvg` and the PPO role-lane code of `local_learner` diff --git a/syncer/src/elastic.rs b/syncer/src/elastic.rs deleted file mode 100644 index 56f55d4b2..000000000 --- a/syncer/src/elastic.rs +++ /dev/null @@ -1,1879 +0,0 @@ -//! Inter-island scheduling, `elastic` mode only (change rl-inter-island-scheduling, -//! design D-S1..S4, D-S6 parts). The legacy (default) mode never touches this -//! module: no new message type is registered by the legacy server, and the -//! legacy merge/quorum/grace/strict path in `server.rs` is unchanged. -//! -//! Reference implementation mirrored here: Python -//! `yeto/rl/engine/island_ledger.CrossIslandLedger` (elastic branch). -//! -//! The coordinator runs behind the whole-vector service loop in -//! `elastic_server.rs` (frames, HMAC, fencing, leases, carry, checkpoint). -//! Its membership half is also ported to the fragment-round server as -//! `membership.rs`; this module is kept until the LoRA path runs there. - -use std::collections::BTreeMap; - -use anyhow::{bail, ensure, Context, Result}; -use sha2::{Digest, Sha256}; - -// --------------------------------------------------------------------------- -// Mode and parameters -// --------------------------------------------------------------------------- - -pub const MODE_LEGACY: &str = "legacy"; -pub const MODE_ELASTIC: &str = "elastic"; - -/// Defaults: design.md P4, marked [待真机校准] (not yet hardware-calibrated). -pub const DEFAULT_QUORUM_THETA: f64 = 0.75; -pub const DEFAULT_CARRY_GAMMA: f64 = 0.5; -pub const DEFAULT_Q_MIN: u32 = 1; -pub const DEFAULT_MAX_CARRY_LAG: u32 = 2; - -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct ElasticParams { - pub quorum_theta: f64, - pub carry_gamma: f64, - pub soft_deadline_s: u64, - pub q_min: u32, - pub max_carry_lag: u32, -} - -impl ElasticParams { - pub fn validate(&self) -> Result<()> { - ensure!( - self.quorum_theta > 0.0 && self.quorum_theta <= 1.0, - "quorum_theta must be in (0, 1], got {}", - self.quorum_theta - ); - ensure!( - (0.0..=1.0).contains(&self.carry_gamma), - "carry_gamma must be in [0, 1], got {}", - self.carry_gamma - ); - ensure!(self.q_min >= 1, "q_min must be >= 1"); - Ok(()) - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Default)] -pub enum IslandSchedulingMode { - /// Existing strict syncer behavior; contributes nothing to the contract - /// hash so every pre-existing semantic profile hash is unchanged. - #[default] - Legacy, - Elastic(ElasticParams), -} - -const ISLAND_SCHEDULING_DOMAIN: &[u8] = b"yeto-syncer-island-scheduling-v1\0"; - -impl IslandSchedulingMode { - pub fn parse( - mode: &str, - quorum_theta: f64, - carry_gamma: f64, - soft_deadline_s: u64, - q_min: u32, - max_carry_lag: u32, - ) -> Result { - match mode { - MODE_LEGACY => Ok(Self::Legacy), - MODE_ELASTIC => { - let p = ElasticParams { - quorum_theta, - carry_gamma, - soft_deadline_s, - q_min, - max_carry_lag, - }; - p.validate()?; - Ok(Self::Elastic(p)) - } - other => bail!("--island-scheduling-mode must be legacy|elastic, got {other:?}"), - } - } - - /// Bytes appended to the semantic profile encoding. Legacy appends - /// nothing (byte-identical legacy hash). Elastic appends the domain, the - /// field name `island_scheduling_mode`, the mode string and all P4 - /// parameters, so a legacy HELLO can never match an elastic server and - /// vice versa. Layout (little endian): - /// domain | u8 len | "island_scheduling_mode" | u8 len | "elastic" | - /// f64 quorum_theta | f64 carry_gamma | u64 soft_deadline_s | u32 q_min | - /// u32 max_carry_lag - pub fn encode_contract(&self, encoded: &mut Vec) { - let Self::Elastic(p) = self else { return }; - encoded.extend_from_slice(ISLAND_SCHEDULING_DOMAIN); - let field = b"island_scheduling_mode"; - encoded.push(field.len() as u8); - encoded.extend_from_slice(field); - encoded.push(MODE_ELASTIC.len() as u8); - encoded.extend_from_slice(MODE_ELASTIC.as_bytes()); - encoded.extend_from_slice(&p.quorum_theta.to_bits().to_le_bytes()); - encoded.extend_from_slice(&p.carry_gamma.to_bits().to_le_bytes()); - encoded.extend_from_slice(&p.soft_deadline_s.to_le_bytes()); - encoded.extend_from_slice(&p.q_min.to_le_bytes()); - encoded.extend_from_slice(&p.max_carry_lag.to_le_bytes()); - } -} - -// --------------------------------------------------------------------------- -// HMAC-SHA256 (RFC 2104), only for the elastic message types -// --------------------------------------------------------------------------- - -pub const HMAC_LEN: usize = 32; - -pub fn hmac_sha256(key: &[u8], msg: &[u8]) -> [u8; 32] { - const B: usize = 64; - let mut k = [0u8; B]; - if key.len() > B { - k[..32].copy_from_slice(&Sha256::digest(key)); - } else { - k[..key.len()].copy_from_slice(key); - } - let mut inner = Sha256::new(); - inner.update(k.map(|b| b ^ 0x36)); - inner.update(msg); - let ih = inner.finalize(); - let mut outer = Sha256::new(); - outer.update(k.map(|b| b ^ 0x5c)); - outer.update(ih); - outer.finalize().into() -} - -fn ct_eq(a: &[u8], b: &[u8]) -> bool { - a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0 -} - -/// Append the MAC over `msg_type || body`. -pub fn seal(key: &[u8], msg_type: u8, mut body: Vec) -> Vec { - let mut mac_input = Vec::with_capacity(body.len() + 1); - mac_input.push(msg_type); - mac_input.extend_from_slice(&body); - body.extend_from_slice(&hmac_sha256(key, &mac_input)); - body -} - -/// Verify and strip the trailing MAC; returns the authenticated body. -pub fn open<'a>(key: &[u8], msg_type: u8, payload: &'a [u8]) -> Result<&'a [u8]> { - ensure!(payload.len() >= HMAC_LEN, "elastic frame shorter than its HMAC"); - let (body, mac) = payload.split_at(payload.len() - HMAC_LEN); - let mut mac_input = Vec::with_capacity(body.len() + 1); - mac_input.push(msg_type); - mac_input.extend_from_slice(body); - ensure!( - ct_eq(&hmac_sha256(key, &mac_input), mac), - "elastic frame HMAC mismatch" - ); - Ok(body) -} - -// --------------------------------------------------------------------------- -// Frames (D-S1, D-S2, D-S4). Every frame starts with syncer_epoch (u64). -// --------------------------------------------------------------------------- - -#[derive(Debug, Clone, PartialEq)] -pub enum ElasticMsg { - /// Body: syncer_epoch u64 | island_id u32 | incarnation u64 | - /// capacity f64 | backend_identity [32] (yeto-framework-decoupling 6.2a: - /// sha256 of the island's BackendIdentity; all zero = not declared). - /// The 32-byte field was added in s17-elastic-identity; older JOIN frames - /// (28-byte body) are refused, so old and new islands/syncers cannot mix. - Join { - syncer_epoch: u64, - island_id: u32, - incarnation: u64, - capacity: f64, - backend_identity: [u8; 32], - /// yeto-framework-decoupling 7.7c: "-" (e.g. - /// "nvidia-h100"), u32 length + utf-8 after backend_identity; empty = - /// not declared. Added in s19-compat (version boundary: JOIN frames - /// without it are refused). - compat_group: String, - }, - JoinAck { - syncer_epoch: u64, - learner_slot: u32, - membership_epoch: u64, - base_version: u64, - policy_hash: [u8; 32], - catch_up: bool, - }, - Leave { - syncer_epoch: u64, - island_id: u32, - reason: u8, - }, - LeaseHeartbeat { - syncer_epoch: u64, - island_id: u32, - membership_epoch: u64, - inner_step: u64, - round_wall_s: f64, - }, - /// P3: the coordinator validates, deduplicates and judges the index - /// entry; samples stay in object storage and are never forwarded. - SampleIndex { - syncer_epoch: u64, - entry: SampleIndexEntry, - }, - /// Coordinator: the run reached total_steps; islands should exit 0. - Finished { - syncer_epoch: u64, - outer_version: u64, - policy_hash: [u8; 32], - }, - /// Coordinator reply to SAMPLE_INDEX. - SampleVerdict { - syncer_epoch: u64, - verdict: u8, - reason: String, - outer_lag: u64, - }, - DeltaReady { - syncer_epoch: u64, - island_id: u32, - base_version: u64, - c_tokens: u64, - c_steps: u64, - }, - ElasticInit { - syncer_epoch: u64, - island_id: u32, - params: Vec, - }, - /// `update` is θ − base as sent by the island (same sign convention as - /// PUSH_FRAGMENT); the server negates it into an outer gradient. - DeltaTensor { - syncer_epoch: u64, - island_id: u32, - base_version: u64, - c_tokens: u64, - c_steps: u64, - update: Vec, - }, - ElasticBase { - syncer_epoch: u64, - outer_version: u64, - params: Vec, - }, -} - -fn put_f32s(b: &mut Vec, v: &[f32]) { - b.extend_from_slice(&(v.len() as u64).to_le_bytes()); - for x in v { - b.extend_from_slice(&x.to_le_bytes()); - } -} - -fn get_f32s(r: &mut crate::protocol::Reader) -> Result> { - let n = usize::try_from(r.u64()?).context("tensor length")?; - let bytes = r.take(n.checked_mul(4).context("tensor length overflow")?)?; - let mut out = Vec::new(); - crate::protocol::decode_tensor(crate::protocol::DTYPE_F32, bytes, &mut out)?; - Ok(out) -} - -/// LEAVE reason codes (one wire byte). Must equal `LeaveReason` in -/// `yeto/rl/elastic_client.py`; a parity test pins both. Only the syncer -/// itself assigns `LEAVE_REASON_LEASE_EXPIRED`; islands send the others. -pub const LEAVE_REASON_FINISHED: u8 = 0; -pub const LEAVE_REASON_LEASE_EXPIRED: u8 = 1; -pub const LEAVE_REASON_STOPPED: u8 = 2; -pub const LEAVE_REASON_PREEMPTED: u8 = 3; -pub const LEAVE_REASON_INFRA_FAILED: u8 = 4; -pub const LEAVE_REASON_TRAINING_FAILED: u8 = 5; -pub const LEAVE_REASON_CONFIG_MISMATCH: u8 = 6; - -/// Log name of a LEAVE reason; an unknown code is logged, never refused. -pub fn leave_reason_name(reason: u8) -> String { - match reason { - LEAVE_REASON_FINISHED => "finished".into(), - LEAVE_REASON_LEASE_EXPIRED => "lease_expired".into(), - LEAVE_REASON_STOPPED => "stopped".into(), - LEAVE_REASON_PREEMPTED => "preempted".into(), - LEAVE_REASON_INFRA_FAILED => "infra_failed".into(), - LEAVE_REASON_TRAINING_FAILED => "training_failed".into(), - LEAVE_REASON_CONFIG_MISMATCH => "config_mismatch".into(), - other => format!("unknown({other})"), - } -} - -impl ElasticMsg { - pub fn msg_type(&self) -> u8 { - use crate::protocol::*; - match self { - Self::Join { .. } => MSG_JOIN, - Self::JoinAck { .. } => MSG_JOIN_ACK, - Self::Leave { .. } => MSG_LEAVE, - Self::LeaseHeartbeat { .. } => MSG_LEASE_HEARTBEAT, - Self::SampleIndex { .. } => MSG_SAMPLE_INDEX, - Self::SampleVerdict { .. } => MSG_SAMPLE_VERDICT, - Self::Finished { .. } => MSG_FINISHED, - Self::DeltaReady { .. } => MSG_DELTA_READY, - Self::ElasticInit { .. } => MSG_ELASTIC_INIT, - Self::DeltaTensor { .. } => MSG_DELTA_TENSOR, - Self::ElasticBase { .. } => MSG_ELASTIC_BASE, - } - } - - pub fn syncer_epoch(&self) -> u64 { - match self { - Self::Join { syncer_epoch, .. } - | Self::JoinAck { syncer_epoch, .. } - | Self::Leave { syncer_epoch, .. } - | Self::LeaseHeartbeat { syncer_epoch, .. } - | Self::SampleIndex { syncer_epoch, .. } - | Self::SampleVerdict { syncer_epoch, .. } - | Self::Finished { syncer_epoch, .. } - | Self::DeltaReady { syncer_epoch, .. } - | Self::ElasticInit { syncer_epoch, .. } - | Self::DeltaTensor { syncer_epoch, .. } - | Self::ElasticBase { syncer_epoch, .. } => *syncer_epoch, - } - } - - /// Encode body and append the HMAC. Returns (msg_type, payload). - pub fn encode(&self, key: &[u8]) -> (u8, Vec) { - let mut b = Vec::new(); - b.extend_from_slice(&self.syncer_epoch().to_le_bytes()); - match self { - Self::Join { island_id, incarnation, capacity, backend_identity, compat_group, .. } => { - b.extend_from_slice(&island_id.to_le_bytes()); - b.extend_from_slice(&incarnation.to_le_bytes()); - b.extend_from_slice(&capacity.to_bits().to_le_bytes()); - b.extend_from_slice(backend_identity); - put_str(&mut b, compat_group); - } - Self::JoinAck { - learner_slot, - membership_epoch, - base_version, - policy_hash, - catch_up, - .. - } => { - b.extend_from_slice(&learner_slot.to_le_bytes()); - b.extend_from_slice(&membership_epoch.to_le_bytes()); - b.extend_from_slice(&base_version.to_le_bytes()); - b.extend_from_slice(policy_hash); - b.push(u8::from(*catch_up)); - } - Self::Leave { island_id, reason, .. } => { - b.extend_from_slice(&island_id.to_le_bytes()); - b.push(*reason); - } - Self::LeaseHeartbeat { island_id, membership_epoch, inner_step, round_wall_s, .. } => { - b.extend_from_slice(&island_id.to_le_bytes()); - b.extend_from_slice(&membership_epoch.to_le_bytes()); - b.extend_from_slice(&inner_step.to_le_bytes()); - b.extend_from_slice(&round_wall_s.to_bits().to_le_bytes()); - } - Self::SampleIndex { entry, .. } => entry.encode(&mut b), - Self::Finished { outer_version, policy_hash, .. } => { - b.extend_from_slice(&outer_version.to_le_bytes()); - b.extend_from_slice(policy_hash); - } - Self::SampleVerdict { verdict, reason, outer_lag, .. } => { - b.push(*verdict); - put_str(&mut b, reason); - b.extend_from_slice(&outer_lag.to_le_bytes()); - } - Self::DeltaReady { island_id, base_version, c_tokens, c_steps, .. } => { - b.extend_from_slice(&island_id.to_le_bytes()); - b.extend_from_slice(&base_version.to_le_bytes()); - b.extend_from_slice(&c_tokens.to_le_bytes()); - b.extend_from_slice(&c_steps.to_le_bytes()); - } - Self::ElasticInit { island_id, params, .. } => { - b.extend_from_slice(&island_id.to_le_bytes()); - put_f32s(&mut b, params); - } - Self::DeltaTensor { island_id, base_version, c_tokens, c_steps, update, .. } => { - b.extend_from_slice(&island_id.to_le_bytes()); - b.extend_from_slice(&base_version.to_le_bytes()); - b.extend_from_slice(&c_tokens.to_le_bytes()); - b.extend_from_slice(&c_steps.to_le_bytes()); - put_f32s(&mut b, update); - } - Self::ElasticBase { outer_version, params, .. } => { - b.extend_from_slice(&outer_version.to_le_bytes()); - put_f32s(&mut b, params); - } - } - let t = self.msg_type(); - (t, seal(key, t, b)) - } - - /// Authenticate and decode. Unknown types are an error. - pub fn decode(key: &[u8], msg_type: u8, payload: &[u8]) -> Result { - use crate::protocol::*; - let body = open(key, msg_type, payload)?; - let mut r = Reader(body); - let syncer_epoch = r.u64()?; - let f64_ = |r: &mut Reader| -> Result { Ok(f64::from_bits(r.u64()?)) }; - let hash = |r: &mut Reader| -> Result<[u8; 32]> { - Ok(r.take(32)?.try_into().context("policy hash")?) - }; - let flag = |r: &mut Reader| -> Result { - match r.u8()? { - 0 => Ok(false), - 1 => Ok(true), - v => bail!("invalid bool byte {v}"), - } - }; - let msg = match msg_type { - MSG_JOIN => Self::Join { - syncer_epoch, - island_id: r.u32()?, - incarnation: r.u64()?, - capacity: f64_(&mut r)?, - backend_identity: r - .take(32) - .context("JOIN without backend identity (island older than the syncer?)")? - .try_into() - .context("backend identity")?, - compat_group: get_str(&mut r) - .context("JOIN without compat_group (island older than the syncer?)")?, - }, - MSG_JOIN_ACK => Self::JoinAck { - syncer_epoch, - learner_slot: r.u32()?, - membership_epoch: r.u64()?, - base_version: r.u64()?, - policy_hash: hash(&mut r)?, - catch_up: flag(&mut r)?, - }, - MSG_LEAVE => Self::Leave { syncer_epoch, island_id: r.u32()?, reason: r.u8()? }, - MSG_LEASE_HEARTBEAT => Self::LeaseHeartbeat { - syncer_epoch, - island_id: r.u32()?, - membership_epoch: r.u64()?, - inner_step: r.u64()?, - round_wall_s: f64_(&mut r)?, - }, - MSG_SAMPLE_INDEX => { - Self::SampleIndex { syncer_epoch, entry: SampleIndexEntry::decode(&mut r)? } - } - MSG_FINISHED => Self::Finished { - syncer_epoch, - outer_version: r.u64()?, - policy_hash: hash(&mut r)?, - }, - MSG_SAMPLE_VERDICT => Self::SampleVerdict { - syncer_epoch, - verdict: r.u8()?, - reason: get_str(&mut r)?, - outer_lag: r.u64()?, - }, - MSG_DELTA_READY => Self::DeltaReady { - syncer_epoch, - island_id: r.u32()?, - base_version: r.u64()?, - c_tokens: r.u64()?, - c_steps: r.u64()?, - }, - MSG_ELASTIC_INIT => { - Self::ElasticInit { syncer_epoch, island_id: r.u32()?, params: get_f32s(&mut r)? } - } - MSG_DELTA_TENSOR => Self::DeltaTensor { - syncer_epoch, - island_id: r.u32()?, - base_version: r.u64()?, - c_tokens: r.u64()?, - c_steps: r.u64()?, - update: get_f32s(&mut r)?, - }, - MSG_ELASTIC_BASE => Self::ElasticBase { - syncer_epoch, - outer_version: r.u64()?, - params: get_f32s(&mut r)?, - }, - other => bail!("message type {other} is not an elastic message"), - }; - ensure!(r.remaining() == 0, "trailing bytes in elastic frame type {msg_type}"); - if let Self::Join { capacity, .. } = &msg { - ensure!(capacity.is_finite() && *capacity > 0.0, "JOIN capacity must be > 0"); - } - Ok(msg) - } -} - -// --------------------------------------------------------------------------- -// Coordinator state machine (mirrors CrossIslandLedger, elastic branch) -// --------------------------------------------------------------------------- - -/// Syncer merge weight: `merge::learner_weight` (tokens^2 / steps), with the -/// elastic zero-tokens / zero-steps branch kept explicit. -pub fn merge_weight(c_tokens: u64, c_steps: u64) -> f64 { - if c_tokens == 0 || c_steps == 0 { - return 0.0; - } - crate::merge::learner_weight(c_tokens, c_steps) -} - -#[derive(Debug, Clone, PartialEq)] -struct Member { - joined_at: u64, - catch_up: bool, - last_heartbeat: f64, - capacity: f64, - /// D-S5: EMA of the island-reported round wall time (LEASE_HEARTBEAT). - round_wall_ema_s: Option, - /// D-S5: arrived-on-time flags of the most recent steps (newest last). - arrivals: std::collections::VecDeque, -} - -impl Member { - fn new(joined_at: u64, catch_up: bool, now: f64, capacity: f64) -> Self { - Self { - joined_at, - catch_up, - last_heartbeat: now, - capacity, - round_wall_ema_s: None, - arrivals: Default::default(), - } - } -} - -/// D-S5 EMA smoothing factor and arrival history length. -pub const ROUND_WALL_EMA_ALPHA: f64 = 0.2; -pub const ARRIVAL_HISTORY: usize = 8; - -/// Sample staleness bound (design: user ruling 2026-10-07, at most 2 outer -/// steps; same as Python StalenessPolicy.max_outer_lag). Not a CLI flag yet. -pub const MAX_SAMPLE_OUTER_LAG: u64 = 2; - -pub const SAMPLE_INDEX_SCHEMA: &str = "yeto.rl.sample-index/v1"; -pub const SAMPLE_URI_SCHEMES: [&str; 3] = ["s3://", "modal-volume://", "nebius-os://"]; -pub const VERDICT_ACCEPT: u8 = 0; -pub const VERDICT_ACCEPT_IS: u8 = 1; -pub const VERDICT_REJECT: u8 = 2; - -pub fn verdict_name(v: u8) -> &'static str { - match v { - VERDICT_ACCEPT => "ACCEPT", - VERDICT_ACCEPT_IS => "ACCEPT_IS", - _ => "REJECT", - } -} - -fn put_str(b: &mut Vec, v: &str) { - b.extend_from_slice(&(v.len() as u32).to_le_bytes()); - b.extend_from_slice(v.as_bytes()); -} - -fn get_str(r: &mut crate::protocol::Reader) -> Result { - let n = r.u32()? as usize; - String::from_utf8(r.take(n)?.to_vec()).context("string is not utf-8") -} - -/// Mirrors Python `yeto/rl/engine/sample_pool.SampleIndexEntry` -/// (schema yeto.rl.sample-index/v1). Wire order (little endian; str = u32 -/// length + utf-8): island_id u32 | outer_version u64 | inner_step u64 | -/// policy_hash [32] (raw sha256; Python hex string decoded) | group_id str | -/// prompt_id str | n u64 | uri str | size_bytes u64 | sha256 str (64 hex) | -/// has_behavior_logprob u8 | advantage_included u8 | created_at f64 | schema str -#[derive(Debug, Clone, PartialEq)] -pub struct SampleIndexEntry { - pub island_id: u32, - pub outer_version: u64, - pub inner_step: u64, - pub policy_hash: [u8; 32], - pub group_id: String, - pub prompt_id: String, - pub n: u64, - pub uri: String, - pub size_bytes: u64, - pub sha256: String, - pub has_behavior_logprob: bool, - pub advantage_included: bool, - pub created_at: f64, - pub schema: String, -} - -fn hex(b: &[u8]) -> String { - b.iter().map(|x| format!("{x:02x}")).collect() -} - -fn json_str(v: &str) -> String { - let mut o = String::from("\""); - for c in v.chars() { - match c { - '"' => o.push_str("\\\""), - '\\' => o.push_str("\\\\"), - c if (c as u32) < 0x20 => o.push_str(&format!("\\u{:04x}", c as u32)), - c => o.push(c), - } - } - o.push('"'); - o -} - -impl SampleIndexEntry { - pub fn encode(&self, b: &mut Vec) { - b.extend_from_slice(&self.island_id.to_le_bytes()); - b.extend_from_slice(&self.outer_version.to_le_bytes()); - b.extend_from_slice(&self.inner_step.to_le_bytes()); - b.extend_from_slice(&self.policy_hash); - put_str(b, &self.group_id); - put_str(b, &self.prompt_id); - b.extend_from_slice(&self.n.to_le_bytes()); - put_str(b, &self.uri); - b.extend_from_slice(&self.size_bytes.to_le_bytes()); - put_str(b, &self.sha256); - b.push(u8::from(self.has_behavior_logprob)); - b.push(u8::from(self.advantage_included)); - b.extend_from_slice(&self.created_at.to_bits().to_le_bytes()); - put_str(b, &self.schema); - } - - pub fn decode(r: &mut crate::protocol::Reader) -> Result { - let flag = |r: &mut crate::protocol::Reader| -> Result { - match r.u8()? { - 0 => Ok(false), - 1 => Ok(true), - v => bail!("invalid bool byte {v}"), - } - }; - Ok(Self { - island_id: r.u32()?, - outer_version: r.u64()?, - inner_step: r.u64()?, - policy_hash: r.take(32)?.try_into()?, - group_id: get_str(r)?, - prompt_id: get_str(r)?, - n: r.u64()?, - uri: get_str(r)?, - size_bytes: r.u64()?, - sha256: get_str(r)?, - has_behavior_logprob: flag(r)?, - advantage_included: flag(r)?, - created_at: f64::from_bits(r.u64()?), - schema: get_str(r)?, - }) - } - - /// Same checks as the Python `__post_init__` (sha256 additionally must - /// be lowercase hex). - pub fn validate(&self) -> Result<()> { - ensure!(self.schema == SAMPLE_INDEX_SCHEMA, "schema {:?} != {SAMPLE_INDEX_SCHEMA}", self.schema); - ensure!( - SAMPLE_URI_SCHEMES.iter().any(|p| self.uri.starts_with(p)), - "uri {:?} is not one of {SAMPLE_URI_SCHEMES:?}", - self.uri - ); - ensure!(self.n >= 1, "n >= 1 required"); - ensure!( - self.sha256.len() == 64 && self.sha256.bytes().all(|c| c.is_ascii_digit() || (b'a'..=b'f').contains(&c)), - "sha256 must be 64 hex chars" - ); - ensure!(self.advantage_included, "group rule: the producer must include the group advantage"); - ensure!(!self.group_id.is_empty(), "group_id must not be empty"); - Ok(()) - } - - /// JSON with Python's field names (island_id as the decimal string). - pub fn to_json(&self) -> String { - format!( - "{{\"advantage_included\":{},\"created_at\":{},\"group_id\":{},\"has_behavior_logprob\":{},\"inner_step\":{},\"island_id\":{},\"n\":{},\"outer_version\":{},\"policy_hash\":\"{}\",\"prompt_id\":{},\"schema\":{},\"sha256\":{},\"size_bytes\":{},\"uri\":{}}}", - self.advantage_included, - f64_json(self.created_at), - json_str(&self.group_id), - self.has_behavior_logprob, - self.inner_step, - json_str(&self.island_id.to_string()), - self.n, - self.outer_version, - hex(&self.policy_hash), - json_str(&self.prompt_id), - json_str(&self.schema), - json_str(&self.sha256), - self.size_bytes, - json_str(&self.uri), - ) - } -} - -/// One accepted-into-index record (verdict ACCEPT or ACCEPT_IS). -#[derive(Debug, Clone, PartialEq)] -pub struct IndexedSample { - pub entry: SampleIndexEntry, - pub verdict: u8, - pub reason: &'static str, - pub outer_lag: u64, -} - -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct Delta { - pub island_id: u32, - pub base_version: u64, - pub c_tokens: u64, - pub c_steps: u64, -} - -/// Tape records (JSONL kinds match the Python ledger's `kind` strings). -#[derive(Debug, Clone, PartialEq)] -pub enum TapeEvent { - PoolJoin { island_id: u32, catch_up: bool, base_version: u64, membership_epoch: u64 }, - PoolLeave { - island_id: u32, - reason: u8, - dropped_uncommitted: Option, - membership_epoch: u64, - }, - DeltaAccepted { island_id: u32 }, - DeltaCarriedOver { island_id: u32, base: u64, lag: u64, discount: f64 }, - DeltaRejected { island_id: u32, reason: &'static str }, - RoundIdle { arrived: usize, cap_arrived: f64, cap_total: f64 }, - SampleIndexed { island_id: u32, group_id: String, verdict: u8, reason: &'static str, outer_lag: u64, duplicate: bool }, - OuterStep { - base_version: u64, - arrived: usize, - cap_arrived: f64, - cap_total: f64, - timed_out: bool, - /// (island, base_version, weight before normalization); carried - /// entries have base_version < this step's base_version. - raw_weights: Vec<(u32, u64, f64)>, - /// Same entries normalized to sum 1 (what merge_avg applies). - weights: Vec<(u32, u64, f64)>, - absent: Vec, - /// Deltas rejected since the previous outer step: (island, reason). - rejected: Vec<(u32, &'static str)>, - }, -} - -fn delta_json(d: &Delta) -> String { - format!( - "{{\"island_id\":{},\"outer_version\":{},\"c_tokens\":{},\"c_steps\":{}}}", - d.island_id, d.base_version, d.c_tokens, d.c_steps - ) -} - -fn f64_json(v: f64) -> String { - if v.is_finite() { format!("{v:?}") } else { "null".into() } -} - -impl TapeEvent { - /// One JSONL record; `kind` strings match the Python ledger. - pub fn to_json(&self, syncer_epoch: u64) -> String { - let body = match self { - Self::PoolJoin { island_id, catch_up, base_version, membership_epoch } => format!( - "\"island_id\":{island_id},\"catch_up\":{catch_up},\"base_version\":{base_version},\"membership_epoch\":{membership_epoch}" - ), - Self::PoolLeave { island_id, reason, dropped_uncommitted, membership_epoch } => format!( - "\"island_id\":{island_id},\"reason\":\"{}\",\"dropped_uncommitted\":{},\"membership_epoch\":{membership_epoch}", - leave_reason_name(*reason), - dropped_uncommitted.as_ref().map(delta_json).unwrap_or_else(|| "null".into()) - ), - Self::DeltaAccepted { island_id } => format!("\"island_id\":{island_id}"), - Self::SampleIndexed { island_id, group_id, verdict, reason, outer_lag, duplicate } => format!( - "\"island_id\":{island_id},\"group_id\":{},\"verdict\":\"{}\",\"reason\":\"{reason}\",\"outer_lag\":{outer_lag},\"duplicate\":{duplicate}", - json_str(group_id), - verdict_name(*verdict) - ), - Self::DeltaCarriedOver { island_id, base, lag, discount } => format!( - "\"island_id\":{island_id},\"base\":{base},\"lag\":{lag},\"discount\":{}", - f64_json(*discount) - ), - Self::DeltaRejected { island_id, reason } => { - format!("\"island_id\":{island_id},\"reason\":\"{reason}\"") - } - Self::RoundIdle { arrived, cap_arrived, cap_total } => format!( - "\"arrived\":{arrived},\"cap_arrived\":{},\"cap_total\":{}", - f64_json(*cap_arrived), - f64_json(*cap_total) - ), - Self::OuterStep { base_version, arrived, cap_arrived, cap_total, timed_out, raw_weights, weights, absent, rejected } => { - let fmt = |v: &Vec<(u32, u64, f64)>| -> String { - v.iter().map(|(i, b, w)| format!("[{i},{b},{}]", f64_json(*w))).collect::>().join(",") - }; - let w = fmt(raw_weights); - let nw = fmt(weights); - let a: Vec = absent.iter().map(u32::to_string).collect(); - let rj: Vec = rejected.iter().map(|(i, r)| format!("[{i},\"{r}\"]")).collect(); - format!( - "\"base_version\":{base_version},\"arrived\":{arrived},\"cap_arrived\":{},\"cap_total\":{},\"timed_out\":{timed_out},\"raw_weights\":[{}],\"weights\":[{}],\"absent\":[{}],\"rejected\":[{}]", - f64_json(*cap_arrived), - f64_json(*cap_total), - w, - nw, - a.join(","), - rj.join(",") - ) - } - }; - format!("{{\"kind\":\"{}\",\"syncer_epoch\":{syncer_epoch},{body}}}", self.kind()) - } - - pub fn kind(&self) -> &'static str { - match self { - Self::PoolJoin { .. } => "pool_join", - Self::PoolLeave { .. } => "pool_leave", - Self::DeltaAccepted { .. } => "delta_accepted", - Self::DeltaCarriedOver { .. } => "delta_carried_over", - Self::DeltaRejected { .. } => "delta_rejected", - Self::RoundIdle { .. } => "round_idle", - Self::SampleIndexed { .. } => "sample_indexed", - Self::OuterStep { .. } => "outer_step", - } - } -} - -/// Wall-clock fields added to each tape line when it is written (L5.7). -/// -/// The coordinator stays clock-free (its tape is compared with the Python -/// ledger); the server stamps each event as it flushes it, under the same -/// lock and right after the event was created. The round starts at the -/// previous `outer_step` line (or when the clock was made). -#[derive(Debug, Clone, PartialEq)] -pub struct TapeClock { - round_start_unix: f64, - last_arrival: Option, -} - -impl TapeClock { - pub fn new(now_unix: f64) -> Self { - Self { round_start_unix: now_unix, last_arrival: None } - } - - /// Extra JSON fields for `ev` written at `now_unix` (leading comma - /// included): `time_unix` on every event; `arrival_offset_s` on - /// accepted and carried deltas; `step_seconds` and - /// `last_arrival_island` on `outer_step`, which also starts a new round. - pub fn stamp(&mut self, ev: &TapeEvent, now_unix: f64) -> String { - let mut out = format!(",\"time_unix\":{}", f64_json(now_unix)); - let offset = now_unix - self.round_start_unix; - match ev { - TapeEvent::DeltaAccepted { island_id } | TapeEvent::DeltaCarriedOver { island_id, .. } => { - self.last_arrival = Some(*island_id); - out += &format!(",\"arrival_offset_s\":{}", f64_json(offset)); - } - TapeEvent::OuterStep { .. } => { - let last = self.last_arrival.take().map(|i| i.to_string()).unwrap_or_else(|| "null".into()); - out += &format!(",\"step_seconds\":{},\"last_arrival_island\":{last}", f64_json(offset)); - self.round_start_unix = now_unix; - } - _ => {} - } - out - } -} - -impl TapeEvent { - /// `to_json` plus the clock fields from `TapeClock::stamp`. - pub fn to_json_stamped(&self, syncer_epoch: u64, clock: &mut TapeClock, now_unix: f64) -> String { - let mut line = self.to_json(syncer_epoch); - line.pop(); // closing brace - line + &clock.stamp(self, now_unix) + "}" - } -} - -#[derive(Debug, Clone, PartialEq)] -pub struct StepPlan { - pub base_version: u64, - /// Normalized merge weights over current + carried deltas. - pub weights: Vec<(u32, u64, f64)>, -} - -pub struct ElasticCoordinator { - pub params: ElasticParams, - pub lease_s: f64, - pub syncer_epoch: u64, - pub membership_epoch: u64, - pub outer_version: u64, - members: BTreeMap, - pending: BTreeMap, - carried: BTreeMap, - /// Rejections since the last outer step (L5.5 merge metrics; not - /// checkpointed, a restart starts the count again). - rejected: Vec<(u32, &'static str)>, - pub tape: Vec, - /// Ledger of uncommitted deltas dropped by LEAVE / lease expiry - /// (persisted by the D-S6 checkpoint). - pub dropped: Vec, - /// Published policy hash per outer version (sha256 of the f32 base); - /// set by the server after each step. Empty in control-plane-only runs. - pub published: BTreeMap, - /// Sample index keyed by (island, group_id); never forwarded. - pub sample_index: BTreeMap<(u32, String), IndexedSample>, -} - -const COORD_STATE_MAGIC: &[u8; 8] = b"YELCST1\0"; - -fn put_delta(b: &mut Vec, d: &Delta) { - b.extend_from_slice(&d.island_id.to_le_bytes()); - b.extend_from_slice(&d.base_version.to_le_bytes()); - b.extend_from_slice(&d.c_tokens.to_le_bytes()); - b.extend_from_slice(&d.c_steps.to_le_bytes()); -} - -fn get_delta(r: &mut crate::protocol::Reader) -> Result { - Ok(Delta { island_id: r.u32()?, base_version: r.u64()?, c_tokens: r.u64()?, c_steps: r.u64()? }) -} - -impl ElasticCoordinator { - fn observe_round_wall(&mut self, island_id: u32, round_wall_s: f64) { - if let Some(m) = self.members.get_mut(&island_id) { - if round_wall_s.is_finite() && round_wall_s > 0.0 { - m.round_wall_ema_s = Some(match m.round_wall_ema_s { - None => round_wall_s, - Some(e) => e + ROUND_WALL_EMA_ALPHA * (round_wall_s - e), - }); - } - } - } - - /// P3 index judgement relative to the current outer version (the - /// consumer is unknown at index time). Mirrors the ledger's `judge` - /// minus consumer-specific rules. Returns (verdict, reason, outer_lag); - /// errors only for malformed entries. Accepted entries are indexed; - /// a byte-identical resubmission is idempotent. - pub fn index_sample(&mut self, e: &SampleIndexEntry) -> Result<(u8, &'static str, u64)> { - e.validate()?; - let key = (e.island_id, e.group_id.clone()); - if let Some(prev) = self.sample_index.get(&key) { - ensure!(prev.entry == *e, "group {} of island {} already indexed with different content", e.group_id, e.island_id); - let (v, r, l) = (prev.verdict, prev.reason, prev.outer_lag); - self.tape.push(TapeEvent::SampleIndexed { island_id: e.island_id, group_id: e.group_id.clone(), verdict: v, reason: r, outer_lag: l, duplicate: true }); - return Ok((v, r, l)); - } - let lag = self.outer_version.checked_sub(e.outer_version); - let (verdict, reason) = if !self.is_member(e.island_id) { - (VERDICT_REJECT, "producer_not_member") - } else if lag.is_none() || !self.published.contains_key(&e.outer_version) { - (VERDICT_REJECT, "unknown_version") - } else if self.published[&e.outer_version] != e.policy_hash { - (VERDICT_REJECT, "policy_hash_mismatch") - } else if lag == Some(0) { - (VERDICT_ACCEPT, "same_version") - } else if lag.unwrap() > MAX_SAMPLE_OUTER_LAG { - (VERDICT_REJECT, "outer_lag_exceeded") - } else if !e.has_behavior_logprob { - (VERDICT_REJECT, "missing_behavior_logprob") - } else { - (VERDICT_ACCEPT_IS, "stale_within_bound") - }; - let outer_lag = lag.unwrap_or(0); - if verdict != VERDICT_REJECT { - self.sample_index.insert(key, IndexedSample { entry: e.clone(), verdict, reason, outer_lag }); - } - self.tape.push(TapeEvent::SampleIndexed { island_id: e.island_id, group_id: e.group_id.clone(), verdict, reason, outer_lag, duplicate: false }); - Ok((verdict, reason, outer_lag)) - } - - /// D-S5 status snapshot (JSON, schema yeto.syncer.elastic-status/v1). - pub fn status_json(&self, now: f64, soft_deadline_remaining_s: f64, state: &str) -> String { - let (cap_arrived, cap_total) = self.capacity_arrived(); - let islands: Vec = self - .members - .iter() - .map(|(id, m)| { - let arrivals: Vec<&str> = m.arrivals.iter().map(|a| if *a { "true" } else { "false" }).collect(); - format!( - "{}:{{\"capacity\":{},\"joined_at\":{},\"catch_up\":{},\"round_wall_ema_s\":{},\"lease_remaining_s\":{},\"arrival_history\":[{}],\"pending\":{},\"carried_over_lag\":{}}}", - json_str(&id.to_string()), - f64_json(m.capacity), - m.joined_at, - m.catch_up && m.joined_at == self.outer_version, - m.round_wall_ema_s.map(f64_json).unwrap_or_else(|| "null".into()), - f64_json((self.lease_s - (now - m.last_heartbeat)).max(0.0)), - arrivals.join(","), - self.pending.contains_key(id), - self.carried.get(id).map(|(_, l)| l.to_string()).unwrap_or_else(|| "null".into()), - ) - }) - .collect(); - format!( - "{{\"schema\":\"yeto.syncer.elastic-status/v1\",\"state\":{},\"syncer_epoch\":{},\"membership_epoch\":{},\"outer_version\":{},\"policy_hash\":{},\"cap_arrived\":{},\"cap_total\":{},\"soft_deadline_remaining_s\":{},\"pending_count\":{},\"carried_over_count\":{},\"dropped_uncommitted_count\":{},\"sample_index_count\":{},\"islands\":{{{}}}}}", - json_str(state), - self.syncer_epoch, - self.membership_epoch, - self.outer_version, - self.published.get(&self.outer_version).map(|h| format!("\"{}\"", hex(h))).unwrap_or_else(|| "null".into()), - f64_json(cap_arrived), - f64_json(cap_total), - f64_json(soft_deadline_remaining_s.max(0.0)), - self.pending.len(), - self.carried.len(), - self.dropped.len(), - self.sample_index.len(), - islands.join(","), - ) - } - - /// D-S6: serialize membership, epochs, carried-over and dropped ledger. - /// Pending current-round deltas are not persisted (islands resend); - /// the elastic contract encoding is embedded and checked on restore. - pub fn encode_state(&self) -> Vec { - let mut b = COORD_STATE_MAGIC.to_vec(); - let mut contract = Vec::new(); - IslandSchedulingMode::Elastic(self.params).encode_contract(&mut contract); - b.extend_from_slice(&(contract.len() as u32).to_le_bytes()); - b.extend_from_slice(&contract); - b.extend_from_slice(&self.syncer_epoch.to_le_bytes()); - b.extend_from_slice(&self.membership_epoch.to_le_bytes()); - b.extend_from_slice(&self.outer_version.to_le_bytes()); - b.extend_from_slice(&(self.members.len() as u32).to_le_bytes()); - for (id, m) in &self.members { - b.extend_from_slice(&id.to_le_bytes()); - b.extend_from_slice(&m.joined_at.to_le_bytes()); - b.push(u8::from(m.catch_up)); - b.extend_from_slice(&m.capacity.to_bits().to_le_bytes()); - } - b.extend_from_slice(&(self.carried.len() as u32).to_le_bytes()); - for (d, lag) in self.carried.values() { - put_delta(&mut b, d); - b.extend_from_slice(&lag.to_le_bytes()); - } - b.extend_from_slice(&(self.dropped.len() as u32).to_le_bytes()); - for d in &self.dropped { - put_delta(&mut b, d); - } - b - } - - /// Restore from `encode_state`. Members get a fresh lease starting at - /// `now`; the caller decides the new syncer_epoch (saved + 1). - pub fn decode_state(params: ElasticParams, lease_s: f64, bytes: &[u8], now: f64) -> Result { - let mut r = crate::protocol::Reader(bytes); - ensure!(r.take(8)? == COORD_STATE_MAGIC, "not an elastic coordinator state"); - let n = r.u32()? as usize; - let mut expected = Vec::new(); - IslandSchedulingMode::Elastic(params).encode_contract(&mut expected); - ensure!(r.take(n)? == expected.as_slice(), "checkpoint elastic parameters differ from this launch"); - let mut c = Self::new(params, lease_s, r.u64()?)?; - c.membership_epoch = r.u64()?; - c.outer_version = r.u64()?; - for _ in 0..r.u32()? { - let id = r.u32()?; - let joined_at = r.u64()?; - let catch_up = r.u8()? != 0; - let capacity = f64::from_bits(r.u64()?); - c.members.insert(id, Member::new(joined_at, catch_up, now, capacity)); - } - for _ in 0..r.u32()? { - let d = get_delta(&mut r)?; - let lag = r.u64()?; - c.carried.insert(d.island_id, (d, lag)); - } - for _ in 0..r.u32()? { - c.dropped.push(get_delta(&mut r)?); - } - ensure!(r.remaining() == 0, "trailing bytes in elastic coordinator state"); - Ok(c) - } - - pub fn new(params: ElasticParams, lease_s: f64, syncer_epoch: u64) -> Result { - params.validate()?; - Ok(Self { - params, - lease_s, - syncer_epoch, - membership_epoch: 0, - outer_version: 0, - members: BTreeMap::new(), - pending: BTreeMap::new(), - carried: BTreeMap::new(), - rejected: Vec::new(), - tape: Vec::new(), - dropped: Vec::new(), - published: BTreeMap::new(), - sample_index: BTreeMap::new(), - }) - } - - /// P9 fencing: refuse anything stamped by an older coordinator incarnation. - pub fn check_fence(&self, msg_epoch: u64) -> Result<()> { - ensure!( - msg_epoch >= self.syncer_epoch, - "fenced: syncer_epoch {msg_epoch} < {}", - self.syncer_epoch - ); - Ok(()) - } - - pub fn member_count(&self) -> usize { - self.members.len() - } - - pub fn is_member(&self, island_id: u32) -> bool { - self.members.contains_key(&island_id) - } - - pub fn join(&mut self, island_id: u32, capacity: f64, now: f64) -> Result { - ensure!(!self.members.contains_key(&island_id), "{island_id} is already a member"); - ensure!(capacity.is_finite() && capacity > 0.0, "capacity must be > 0"); - // Founding round (nothing published yet) is not catch-up. - let catch_up = self.outer_version > 0; - self.members.insert( - island_id, - Member::new(self.outer_version, catch_up, now, capacity), - ); - self.membership_epoch += 1; - self.tape.push(TapeEvent::PoolJoin { - island_id, - catch_up, - base_version: self.outer_version, - membership_epoch: self.membership_epoch, - }); - Ok(catch_up) - } - - pub fn leave(&mut self, island_id: u32, reason: u8) -> Result<()> { - ensure!(self.members.remove(&island_id).is_some(), "{island_id} is not a member"); - let mut dropped = self.pending.remove(&island_id); - let carried = self.carried.remove(&island_id); - if dropped.is_none() { - dropped = carried.map(|(d, _)| d); - } - self.dropped.extend(dropped); - self.membership_epoch += 1; - self.tape.push(TapeEvent::PoolLeave { - island_id, - reason, - dropped_uncommitted: dropped, - membership_epoch: self.membership_epoch, - }); - Ok(()) - } - - pub fn heartbeat(&mut self, island_id: u32, now: f64) -> Result<()> { - let m = self - .members - .get_mut(&island_id) - .with_context(|| format!("{island_id} is not a member (rejoin required)"))?; - m.last_heartbeat = now; - Ok(()) - } - - pub fn expire_leases(&mut self, now: f64) -> Vec { - let gone: Vec = self - .members - .iter() - .filter(|(_, m)| now - m.last_heartbeat > self.lease_s) - .map(|(&i, _)| i) - .collect(); - for &i in &gone { - self.leave(i, LEAVE_REASON_LEASE_EXPIRED).expect("member present"); - } - gone - } - - /// Apply one authenticated elastic message (fenced first). - pub fn apply(&mut self, msg: &ElasticMsg, now: f64) -> Result> { - self.check_fence(msg.syncer_epoch())?; - match msg { - ElasticMsg::Join { island_id, capacity, .. } => { - let catch_up = self.join(*island_id, *capacity, now)?; - Ok(Some(ElasticMsg::JoinAck { - syncer_epoch: self.syncer_epoch, - learner_slot: *island_id, - membership_epoch: self.membership_epoch, - base_version: self.outer_version, - policy_hash: self.published.get(&self.outer_version).copied().unwrap_or([0; 32]), - catch_up, - })) - } - ElasticMsg::Leave { island_id, reason, .. } => { - self.leave(*island_id, *reason)?; - Ok(None) - } - ElasticMsg::LeaseHeartbeat { island_id, round_wall_s, .. } => { - self.heartbeat(*island_id, now)?; - self.observe_round_wall(*island_id, *round_wall_s); - Ok(None) - } - ElasticMsg::SampleIndex { entry, .. } => { - let (verdict, reason, outer_lag) = self.index_sample(entry)?; - Ok(Some(ElasticMsg::SampleVerdict { - syncer_epoch: self.syncer_epoch, - verdict, - reason: reason.to_string(), - outer_lag, - })) - } - ElasticMsg::DeltaReady { island_id, base_version, c_tokens, c_steps, .. } - | ElasticMsg::DeltaTensor { island_id, base_version, c_tokens, c_steps, .. } => { - // A delivered delta also proves liveness (renews the lease). - if self.is_member(*island_id) { - self.heartbeat(*island_id, now)?; - } - let ev = self.submit(Delta { - island_id: *island_id, - base_version: *base_version, - c_tokens: *c_tokens, - c_steps: *c_steps, - }); - if let TapeEvent::DeltaRejected { reason, .. } = ev { - bail!("delta rejected: {reason}"); - } - Ok(None) - } - ElasticMsg::JoinAck { .. } - | ElasticMsg::ElasticBase { .. } - | ElasticMsg::SampleVerdict { .. } - | ElasticMsg::Finished { .. } => { - bail!("coordinator-to-island message received by the coordinator") - } - ElasticMsg::ElasticInit { .. } => bail!("ELASTIC_INIT is handled by the server"), - } - } - - pub fn capacity_arrived(&self) -> (f64, f64) { - let total = self.members.values().map(|m| m.capacity).sum(); - let arrived = self.pending.keys().map(|i| self.members[i].capacity).sum(); - (arrived, total) - } - - pub fn submit(&mut self, d: Delta) -> &TapeEvent { - let ev = if !self.members.contains_key(&d.island_id) { - TapeEvent::DeltaRejected { island_id: d.island_id, reason: "not_member" } - } else if d.base_version > self.outer_version { - TapeEvent::DeltaRejected { island_id: d.island_id, reason: "unknown_version" } - } else { - let lag = self.outer_version - d.base_version; - if lag == 0 { - self.pending.insert(d.island_id, d); - TapeEvent::DeltaAccepted { island_id: d.island_id } - } else if lag > u64::from(self.params.max_carry_lag) { - TapeEvent::DeltaRejected { island_id: d.island_id, reason: "carry_lag_exceeded" } - } else { - self.carried.insert(d.island_id, (d, lag)); - TapeEvent::DeltaCarriedOver { - island_id: d.island_id, - base: d.base_version, - lag, - discount: self.params.carry_gamma.powi(lag as i32), - } - } - }; - if let TapeEvent::DeltaRejected { island_id, reason } = ev { - self.rejected.push((island_id, reason)); - } - self.tape.push(ev); - self.tape.last().unwrap() - } - - fn weight_of(&self, d: &Delta) -> f64 { - let m = &self.members[&d.island_id]; - if m.catch_up && m.joined_at == self.outer_version { - return 0.0; - } - merge_weight(d.c_tokens, d.c_steps) - } - - /// P4 step. `timed_out` means the soft deadline (soft_deadline_s) passed. - pub fn try_advance(&mut self, timed_out: bool) -> Option { - let (cap_arrived, cap_total) = self.capacity_arrived(); - let arrived = self.pending.len(); - let reached = cap_total > 0.0 && cap_arrived >= self.params.quorum_theta * cap_total; - if !reached && !timed_out { - return None; - } - if arrived < self.params.q_min as usize { - if timed_out { - self.tape.push(TapeEvent::RoundIdle { arrived, cap_arrived, cap_total }); - } - return None; - } - let base_version = self.outer_version; - let mut raw: Vec<(u32, u64, f64)> = - self.pending.values().map(|d| (d.island_id, d.base_version, self.weight_of(d))).collect(); - for (i, (d, lag)) in &self.carried { - if self.members.contains_key(i) { - let w = merge_weight(d.c_tokens, d.c_steps) - * self.params.carry_gamma.powi(*lag as i32); - raw.push((*i, d.base_version, w)); - } - } - self.carried.clear(); - let total: f64 = raw.iter().map(|x| x.2).sum(); - let weights: Vec<(u32, u64, f64)> = raw - .iter() - .map(|&(i, b, w)| (i, b, if total > 0.0 { w / total } else { 0.0 })) - .collect(); - let absent = self.members.keys().filter(|i| !self.pending.contains_key(i)).copied().collect(); - let arrived_ids: Vec = self.pending.keys().copied().collect(); - for (id, m) in self.members.iter_mut() { - m.arrivals.push_back(arrived_ids.contains(id)); - while m.arrivals.len() > ARRIVAL_HISTORY { - m.arrivals.pop_front(); - } - } - self.outer_version += 1; - self.pending.clear(); - self.tape.push(TapeEvent::OuterStep { - base_version, - arrived, - cap_arrived, - cap_total, - timed_out, - raw_weights: raw, - weights: weights.clone(), - absent, - rejected: std::mem::take(&mut self.rejected), - }); - Some(StepPlan { base_version, weights }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::protocol::*; - - /// rl-island-graceful-leave D7 Q5: every code has a log name; an unknown - /// code is logged as unknown() and still removes the island. - #[test] - fn leave_reason_names_and_unknown_code() { - let names: Vec = (0..=6).map(leave_reason_name).collect(); - assert_eq!(names, ["finished", "lease_expired", "stopped", "preempted", "infra_failed", - "training_failed", "config_mismatch"]); - assert_eq!(leave_reason_name(42), "unknown(42)"); - let mut c = ElasticCoordinator::new(params(), 30.0, 0).unwrap(); - c.join(1, 1.0, 0.0).unwrap(); - c.leave(1, 42).unwrap(); - assert!(!c.is_member(1)); - } - - /// rl-code-quality-cleanup 2.4: the elastic weight is merge::learner_weight, - /// bit for bit, and keeps its zero branch; same values as the old - /// tokens^2 / steps formula. - #[test] - fn merge_weight_is_merge_learner_weight() { - for &(t, s) in &[(1u64, 1u64), (37, 4), (91, 4), (4096, 3), (1 << 20, 7), (u32::MAX as u64, 1), - (5, u32::MAX as u64 + 9)] { - assert_eq!(merge_weight(t, s).to_bits(), crate::merge::learner_weight(t, s).to_bits()); - assert_eq!(merge_weight(t, s).to_bits(), ((t as f64).powi(2) / s as f64).to_bits()); - } - assert_eq!(merge_weight(0, 4), 0.0); - assert_eq!(merge_weight(7, 0), 0.0); - } - - /// Golden replay: tests/test_rl_inter_island_rust_golden.py drives the - /// Python CrossIslandLedger (elastic) through a fixed scenario and writes - /// the expected step/weight/carried/rejected/dropped sequence; the Rust - /// coordinator must reproduce it op by op. - #[test] - fn python_ledger_golden_tape_matches() { - let text = include_str!("../tests/fixtures/elastic_ledger_golden.txt"); - let close = |a: f64, b: f64| (a - b).abs() <= 1e-12; - let mut c: Option = None; - let mut checked = 0; - for (n, line) in text.lines().enumerate() { - if line.starts_with('#') || line.trim().is_empty() { - continue; - } - let t: Vec<&str> = line.split_whitespace().collect(); - if t[0] == "params" { - let p = ElasticParams { - quorum_theta: t[1].parse().unwrap(), - carry_gamma: t[2].parse().unwrap(), - soft_deadline_s: 900, - q_min: t[3].parse().unwrap(), - max_carry_lag: t[4].parse().unwrap(), - }; - c = Some(ElasticCoordinator::new(p, t[5].parse().unwrap(), 0).unwrap()); - continue; - } - let c = c.as_mut().expect("params line first"); - let (lhs, rhs) = line.split_once(" => ").unwrap(); - let a: Vec<&str> = lhs.split_whitespace().collect(); - let e: Vec<&str> = rhs.split_whitespace().collect(); - let id = |s: &str| s.parse::().unwrap(); - let u = |s: &str| s.parse::().unwrap(); - let ctx = format!("fixture line {}: {line}", n + 1); - match a[0] { - "join" => { - let catch_up = c.join(id(a[1]), a[2].parse().unwrap(), a[3].parse().unwrap()).unwrap(); - assert_eq!(catch_up.to_string(), e[1], "{ctx}"); - assert_eq!(c.outer_version, u(e[2]), "{ctx}"); - } - "leave" => { - c.leave(id(a[1]), LEAVE_REASON_FINISHED).unwrap(); - match c.tape.last().unwrap() { - TapeEvent::PoolLeave { dropped_uncommitted, .. } => { - let got = dropped_uncommitted - .map(|d| format!("{} {} {}", d.base_version, d.c_tokens, d.c_steps)) - .unwrap_or_else(|| "none".into()); - assert_eq!(got, e[1..].join(" "), "{ctx}"); - } - ev => panic!("{ctx}: {ev:?}"), - } - } - "submit" => { - let ev = c.submit(Delta { - island_id: id(a[1]), - base_version: u(a[2]), - c_tokens: u(a[3]), - c_steps: u(a[4]), - }); - let got = match ev { - TapeEvent::DeltaAccepted { .. } => "delta_accepted".to_string(), - TapeEvent::DeltaCarriedOver { lag, .. } => format!("delta_carried_over {lag}"), - TapeEvent::DeltaRejected { reason, .. } => format!("delta_rejected {reason}"), - ev => panic!("{ctx}: {ev:?}"), - }; - assert_eq!(got, rhs, "{ctx}"); - } - "advance" => match (c.try_advance(a[1] == "1"), e[0]) { - (None, "no_step") => assert_eq!(c.outer_version, u(e[1]), "{ctx}"), - (Some(plan), "outer_step") => { - assert_eq!(plan.base_version, u(e[1]), "{ctx}"); - let TapeEvent::OuterStep { arrived, cap_arrived, cap_total, absent, .. } = - c.tape.last().unwrap().clone() - else { - panic!("{ctx}") - }; - assert_eq!(arrived, e[2].parse::().unwrap(), "{ctx}"); - assert!(close(cap_arrived, e[3].parse().unwrap()), "{ctx}"); - assert!(close(cap_total, e[4].parse().unwrap()), "{ctx}"); - let want_absent: Vec = - if e[5] == "-" { vec![] } else { e[5].split(',').map(id).collect() }; - assert_eq!(absent, want_absent, "{ctx}"); - let mut got = plan.weights.clone(); - got.sort_by(|x, y| (x.0, x.1).cmp(&(y.0, y.1))); - let want: Vec<(u32, u64, f64)> = e[6] - .split(',') - .map(|w| { - let f: Vec<&str> = w.split(':').collect(); - (id(f[0]), u(f[1]), f[2].parse().unwrap()) - }) - .collect(); - assert_eq!(got.len(), want.len(), "{ctx}"); - for (g, w) in got.iter().zip(&want) { - assert!(g.0 == w.0 && g.1 == w.1 && close(g.2, w.2), "{ctx}: {g:?} vs {w:?}"); - } - } - (got, _) => panic!("{ctx}: got {got:?}"), - }, - op => panic!("{ctx}: unknown op {op}"), - } - checked += 1; - } - assert_eq!(checked, 23, "fixture op count"); - } - - fn params() -> ElasticParams { - ElasticParams { - quorum_theta: DEFAULT_QUORUM_THETA, - carry_gamma: DEFAULT_CARRY_GAMMA, - soft_deadline_s: 900, - q_min: DEFAULT_Q_MIN, - max_carry_lag: DEFAULT_MAX_CARRY_LAG, - } - } - - fn delta(island_id: u32, base_version: u64) -> Delta { - Delta { island_id, base_version, c_tokens: 10, c_steps: 1 } - } - - fn entry(island_id: u32, outer_version: u64, policy_hash: [u8; 32], group: &str) -> SampleIndexEntry { - SampleIndexEntry { - island_id, - outer_version, - inner_step: 3, - policy_hash, - group_id: group.into(), - prompt_id: "p".into(), - n: 8, - uri: "s3://bucket/g.parquet".into(), - size_bytes: 1024, - sha256: "a".repeat(64), - has_behavior_logprob: true, - advantage_included: true, - created_at: 1.5, - schema: SAMPLE_INDEX_SCHEMA.into(), - } - } - - #[test] - fn sample_index_validates_dedups_and_judges_three_ways() { - let mut c = ElasticCoordinator::new(params(), 30.0, 0).unwrap(); - c.join(1, 1.0, 0.0).unwrap(); - c.published.insert(0, [0; 32]); - for v in 1..=3u64 { - c.submit(delta(1, v - 1)); - c.try_advance(false).unwrap(); - c.published.insert(v, [v as u8; 32]); - } - // now outer_version 3 - assert_eq!(c.index_sample(&entry(1, 3, [3; 32], "a")).unwrap(), (VERDICT_ACCEPT, "same_version", 0)); - assert_eq!(c.index_sample(&entry(1, 2, [2; 32], "b")).unwrap(), (VERDICT_ACCEPT_IS, "stale_within_bound", 1)); - assert_eq!(c.index_sample(&entry(1, 0, [0; 32], "c")).unwrap().1, "outer_lag_exceeded"); - assert_eq!(c.index_sample(&entry(1, 2, [9; 32], "d")).unwrap().1, "policy_hash_mismatch"); - assert_eq!(c.index_sample(&entry(1, 7, [3; 32], "e")).unwrap().1, "unknown_version"); - assert_eq!(c.index_sample(&entry(2, 3, [3; 32], "f")).unwrap().1, "producer_not_member"); - let mut no_lp = entry(1, 2, [2; 32], "g"); - no_lp.has_behavior_logprob = false; - assert_eq!(c.index_sample(&no_lp).unwrap().1, "missing_behavior_logprob"); - // Duplicate is idempotent; conflicting content is an error. - assert_eq!(c.index_sample(&entry(1, 2, [2; 32], "b")).unwrap().0, VERDICT_ACCEPT_IS); - let mut conflict = entry(1, 2, [2; 32], "b"); - conflict.n = 9; - assert!(c.index_sample(&conflict).is_err()); - assert_eq!(c.sample_index.len(), 2); - // Validation mirrors Python __post_init__. - for bad in [ - SampleIndexEntry { schema: "x".into(), ..entry(1, 3, [3; 32], "h") }, - SampleIndexEntry { uri: "http://x".into(), ..entry(1, 3, [3; 32], "h") }, - SampleIndexEntry { n: 0, ..entry(1, 3, [3; 32], "h") }, - SampleIndexEntry { sha256: "z".repeat(64), ..entry(1, 3, [3; 32], "h") }, - SampleIndexEntry { advantage_included: false, ..entry(1, 3, [3; 32], "h") }, - ] { - assert!(c.index_sample(&bad).is_err(), "{bad:?}"); - } - assert!(c.tape.iter().any(|e| e.kind() == "sample_indexed")); - assert!(entry(1, 3, [3; 32], "a").to_json().starts_with("{\"advantage_included\":true,")); - } - - #[test] - fn status_snapshot_exports_scheduling_fields() { - let mut c = ElasticCoordinator::new(params(), 30.0, 2).unwrap(); - c.join(1, 3.0, 0.0).unwrap(); - c.join(2, 1.0, 0.0).unwrap(); - c.observe_round_wall(1, 10.0); - c.observe_round_wall(1, 20.0); - c.submit(delta(1, 0)); - c.try_advance(false).unwrap(); - c.submit(delta(2, 0)); // carried, lag 1 - let s = c.status_json(5.0, 12.0, "running"); - assert!(s.contains("\"state\":\"running\"")); - for needle in [ - "\"syncer_epoch\":2", - "\"outer_version\":1", - "\"carried_over_count\":1", - "\"1\":{\"capacity\":3.0", - "\"round_wall_ema_s\":12.0", - "\"lease_remaining_s\":25.0", - "\"arrival_history\":[true]", - "\"arrival_history\":[false],\"pending\":false,\"carried_over_lag\":1", - "\"soft_deadline_remaining_s\":12.0", - ] { - assert!(s.contains(needle), "{needle} not in {s}"); - } - } - - #[test] - fn hmac_matches_rfc4231_case_2() { - let mac = hmac_sha256(b"Jefe", b"what do ya want for nothing?"); - let hex: String = mac.iter().map(|b| format!("{b:02x}")).collect(); - assert_eq!(hex, "5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843"); - } - - #[test] - fn mode_parsing_and_validation() { - assert_eq!(IslandSchedulingMode::parse("legacy", 9.0, 9.0, 0, 0, 0).unwrap(), IslandSchedulingMode::Legacy); - assert!(IslandSchedulingMode::parse("elastic", 0.75, 0.5, 900, 1, 2).is_ok()); - assert!(IslandSchedulingMode::parse("elastic", 0.0, 0.5, 900, 1, 2).is_err()); - assert!(IslandSchedulingMode::parse("elastic", 0.75, 1.5, 900, 1, 2).is_err()); - assert!(IslandSchedulingMode::parse("elastic", 0.75, 0.5, 900, 0, 2).is_err()); - assert!(IslandSchedulingMode::parse("Elastic", 0.75, 0.5, 900, 1, 2).is_err()); - let mut v = Vec::new(); - IslandSchedulingMode::Legacy.encode_contract(&mut v); - assert!(v.is_empty(), "legacy must not change the contract encoding"); - } - - #[test] - fn all_frames_roundtrip_and_reject_tampering_and_wrong_key() { - let msgs = vec![ - ElasticMsg::Join { syncer_epoch: 3, island_id: 1, incarnation: 7, capacity: 2.5, backend_identity: [0xab; 32], compat_group: String::new() }, - ElasticMsg::JoinAck { - syncer_epoch: 3, - learner_slot: 1, - membership_epoch: 4, - base_version: 9, - policy_hash: [7; 32], - catch_up: true, - }, - ElasticMsg::Leave { syncer_epoch: 3, island_id: 1, reason: LEAVE_REASON_FINISHED }, - ElasticMsg::LeaseHeartbeat { - syncer_epoch: 3, - island_id: 1, - membership_epoch: 4, - inner_step: 11, - round_wall_s: 1.25, - }, - ElasticMsg::SampleIndex { syncer_epoch: 3, entry: entry(1, 9, [1; 32], "g1") }, - ElasticMsg::SampleVerdict { - syncer_epoch: 3, - verdict: VERDICT_ACCEPT_IS, - reason: "stale_within_bound".into(), - outer_lag: 1, - }, - ElasticMsg::Finished { syncer_epoch: 3, outer_version: 8, policy_hash: [5; 32] }, - ElasticMsg::DeltaReady { - syncer_epoch: 3, - island_id: 1, - base_version: 9, - c_tokens: 100, - c_steps: 4, - }, - ElasticMsg::ElasticInit { syncer_epoch: 3, island_id: 1, params: vec![1.0, -2.0] }, - ElasticMsg::DeltaTensor { - syncer_epoch: 3, - island_id: 1, - base_version: 2, - c_tokens: 5, - c_steps: 1, - update: vec![0.5, 0.25, -1.0], - }, - ElasticMsg::ElasticBase { syncer_epoch: 3, outer_version: 4, params: vec![3.0] }, - ]; - for m in msgs { - let (t, p) = m.encode(b"k1"); - assert_eq!(ElasticMsg::decode(b"k1", t, &p).unwrap(), m); - assert!(ElasticMsg::decode(b"k2", t, &p).is_err(), "wrong key accepted"); - let mut bad = p.clone(); - bad[0] ^= 1; - assert!(ElasticMsg::decode(b"k1", t, &bad).is_err(), "tampered body accepted"); - // MAC binds the type byte too. - let other = if t == MSG_LEAVE { MSG_JOIN } else { MSG_LEAVE }; - assert!(ElasticMsg::decode(b"k1", other, &p).is_err()); - } - assert!(ElasticMsg::decode(b"k1", MSG_HEARTBEAT, &seal(b"k1", MSG_HEARTBEAT, vec![0; 8])).is_err()); - } - - /// Golden JOIN frame (key b"k1"), the bytes the Python client - /// (tests/test_rl_inter_island_elastic_client.py RUST_GOLDEN[15]) must - /// produce; and the pre-identity 28-byte JOIN body is refused. - #[test] - fn join_frame_with_backend_identity_golden_and_old_format_refused() { - let m = ElasticMsg::Join { - syncer_epoch: 3, - island_id: 1, - incarnation: 7, - capacity: 2.5, - backend_identity: [0xab; 32], - compat_group: "nvidia-h100".to_string(), - }; - let (t, p) = m.encode(b"k1"); - assert_eq!(t, MSG_JOIN); - println!("JOIN_GOLDEN {}", hex(&p)); - assert_eq!(hex(&p), JOIN_GOLDEN_HEX); - let mut old = Vec::new(); - old.extend_from_slice(&3u64.to_le_bytes()); - old.extend_from_slice(&1u32.to_le_bytes()); - old.extend_from_slice(&7u64.to_le_bytes()); - old.extend_from_slice(&2.5f64.to_bits().to_le_bytes()); - let err = ElasticMsg::decode(b"k1", MSG_JOIN, &seal(b"k1", MSG_JOIN, old.clone())).unwrap_err(); - assert!(format!("{err:#}").contains("without backend identity"), "{err:#}"); - // s19-compat: a JOIN with the identity but without compat_group is refused. - old.extend_from_slice(&[0xab; 32]); - let err = ElasticMsg::decode(b"k1", MSG_JOIN, &seal(b"k1", MSG_JOIN, old)).unwrap_err(); - assert!(format!("{err:#}").contains("without compat_group"), "{err:#}"); - } - - const JOIN_GOLDEN_HEX: &str = "03000000000000000100000007000000000000000000000000000440abababababababababababababababababababababababababababababababab0b0000006e76696469612d68313030d76ccd494a9ee9bb15d43cdcafb727ab8d0c651b9b81c6481e80ff4084314b4c"; - - fn hex(b: &[u8]) -> String { - b.iter().map(|x| format!("{x:02x}")).collect() - } - - #[test] - fn stale_syncer_epoch_is_fenced() { - let mut c = ElasticCoordinator::new(params(), 30.0, 5).unwrap(); - let old = ElasticMsg::Join { syncer_epoch: 4, island_id: 1, incarnation: 0, capacity: 1.0, backend_identity: [0; 32], compat_group: String::new() }; - assert!(format!("{:#}", c.apply(&old, 0.0).unwrap_err()).contains("fenced")); - assert!(!c.is_member(1)); - let cur = ElasticMsg::Join { syncer_epoch: 5, island_id: 1, incarnation: 0, capacity: 1.0, backend_identity: [0; 32], compat_group: String::new() }; - assert!(matches!(c.apply(&cur, 0.0).unwrap(), Some(ElasticMsg::JoinAck { syncer_epoch: 5, .. }))); - let hb = ElasticMsg::LeaseHeartbeat { syncer_epoch: 4, island_id: 1, membership_epoch: 1, inner_step: 0, round_wall_s: 0.0 }; - assert!(c.apply(&hb, 1.0).is_err()); - } - - #[test] - fn capacity_weighted_quorum_steps_without_waiting_for_small_island() { - let mut c = ElasticCoordinator::new(params(), 30.0, 0).unwrap(); - c.join(1, 3.0, 0.0).unwrap(); - c.join(2, 1.0, 0.0).unwrap(); - // Small island alone: 1/4 < 0.75 -> wait. - c.submit(delta(2, 0)); - assert!(c.try_advance(false).is_none()); - c.pending.clear(); - // Big island alone: 3/4 >= 0.75 -> step, small island absent. - c.submit(delta(1, 0)); - let plan = c.try_advance(false).unwrap(); - assert_eq!(plan.base_version, 0); - assert_eq!(plan.weights, vec![(1, 0, 1.0)]); - assert_eq!(c.outer_version, 1); - match c.tape.last().unwrap() { - TapeEvent::OuterStep { absent, .. } => assert_eq!(absent, &vec![2]), - e => panic!("{e:?}"), - } - } - - #[test] - fn soft_deadline_steps_or_idles_by_q_min() { - let mut p = params(); - p.q_min = 2; - let mut c = ElasticCoordinator::new(p, 30.0, 0).unwrap(); - for i in 1..=3 { - c.join(i, 1.0, 0.0).unwrap(); - } - c.submit(delta(1, 0)); - assert!(c.try_advance(true).is_none()); - assert_eq!(c.tape.last().unwrap().kind(), "round_idle"); - c.submit(delta(2, 0)); - assert!(c.try_advance(false).is_none(), "2/3 < 0.75 before deadline"); - assert!(c.try_advance(true).is_some(), "deadline with arrived >= q_min"); - } - - #[test] - fn late_delta_is_discounted_by_gamma_pow_lag_and_rejected_past_max_lag() { - let mut c = ElasticCoordinator::new(params(), 30.0, 0).unwrap(); - c.join(1, 1.0, 0.0).unwrap(); - c.join(2, 1.0, 0.0).unwrap(); - c.submit(delta(1, 0)); - c.try_advance(true).unwrap(); // v1 - c.submit(delta(1, 1)); - c.try_advance(true).unwrap(); // v2 - // Island 2 finally delivers on base 0: lag 2 -> carried with 0.25. - match c.submit(delta(2, 0)) { - TapeEvent::DeltaCarriedOver { lag, discount, .. } => { - assert_eq!((*lag, *discount), (2, 0.25)) - } - e => panic!("{e:?}"), - } - c.submit(delta(1, 2)); - let plan = c.try_advance(true).unwrap(); - // raw: island1 100, island2 100*0.25=25 -> 0.8 / 0.2 - assert_eq!(plan.weights, vec![(1, 2, 0.8), (2, 0, 0.2)]); - // lag 3 > max_carry_lag 2 -> rejected. - match c.submit(delta(2, 0)) { - TapeEvent::DeltaRejected { reason, .. } => assert_eq!(*reason, "carry_lag_exceeded"), - e => panic!("{e:?}"), - } - } - - #[test] - fn leave_and_lease_expiry_drop_uncommitted_delta() { - let mut c = ElasticCoordinator::new(params(), 10.0, 0).unwrap(); - c.join(1, 1.0, 0.0).unwrap(); - c.join(2, 1.0, 0.0).unwrap(); - c.join(3, 1.0, 0.0).unwrap(); - c.submit(delta(2, 0)); - c.leave(2, LEAVE_REASON_FINISHED).unwrap(); - match c.tape.last().unwrap() { - TapeEvent::PoolLeave { dropped_uncommitted: Some(d), .. } => assert_eq!(d.island_id, 2), - e => panic!("{e:?}"), - } - c.submit(delta(3, 0)); - c.heartbeat(1, 15.0).unwrap(); - assert_eq!(c.expire_leases(15.0), vec![3]); - match c.tape.last().unwrap() { - TapeEvent::PoolLeave { reason, dropped_uncommitted: Some(d), .. } => { - assert_eq!((*reason, d.island_id), (LEAVE_REASON_LEASE_EXPIRED, 3)) - } - e => panic!("{e:?}"), - } - assert_eq!(c.capacity_arrived(), (0.0, 1.0)); - assert_eq!(c.membership_epoch, 5); - } - - #[test] - fn coordinator_state_roundtrips_membership_carry_and_dropped_ledger() { - let mut c = ElasticCoordinator::new(params(), 30.0, 4).unwrap(); - c.join(1, 2.0, 0.0).unwrap(); - c.join(2, 1.0, 0.0).unwrap(); - c.join(3, 1.0, 0.0).unwrap(); - c.submit(delta(1, 0)); - c.try_advance(true).unwrap(); - c.submit(delta(2, 0)); // carried, lag 1 - c.submit(delta(3, 1)); - c.leave(3, LEAVE_REASON_FINISHED).unwrap(); // dropped - let bytes = c.encode_state(); - let r = ElasticCoordinator::decode_state(params(), 30.0, &bytes, 100.0).unwrap(); - assert_eq!((r.syncer_epoch, r.membership_epoch, r.outer_version), (4, 4, 1)); - assert_eq!(r.members.keys().copied().collect::>(), vec![1, 2]); - assert_eq!(r.members[&1].capacity, 2.0); - assert_eq!(r.carried[&2], (delta(2, 0), 1)); - assert_eq!(r.dropped, vec![delta(3, 1)]); - assert_eq!(r.members[&2].last_heartbeat, 100.0); - let mut other = params(); - other.carry_gamma = 0.25; - assert!(ElasticCoordinator::decode_state(other, 30.0, &bytes, 0.0).is_err()); - } - - #[test] - fn joiner_has_zero_weight_in_its_first_round() { - let mut c = ElasticCoordinator::new(params(), 30.0, 0).unwrap(); - assert!(!c.join(1, 1.0, 0.0).unwrap(), "founding member is not catch-up"); - c.submit(delta(1, 0)); - c.try_advance(false).unwrap(); - assert!(c.join(2, 1.0, 0.0).unwrap()); - c.submit(delta(1, 1)); - c.submit(delta(2, 1)); - let plan = c.try_advance(false).unwrap(); - assert_eq!(plan.weights, vec![(1, 1, 1.0), (2, 1, 0.0)]); - c.submit(delta(1, 2)); - c.submit(delta(2, 2)); - let plan = c.try_advance(false).unwrap(); - assert_eq!(plan.weights, vec![(1, 2, 0.5), (2, 2, 0.5)]); - } - - /// L5.5: weights, carried discount and rejected reasons per outer step, - /// against hand-computed values (w = tokens^2 / steps). - #[test] - fn outer_step_records_merge_shares_carry_and_rejections() { - let mut c = ElasticCoordinator::new(params(), 30.0, 0).unwrap(); - for i in 1..=3 { - c.join(i, 1.0, 0.0).unwrap(); - } - c.submit(Delta { island_id: 1, base_version: 0, c_tokens: 10, c_steps: 1 }); - c.submit(Delta { island_id: 2, base_version: 0, c_tokens: 20, c_steps: 2 }); - let plan = c.try_advance(true).unwrap(); // v0 -> v1, island 3 absent - assert_eq!(plan.weights, vec![(1, 0, 100.0 / 300.0), (2, 0, 200.0 / 300.0)]); - match c.tape.last().unwrap() { - TapeEvent::OuterStep { rejected, .. } => assert!(rejected.is_empty()), - e => panic!("{e:?}"), - } - - c.submit(Delta { island_id: 3, base_version: 0, c_tokens: 30, c_steps: 1 }); // lag 1 - c.submit(Delta { island_id: 1, base_version: 1, c_tokens: 10, c_steps: 1 }); - c.submit(Delta { island_id: 9, base_version: 1, c_tokens: 10, c_steps: 1 }); - c.submit(Delta { island_id: 2, base_version: 5, c_tokens: 10, c_steps: 1 }); - let plan = c.try_advance(true).unwrap(); - // raw: island 1 = 100; island 3 carried = 900 * 0.5 = 450. - let mut got = plan.weights.clone(); - got.sort_by_key(|x| x.0); - assert_eq!(got, vec![(1, 1, 100.0 / 550.0), (3, 0, 450.0 / 550.0)]); - let ev = c.tape.last().unwrap(); - let TapeEvent::OuterStep { raw_weights, rejected, .. } = ev else { panic!("{ev:?}") }; - assert!(raw_weights.contains(&(3, 0, 450.0))); - assert_eq!(rejected, &vec![(9, "not_member"), (2, "unknown_version")]); - assert!(ev.to_json(0).ends_with(",\"rejected\":[[9,\"not_member\"],[2,\"unknown_version\"]]}")); - } - - /// L5.7: time_unix on every line; arrival offset on deltas; step - /// duration and last arrival on outer_step, which restarts the round. - #[test] - fn tape_clock_stamps_arrival_offsets_and_step_duration() { - let mut clock = TapeClock::new(100.0); - let join = TapeEvent::PoolJoin { island_id: 1, catch_up: false, base_version: 0, membership_epoch: 1 }; - assert_eq!(clock.stamp(&join, 100.5), ",\"time_unix\":100.5"); - let acc = TapeEvent::DeltaAccepted { island_id: 1 }; - assert_eq!(clock.stamp(&acc, 101.5), ",\"time_unix\":101.5,\"arrival_offset_s\":1.5"); - let carry = TapeEvent::DeltaCarriedOver { island_id: 3, base: 0, lag: 1, discount: 0.5 }; - assert_eq!(clock.stamp(&carry, 103.0), ",\"time_unix\":103.0,\"arrival_offset_s\":3.0"); - let step = TapeEvent::OuterStep { - base_version: 1, - arrived: 1, - cap_arrived: 1.0, - cap_total: 2.0, - timed_out: false, - raw_weights: vec![], - weights: vec![], - absent: vec![], - rejected: vec![], - }; - let line = step.to_json_stamped(7, &mut clock, 104.0); - assert!(line.starts_with("{\"kind\":\"outer_step\",\"syncer_epoch\":7,")); - assert!(line.ends_with(",\"time_unix\":104.0,\"step_seconds\":4.0,\"last_arrival_island\":3}")); - assert_eq!(clock.stamp(&acc, 104.25), ",\"time_unix\":104.25,\"arrival_offset_s\":0.25"); - clock.stamp(&step, 105.0); - assert!(clock.stamp(&step, 106.0).ends_with("\"step_seconds\":1.0,\"last_arrival_island\":null")); - } -} diff --git a/syncer/src/elastic_server.rs b/syncer/src/elastic_server.rs deleted file mode 100644 index 9a9f9966c..000000000 --- a/syncer/src/elastic_server.rs +++ /dev/null @@ -1,1646 +0,0 @@ -//! Elastic-mode server: runs `ElasticCoordinator` behind a TCP listener and -//! applies its `StepPlan` to a flat f32 parameter vector. Used only when -//! `--island-scheduling-mode elastic`; the legacy server path in -//! `server.rs` never reaches this file. -//! -//! Messages: JOIN/JOIN_ACK, LEAVE, LEASE_HEARTBEAT, ELASTIC_INIT (seed the -//! base), DELTA_TENSOR (update + merge-weight inputs), ELASTIC_BASE -//! (broadcast after each step), and the metadata-only DELTA_READY (control -//! plane only; refused for merging once a base exists). -//! Merge: plan weights (γ^lag for carried deltas, 0 for a joiner's first -//! round, departed islands' deltas never stored into the plan) feed -//! `merge::merge_avg`, then `merge::nesterov_step` with --outer-lr / -//! --outer-momentum. The legacy per-fragment RDA/ISO/HeLoCo pipeline and -//! layout/fragment framing are not used here. -//! Not yet: SAMPLE_INDEX handling (refused), D-S5 status export, D-S6 -//! checkpoint fields. - -use std::collections::HashMap; -use std::io::Write; -use std::path::PathBuf; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{Arc, Mutex}; -use std::time::{Duration, Instant}; - -use anyhow::{bail, ensure, Context, Result}; -use tokio::io::AsyncWriteExt; -use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::mpsc; -use tracing::{info, warn}; - -use crate::elastic::{ElasticCoordinator, ElasticMsg, ElasticParams}; -use crate::protocol::*; - -pub struct ElasticServerConfig { - pub params: ElasticParams, - pub key: Vec, - pub syncer_epoch: u64, - pub lease_s: f64, - pub total_steps: u64, - pub event_tape: Option, - pub tick: Duration, - pub outer_lr: f32, - pub outer_momentum: f32, - /// Delayed Nesterov period N >= 2; None = plain Nesterov. - pub outer_delayed_nesterov_period: Option, - /// After total_steps: keep serving FINISHED for this long (or until - /// every member has left) so late islands exit cleanly. - pub final_grace: Duration, - /// D-S6 checkpoint (elastic-only format, tmp + fsync + rename). - pub checkpoint_path: Option, - pub checkpoint_every: u64, - pub resume: bool, - /// secret-handling-hardening: backend identity pinned by the head's - /// config. None keeps "first accepted JOIN pins it". - pub expected_backend_identity: Option<[u8; 32]>, -} - -/// v2 (s17-elastic-identity) adds the pinned backend identity; v1 -/// checkpoints are refused (version boundary, see hash-migration.md). -const SERVER_STATE_MAGIC: &[u8; 8] = b"YELSRV2\0"; -const SERVER_STATE_MAGIC_V1: &[u8; 8] = b"YELSRV1\0"; - -fn group_label(g: &str) -> &str { - if g.is_empty() { "未声明" } else { g } -} - -/// Seconds since the Unix epoch (tape `time_unix`). -fn unix_now() -> f64 { - std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_or(f64::NAN, |d| d.as_secs_f64()) -} - -fn hex32(h: &[u8; 32]) -> String { - if h.iter().all(|b| *b == 0) { - return "".into(); - } - h.iter().map(|b| format!("{b:02x}")).collect() -} - -fn put_f32s(b: &mut Vec, v: &[f32]) { - b.extend_from_slice(&(v.len() as u64).to_le_bytes()); - v.iter().for_each(|x| b.extend_from_slice(&x.to_le_bytes())); -} - -fn get_f32s(r: &mut Reader) -> Result> { - let n = usize::try_from(r.u64()?)?; - let bytes = r.take(n.checked_mul(4).context("tensor length overflow")?)?; - let mut out = Vec::new(); - decode_tensor(DTYPE_F32, bytes, &mut out)?; - Ok(out) -} - -fn write_atomic(path: &std::path::Path, bytes: &[u8]) -> Result<()> { - let tmp = path.with_extension("tmp"); - { - let mut f = std::fs::File::create(&tmp)?; - f.write_all(bytes)?; - f.sync_all()?; - } - std::fs::rename(&tmp, path)?; - if let Some(dir) = path.parent().filter(|d| !d.as_os_str().is_empty()) { - std::fs::File::open(dir)?.sync_all()?; - } - Ok(()) -} - -/// D-S5 status snapshot refresh period. -const STATUS_EVERY: Duration = Duration::from_millis(250); - -/// Largest elastic frame accepted (tensor messages included). -const MAX_ELASTIC_FRAME: u64 = 4 << 30; - -type Outbox = mpsc::UnboundedSender<(u8, Vec)>; - -struct Shared { - coord: ElasticCoordinator, - /// yeto-framework-decoupling 6.2a: backend identity hash pinned by the - /// first accepted JOIN (like the strict HELLO session contract); every - /// later JOIN must carry the same 32 bytes. Kept for the whole run and in - /// the checkpoint, so a Miles session never admits a verl island. - backend_identity: Option<[u8; 32]>, - /// yeto-framework-decoupling 7.7c: compat_group ("-") - /// pinned by the first accepted JOIN. A JOIN from another group is refused - /// with both groups named (phase 1: tolerance table empty). Not in the - /// checkpoint: the group is also inside the identity hash, so after - /// --resume the identity pin still refuses it; the readable pin is set - /// again by the first JOIN whose identity matches. - compat_group: Option, - flushed: usize, - round_started: Instant, - tape: Option, - tape_clock: crate::elastic::TapeClock, - key: Vec, - outer_lr: f32, - outer_momentum: f32, - outer_delayed_nesterov_period: Option, - params: Option>, - momentum_buf: Vec, - /// Delayed Nesterov step count and accumulator (empty when disabled). - delayed_nesterov: crate::merge::DelayedNesterovState, - /// Updates (θ − base) keyed by (island, base_version) awaiting a step. - stored: HashMap<(u32, u64), Vec>, - outboxes: HashMap, - checkpoint_path: Option, - checkpoint_every: u64, - /// D-S5 snapshot and P3 index files, next to --event-tape. - status_path: Option, - sample_index_path: Option, - last_status_write: Option, - total_steps: u64, - /// When FINISHED was broadcast (start of the final grace window). - finished_at: Option, -} - -fn policy_hash(params: &[f32]) -> [u8; 32] { - use sha2::{Digest, Sha256}; - let mut h = Sha256::new(); - params.iter().for_each(|x| h.update(x.to_le_bytes())); - h.finalize().into() -} - -impl Shared { - fn finished(&self) -> bool { - self.coord.outer_version >= self.total_steps - } - - fn finished_msg(&self) -> (u8, Vec) { - ElasticMsg::Finished { - syncer_epoch: self.coord.syncer_epoch, - outer_version: self.coord.outer_version, - policy_hash: self.coord.published.get(&self.coord.outer_version).copied().unwrap_or([0; 32]), - } - .encode(&self.key) - } - - fn publish_hash(&mut self) { - if let Some(p) = &self.params { - let h = policy_hash(p); - self.coord.published.insert(self.coord.outer_version, h); - } - } - - fn write_status(&mut self, now: f64, soft_remaining_s: f64) -> Result<()> { - if let Some(path) = &self.status_path { - write_atomic( - path, - self.coord - .status_json(now, soft_remaining_s, if self.finished() { "finished" } else { "running" }) - .as_bytes(), - ) - .with_context(|| format!("write status {}", path.display()))?; - } - self.last_status_write = Some(Instant::now()); - Ok(()) - } - - fn append_sample_index(&self, e: &crate::elastic::SampleIndexEntry, verdict: u8, reason: &str, lag: u64) -> Result<()> { - if let Some(path) = &self.sample_index_path { - let mut f = std::fs::OpenOptions::new().create(true).append(true).open(path)?; - writeln!( - f, - "{{\"syncer_epoch\":{},\"verdict\":\"{}\",\"reason\":\"{reason}\",\"outer_lag\":{lag},\"entry\":{}}}", - self.coord.syncer_epoch, - crate::elastic::verdict_name(verdict), - e.to_json() - )?; - f.flush()?; - } - Ok(()) - } - - fn encode_checkpoint(&self) -> Vec { - let mut b = SERVER_STATE_MAGIC.to_vec(); - match &self.backend_identity { - Some(h) => { - b.push(1); - b.extend_from_slice(h); - } - None => b.push(0), - } - let coord = self.coord.encode_state(); - b.extend_from_slice(&(coord.len() as u64).to_le_bytes()); - b.extend_from_slice(&coord); - match &self.params { - Some(p) => { - b.push(1); - put_f32s(&mut b, p); - put_f32s(&mut b, &self.momentum_buf); - // Only with Delayed Nesterov, so plain checkpoints keep their bytes. - if let Some(period) = self.outer_delayed_nesterov_period { - b.extend_from_slice(&period.to_le_bytes()); - b.extend_from_slice(&self.delayed_nesterov.step.to_le_bytes()); - put_f32s(&mut b, &self.delayed_nesterov.accum); - } - } - None => b.push(0), - } - let mut keys: Vec<_> = self.stored.keys().copied().collect(); - keys.sort_unstable(); - b.extend_from_slice(&(keys.len() as u32).to_le_bytes()); - for k in keys { - b.extend_from_slice(&k.0.to_le_bytes()); - b.extend_from_slice(&k.1.to_le_bytes()); - put_f32s(&mut b, &self.stored[&k]); - } - b - } - - fn write_checkpoint(&self) -> Result<()> { - if let Some(path) = &self.checkpoint_path { - write_atomic(path, &self.encode_checkpoint()) - .with_context(|| format!("write elastic checkpoint {}", path.display()))?; - } - Ok(()) - } - - /// Restore coordinator, base, momentum and carried tensors. - fn restore(&mut self, bytes: &[u8], lease_s: f64, now: f64) -> Result<()> { - let mut r = Reader(bytes); - let magic = r.take(8)?; - ensure!( - magic != SERVER_STATE_MAGIC_V1, - "elastic syncer checkpoint predates the backend identity field (YELSRV1); \ - it cannot be resumed by this syncer (version boundary s17-elastic-identity)" - ); - ensure!(magic == SERVER_STATE_MAGIC, "not an elastic syncer checkpoint"); - self.backend_identity = match r.u8()? { - 0 => None, - 1 => Some(r.take(32)?.try_into().context("backend identity")?), - v => bail!("invalid backend identity flag {v} in elastic checkpoint"), - }; - let n = usize::try_from(r.u64()?)?; - self.coord = ElasticCoordinator::decode_state(self.coord.params, lease_s, r.take(n)?, now)?; - if r.u8()? == 1 { - self.params = Some(get_f32s(&mut r)?); - self.momentum_buf = get_f32s(&mut r)?; - if let Some(period) = self.outer_delayed_nesterov_period { - ensure!( - r.u32()? == period, - "elastic checkpoint delayed Nesterov period does not match --outer-delayed-nesterov-period" - ); - let step = r.u64()?; - let accum = get_f32s(&mut r)?; - ensure!( - accum.len() == self.momentum_buf.len(), - "elastic checkpoint has no matching delayed Nesterov state" - ); - self.delayed_nesterov = crate::merge::DelayedNesterovState { step, accum: accum.into() }; - } - } - for _ in 0..r.u32()? { - let k = (r.u32()?, r.u64()?); - self.stored.insert(k, get_f32s(&mut r)?); - } - ensure!(r.remaining() == 0, "trailing bytes in elastic syncer checkpoint"); - Ok(()) - } - - fn flush_tape(&mut self) -> Result<()> { - let new = &self.coord.tape[self.flushed..]; - if let Some(f) = self.tape.as_mut() { - for ev in new { - writeln!(f, "{}", ev.to_json_stamped(self.coord.syncer_epoch, &mut self.tape_clock, unix_now()))?; - } - f.flush()?; - } - for ev in new { - info!(kind = ev.kind(), outer_version = self.coord.outer_version, "elastic event"); - } - self.flushed = self.coord.tape.len(); - Ok(()) - } - - fn base_msg(&self) -> Option<(u8, Vec)> { - self.params.as_ref().map(|p| { - ElasticMsg::ElasticBase { - syncer_epoch: self.coord.syncer_epoch, - outer_version: self.coord.outer_version, - params: p.clone(), - } - .encode(&self.key) - }) - } - - /// Bandwidth record for one tensor frame (ELASTIC_INIT / DELTA_TENSOR - /// received, ELASTIC_BASE sent): frame bytes incl. the 13-byte header - /// and wall seconds from header to last payload byte. Written straight - /// to the event tape as `kind:"transfer"` (not part of the coordinator - /// tape, so checkpoints and ledger golden cases are unchanged). - fn record_transfer(&mut self, direction: &str, msg_type: u8, island: Option, bytes: u64, seconds: f64) -> Result<()> { - let name = match msg_type { - MSG_ELASTIC_INIT => "elastic_init", - MSG_DELTA_TENSOR => "delta_tensor", - MSG_ELASTIC_BASE => "elastic_base", - _ => "other", - }; - let island_json = island.map(|i| i.to_string()).unwrap_or_else(|| "null".into()); - let secs = if seconds.is_finite() { format!("{seconds:?}") } else { "null".into() }; - let line = format!( - "{{\"kind\":\"transfer\",\"syncer_epoch\":{},\"direction\":\"{direction}\",\"msg\":\"{name}\",\"island_id\":{island_json},\"bytes\":{bytes},\"seconds\":{secs},\"outer_version\":{},\"time_unix\":{:?}}}", - self.coord.syncer_epoch, self.coord.outer_version, unix_now() - ); - if let Some(f) = self.tape.as_mut() { - writeln!(f, "{line}")?; - f.flush()?; - } - info!(direction, msg = name, island = ?island, bytes, seconds, "elastic transfer"); - Ok(()) - } - - /// Forget updates from islands that are no longer members. - fn purge_departed(&mut self) { - let coord = &self.coord; - self.stored.retain(|(i, _), _| coord.is_member(*i)); - self.outboxes.retain(|i, _| coord.is_member(*i)); - } - - /// Step attempt; on a step, merge the planned updates into the base and - /// broadcast it. Resets the soft-deadline clock on a step or timeout. - fn advance(&mut self, timed_out: bool) -> Result<()> { - let Some(plan) = self.coord.try_advance(timed_out) else { - if timed_out { - self.round_started = Instant::now(); - } - return Ok(()); - }; - self.round_started = Instant::now(); - if let Some(params) = self.params.as_mut() { - let mut grads = Vec::with_capacity(plan.weights.len()); - let mut weights = Vec::with_capacity(plan.weights.len()); - for &(island, base, w) in &plan.weights { - let mut u = self.stored.remove(&(island, base)).with_context(|| { - format!("no DELTA_TENSOR for island {island} base {base} in a tensor run") - })?; - // Wire carries θ − base; merge.rs consumes outer gradients. - u.iter_mut().for_each(|x| *x = -*x); - grads.push(u); - weights.push(w); - } - let mut merged = vec![0.0f32; params.len()]; - let refs: Vec<&[f32]> = grads.iter().map(Vec::as_slice).collect(); - crate::merge::merge_avg(&refs, &weights, &mut merged); - match self.outer_delayed_nesterov_period { - None => crate::merge::nesterov_step( - params, - &mut self.momentum_buf, - &merged, - self.outer_lr, - self.outer_momentum, - ), - Some(period) => crate::merge::delayed_nesterov_step( - params, - &mut self.momentum_buf, - &mut self.delayed_nesterov, - &merged, - self.outer_lr, - self.outer_momentum, - period, - ), - } - } - // Every pending/carried entry was consumed or dropped by the step. - self.stored.clear(); - self.publish_hash(); - if self.checkpoint_every > 0 && self.coord.outer_version % self.checkpoint_every == 0 { - self.write_checkpoint()?; - } - if let Some(frame) = self.base_msg() { - self.outboxes.retain(|_, tx| tx.send(frame.clone()).is_ok()); - } - Ok(()) - } - - /// Handle one authenticated message; returns frames for the sender. - fn handle(&mut self, msg: ElasticMsg, now: f64, outbox: &Outbox) -> Result)>> { - self.coord.check_fence(msg.syncer_epoch())?; - let mut replies = Vec::new(); - if self.finished() { - // Final grace window: LEAVE is still applied; everything else - // gets FINISHED so the island can exit 0. - if let ElasticMsg::Leave { .. } = msg { - self.coord.apply(&msg, now)?; - self.purge_departed(); - } else { - replies.push(self.finished_msg()); - } - return Ok(replies); - } - match &msg { - ElasticMsg::ElasticInit { island_id, params, .. } => { - ensure!(self.coord.is_member(*island_id), "{island_id} is not a member"); - if self.params.is_none() { - ensure!(!params.is_empty(), "ELASTIC_INIT with empty parameters"); - self.momentum_buf = vec![0.0; params.len()]; - if self.outer_delayed_nesterov_period.is_some() { - self.delayed_nesterov.accum = vec![0.0; params.len()].into(); - } - self.params = Some(params.clone()); - self.publish_hash(); - } - replies.extend(self.base_msg()); - return Ok(replies); - } - ElasticMsg::DeltaTensor { island_id, base_version, update, .. } => { - let p = self.params.as_ref().context("DELTA_TENSOR before ELASTIC_INIT")?; - ensure!(update.len() == p.len(), "DELTA_TENSOR length {} != {}", update.len(), p.len()); - ensure!(self.coord.is_member(*island_id), "{island_id} is not a member"); - self.stored.insert((*island_id, *base_version), update.clone()); - } - ElasticMsg::DeltaReady { .. } if self.params.is_some() => { - bail!("DELTA_READY carries no tensor; use DELTA_TENSOR once a base exists") - } - ElasticMsg::Join { island_id, backend_identity, compat_group, .. } => { - if let Some(pinned) = &self.compat_group { - ensure!( - pinned == compat_group, - "compat_group mismatch, JOIN refused: 兼容组不同：{} 对 {}，容差未标定 \ - (island {island_id} declares {}, this elastic session is pinned to {} by its \ - first JOIN; tolerance not calibrated)", - group_label(pinned), - group_label(compat_group), - group_label(compat_group), - group_label(pinned) - ); - } - if let Some(pinned) = &self.backend_identity { - ensure!( - pinned == backend_identity, - "backend identity mismatch, JOIN refused: island {island_id} declares {} \ - but this elastic session is pinned to {} (head config or first JOIN); islands with a \ - different training backend (e.g. Miles vs verl), engine commit, device \ - family or parameter-name map cannot be merged", - hex32(backend_identity), - hex32(pinned) - ); - } - } - _ => {} - } - let result = self.coord.apply(&msg, now); - if let ElasticMsg::DeltaTensor { island_id, base_version, .. } = &msg { - if result.is_err() { - self.stored.remove(&(*island_id, *base_version)); - } - } - let reply = result?; - match &msg { - ElasticMsg::Join { island_id, backend_identity, compat_group, .. } => { - self.backend_identity.get_or_insert(*backend_identity); - self.compat_group.get_or_insert_with(|| compat_group.clone()); - self.outboxes.insert(*island_id, outbox.clone()); - replies.extend(reply.map(|m| m.encode(&self.key))); - replies.extend(self.base_msg()); - } - ElasticMsg::Leave { .. } => self.purge_departed(), - ElasticMsg::SampleIndex { entry, .. } => { - if let Some(crate::elastic::TapeEvent::SampleIndexed { verdict, reason, outer_lag, duplicate: false, .. }) = - self.coord.tape.last() - { - if *verdict != crate::elastic::VERDICT_REJECT { - self.append_sample_index(entry, *verdict, reason, *outer_lag)?; - } - } - replies.extend(reply.map(|m| m.encode(&self.key))); - } - ElasticMsg::DeltaReady { .. } | ElasticMsg::DeltaTensor { .. } => self.advance(false)?, - _ => replies.extend(reply.map(|m| m.encode(&self.key))), - } - Ok(replies) - } -} - -fn island_of(cell: &AtomicU64) -> Option { - u32::try_from(cell.load(Ordering::Relaxed)).ok() -} - -fn msg_island(m: &ElasticMsg) -> Option { - match m { - ElasticMsg::Join { island_id, .. } - | ElasticMsg::Leave { island_id, .. } - | ElasticMsg::LeaseHeartbeat { island_id, .. } - | ElasticMsg::DeltaReady { island_id, .. } - | ElasticMsg::ElasticInit { island_id, .. } - | ElasticMsg::DeltaTensor { island_id, .. } => Some(*island_id), - _ => None, - } -} - -async fn send(wr: &mut tokio::net::tcp::OwnedWriteHalf, t: u8, p: &[u8]) -> Result<()> { - let mut header = [0u8; 13]; - header[0..4].copy_from_slice(&MAGIC.to_le_bytes()); - header[4] = t; - header[5..13].copy_from_slice(&(p.len() as u64).to_le_bytes()); - wr.write_all(&header).await?; - wr.write_all(p).await?; - Ok(()) -} - -async fn serve_conn( - stream: TcpStream, - shared: Arc>, - key: Arc>, - epoch0: Instant, - fatal: mpsc::UnboundedSender, -) -> Result<()> { - let (mut rd, mut wr) = stream.into_split(); - let (tx, mut rx) = mpsc::unbounded_channel::<(u8, Vec)>(); - // Island on this connection, learned from the first decoded message. - let island = Arc::new(AtomicU64::new(u64::MAX)); - let (w_shared, w_island) = (shared.clone(), island.clone()); - let writer = tokio::spawn(async move { - while let Some((t, p)) = rx.recv().await { - let started = Instant::now(); - if send(&mut wr, t, &p).await.is_err() { - break; - } - if t == MSG_ELASTIC_BASE { - let secs = started.elapsed().as_secs_f64(); - let isl = island_of(&w_island); - let _ = w_shared.lock().unwrap().record_transfer("send", t, isl, p.len() as u64 + 13, secs); - } - } - }); - let result = async { - loop { - let mut header_at: Option = None; - let frame = match read_frame_limited(&mut rd, |t| { - header_at = Some(Instant::now()); - match t { - MSG_JOIN | MSG_LEAVE | MSG_LEASE_HEARTBEAT | MSG_SAMPLE_INDEX | MSG_DELTA_READY - | MSG_ELASTIC_INIT | MSG_DELTA_TENSOR => Ok(MAX_ELASTIC_FRAME), - other => bail!("message type {other} is not accepted by the elastic server"), - } - }) - .await - { - Ok(f) => f, - Err(e) => { - if e.downcast_ref::() - .is_some_and(|io| io.kind() == std::io::ErrorKind::UnexpectedEof) - { - return Ok(()); - } - let _ = tx.send((MSG_ERROR, format!("{e:#}").into_bytes())); - return Err(e); - } - }; - let now = epoch0.elapsed().as_secs_f64(); - let recv_secs = header_at.map(|t| t.elapsed().as_secs_f64()); - let outcome = ElasticMsg::decode(&key, frame.msg_type, &frame.payload).and_then(|msg| { - if let Some(i) = msg_island(&msg) { - island.store(u64::from(i), Ordering::Relaxed); - } - let mut g = shared.lock().unwrap(); - if matches!(frame.msg_type, MSG_ELASTIC_INIT | MSG_DELTA_TENSOR) { - let bytes = frame.payload.len() as u64 + 13; - g.record_transfer("recv", frame.msg_type, island_of(&island), bytes, recv_secs.unwrap_or(f64::NAN))?; - } - let r = g.handle(msg, now, &tx); - g.flush_tape()?; - r - }); - match outcome { - Ok(frames) => { - for f in frames { - let _ = tx.send(f); - } - } - Err(e) if format!("{e:#}").contains("no DELTA_TENSOR") => { - // Merge invariant broken: stop the whole server. - let _ = fatal.send(anyhow::anyhow!("{e:#}")); - return Err(e); - } - Err(e) => { - // Fenced, unauthenticated or rejected: report and keep - // the connection (the island decides whether to rejoin). - warn!(error = %format!("{e:#}"), "elastic message refused"); - let _ = tx.send((MSG_ERROR, format!("{e:#}").into_bytes())); - } - } - } - } - .await; - drop(tx); - let _ = writer.await; - result -} - -/// Serve until `total_steps` outer steps have been published. Returns the -/// final outer version and parameters (None for a control-plane-only run). -pub async fn run(listener: TcpListener, cfg: ElasticServerConfig) -> Result<(u64, Option>)> { - let tape = match &cfg.event_tape { - Some(p) => Some( - std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(p) - .with_context(|| format!("open event tape {}", p.display()))?, - ), - None => None, - }; - let shared = Arc::new(Mutex::new(Shared { - coord: ElasticCoordinator::new(cfg.params, cfg.lease_s, cfg.syncer_epoch)?, - backend_identity: cfg.expected_backend_identity, - compat_group: None, - flushed: 0, - round_started: Instant::now(), - tape_clock: crate::elastic::TapeClock::new(unix_now()), - tape, - key: cfg.key.clone(), - outer_lr: cfg.outer_lr, - outer_momentum: cfg.outer_momentum, - outer_delayed_nesterov_period: cfg.outer_delayed_nesterov_period, - params: None, - momentum_buf: Vec::new(), - delayed_nesterov: Default::default(), - stored: HashMap::new(), - outboxes: HashMap::new(), - checkpoint_path: cfg.checkpoint_path.clone(), - checkpoint_every: cfg.checkpoint_every, - total_steps: cfg.total_steps, - finished_at: None, - status_path: cfg.event_tape.as_ref().map(|t| t.with_file_name("status.json")), - sample_index_path: cfg.event_tape.as_ref().map(|t| t.with_file_name("sample_index.jsonl")), - last_status_write: None, - })); - if cfg.resume { - let path = cfg.checkpoint_path.as_ref().context("--resume requires --checkpoint-path")?; - let bytes = std::fs::read(path) - .with_context(|| format!("read elastic checkpoint {}", path.display()))?; - let mut g = shared.lock().unwrap(); - g.restore(&bytes, cfg.lease_s, 0.0)?; - if let Some(pinned) = cfg.expected_backend_identity { - match g.backend_identity { - Some(saved) if saved != pinned => bail!( - "elastic checkpoint pins backend identity {} but the head config pins {}", - hex32(&saved), - hex32(&pinned) - ), - _ => g.backend_identity = Some(pinned), - } - } - g.publish_hash(); - // P9: every restart is a new coordinator incarnation. - let epoch = (g.coord.syncer_epoch + 1).max(cfg.syncer_epoch); - g.coord.syncer_epoch = epoch; - g.write_checkpoint()?; - info!(syncer_epoch = epoch, outer_version = g.coord.outer_version, "elastic syncer resumed"); - } - let key = Arc::new(cfg.key); - let epoch0 = Instant::now(); - let soft = Duration::from_secs(cfg.params.soft_deadline_s); - let (fatal_tx, mut fatal_rx) = mpsc::unbounded_channel(); - info!(syncer_epoch = cfg.syncer_epoch, ?soft, lease_s = cfg.lease_s, "elastic syncer listening"); - let mut ticker = tokio::time::interval(cfg.tick); - loop { - tokio::select! { - accepted = listener.accept() => { - let (stream, peer) = accepted?; - stream.set_nodelay(true).ok(); - let (s, k, f) = (shared.clone(), key.clone(), fatal_tx.clone()); - tokio::spawn(async move { - if let Err(e) = serve_conn(stream, s, k, epoch0, f).await { - warn!(%peer, error = %format!("{e:#}"), "elastic connection closed"); - } - }); - } - Some(e) = fatal_rx.recv() => return Err(e), - _ = ticker.tick() => { - let mut g = shared.lock().unwrap(); - if !g.coord.expire_leases(epoch0.elapsed().as_secs_f64()).is_empty() { - g.purge_departed(); - } - let timed_out = g.round_started.elapsed() >= soft; - if !g.finished() { - g.advance(timed_out)?; - } - g.flush_tape()?; - if g.last_status_write.is_none_or(|t| t.elapsed() >= STATUS_EVERY) { - let remaining = soft.saturating_sub(g.round_started.elapsed()).as_secs_f64(); - g.write_status(epoch0.elapsed().as_secs_f64(), remaining)?; - } - if g.finished() { - if g.finished_at.is_none() { - let frame = g.finished_msg(); - g.outboxes.retain(|_, tx| tx.send(frame.clone()).is_ok()); - g.finished_at = Some(Instant::now()); - g.write_status(epoch0.elapsed().as_secs_f64(), 0.0)?; - info!(outer_version = g.coord.outer_version, grace = ?cfg.final_grace, "elastic run finished; final grace window"); - } - let grace_over = g.finished_at.is_some_and(|t| t.elapsed() >= cfg.final_grace); - if grace_over || g.coord.member_count() == 0 { - g.write_status(epoch0.elapsed().as_secs_f64(), 0.0)?; - return Ok((g.coord.outer_version, g.params.clone())); - } - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; - - const KEY: &[u8] = b"test-key"; - const MILES_ID: [u8; 32] = [0x11; 32]; - - struct Island { - id: u32, - epoch: u64, - rd: OwnedReadHalf, - wr: OwnedWriteHalf, - identity: [u8; 32], - compat: String, - } - - impl Island { - async fn connect(port: u16, id: u32, epoch: u64) -> Self { - let (rd, wr) = TcpStream::connect(("127.0.0.1", port)).await.unwrap().into_split(); - Self { id, epoch, rd, wr, identity: MILES_ID, compat: "nvidia-h100".to_string() } - } - async fn send(&mut self, m: ElasticMsg) { - let (t, p) = m.encode(KEY); - send(&mut self.wr, t, &p).await.unwrap(); - } - async fn recv(&mut self) -> (u8, Vec) { - let f = read_frame_limited(&mut self.rd, |_| Ok(MAX_ELASTIC_FRAME)).await.unwrap(); - (f.msg_type, f.payload) - } - async fn join(&mut self) -> ElasticMsg { - self.send(ElasticMsg::Join { - syncer_epoch: self.epoch, - island_id: self.id, - incarnation: 0, - capacity: 1.0, - backend_identity: self.identity, - compat_group: self.compat.clone(), - }) - .await; - let (t, p) = self.recv().await; - ElasticMsg::decode(KEY, t, &p).unwrap() - } - async fn delta(&mut self, base_version: u64) { - self.send(ElasticMsg::DeltaReady { - syncer_epoch: self.epoch, - island_id: self.id, - base_version, - c_tokens: 10, - c_steps: 1, - }) - .await; - } - async fn init(&mut self, params: Vec) { - self.send(ElasticMsg::ElasticInit { syncer_epoch: self.epoch, island_id: self.id, params }).await; - } - async fn update(&mut self, base_version: u64, value: f32) { - self.send(ElasticMsg::DeltaTensor { - syncer_epoch: self.epoch, - island_id: self.id, - base_version, - c_tokens: 10, - c_steps: 1, - update: vec![value; 4], - }) - .await; - } - /// Next ELASTIC_BASE frame (skips other frames; panics on MSG_ERROR). - async fn next_base(&mut self) -> (u64, Vec) { - loop { - let (t, p) = tokio::time::timeout(Duration::from_secs(5), self.recv()).await.unwrap(); - assert_ne!(t, MSG_ERROR, "{}", String::from_utf8_lossy(&p)); - if t == MSG_ELASTIC_BASE { - match ElasticMsg::decode(KEY, t, &p).unwrap() { - ElasticMsg::ElasticBase { outer_version, params, .. } => { - return (outer_version, params) - } - _ => unreachable!(), - } - } - } - } - async fn heartbeat(&mut self) { - self.send(ElasticMsg::LeaseHeartbeat { - syncer_epoch: self.epoch, - island_id: self.id, - membership_epoch: 0, - inner_step: 0, - round_wall_s: 0.0, - }) - .await; - } - } - - fn tape_lines(path: &std::path::Path) -> Vec { - std::fs::read_to_string(path).unwrap_or_default().lines().map(str::to_string).collect() - } - - async fn wait_for(path: &std::path::Path, needle: &str, count: usize) { - for _ in 0..200 { - if tape_lines(path).iter().filter(|l| l.contains(needle)).count() >= count { - return; - } - tokio::time::sleep(Duration::from_millis(20)).await; - } - panic!("timed out waiting for {count}x {needle}: {:#?}", tape_lines(path)); - } - - /// 3 islands; one is late (carried over), one stops heartbeating and is - /// expired with its uncommitted delta dropped, one joins mid-run with - /// zero weight in its first round; a stale-epoch message is fenced. - #[tokio::test] - async fn three_islands_late_expiry_and_midrun_join() { - let dir = std::env::temp_dir().join(format!("elastic-it-{}", std::process::id())); - std::fs::create_dir_all(&dir).unwrap(); - let tape = dir.join("tape.jsonl"); - let _ = std::fs::remove_file(&tape); - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run( - listener, - ElasticServerConfig { - params: ElasticParams { - quorum_theta: 0.6, - carry_gamma: 0.5, - soft_deadline_s: 60, - q_min: 1, - max_carry_lag: 2, - }, - key: KEY.to_vec(), - syncer_epoch: 7, - lease_s: 0.6, - total_steps: 3, - event_tape: Some(tape.clone()), - tick: Duration::from_millis(20), - outer_lr: 1.0, - outer_momentum: 0.0, - outer_delayed_nesterov_period: None, - final_grace: Duration::ZERO, - checkpoint_path: None, - checkpoint_every: 0, - resume: false, - expected_backend_identity: None, - }, - )); - let (mut a, mut b, mut c) = ( - Island::connect(port, 1, 7).await, - Island::connect(port, 2, 7).await, - Island::connect(port, 3, 7).await, - ); - for i in [&mut a, &mut b, &mut c] { - assert!(matches!(i.join().await, ElasticMsg::JoinAck { catch_up: false, .. })); - } - - // Fencing: an older syncer_epoch is refused with MSG_ERROR. - let mut stale = Island::connect(port, 9, 6).await; - stale.send(ElasticMsg::Join { syncer_epoch: 6, island_id: 9, incarnation: 0, capacity: 1.0, backend_identity: MILES_ID, compat_group: String::new() }).await; - let (t, p) = stale.recv().await; - assert_eq!(t, MSG_ERROR); - assert!(String::from_utf8_lossy(&p).contains("fenced")); - - // Round 0: A and B (2/3 >= 0.6) step without C. - a.delta(0).await; - b.delta(0).await; - wait_for(&tape, "\"kind\":\"outer_step\"", 1).await; - // C is late on base 0 (lag 1) -> carried over; round 1 A+B+carried C. - c.delta(0).await; - wait_for(&tape, "delta_carried_over", 1).await; - a.delta(1).await; - b.delta(1).await; - wait_for(&tape, "\"kind\":\"outer_step\"", 2).await; - // C submits on base 2 then goes silent; A, B keep heartbeating. - c.delta(2).await; - wait_for(&tape, "\"kind\":\"delta_accepted\",\"syncer_epoch\":7,\"island_id\":3", 1).await; - for _ in 0..10 { - a.heartbeat().await; - b.heartbeat().await; - tokio::time::sleep(Duration::from_millis(100)).await; - } - wait_for(&tape, "lease_expired", 1).await; - // D joins mid-run (catch-up), then A, B, D finish round 2. - let mut d = Island::connect(port, 4, 7).await; - assert!(matches!(d.join().await, ElasticMsg::JoinAck { catch_up: true, base_version: 2, .. })); - a.delta(2).await; - b.delta(2).await; - d.delta(2).await; - let (final_version, params) = tokio::time::timeout(Duration::from_secs(10), server).await.unwrap().unwrap().unwrap(); - assert!(params.is_none(), "control-plane-only run has no tensors"); - assert!(final_version >= 3); - - let status = std::fs::read_to_string(dir.join("status.json")).unwrap(); - assert!(status.contains("\"schema\":\"yeto.syncer.elastic-status/v1\""), "{status}"); - assert!(status.contains("\"syncer_epoch\":7"), "{status}"); - let lines = tape_lines(&tape); - let steps: Vec<&String> = lines.iter().filter(|l| l.contains("\"kind\":\"outer_step\"")).collect(); - assert!(steps[0].contains("\"absent\":[3]"), "{}", steps[0]); - // Round 1: A,B at 100 each, C carried at 100*0.5 on base 0. - assert!(steps[1].contains("[3,0,50.0]"), "{}", steps[1]); - assert!(steps[1].contains("\"weights\":[[1,1,0.4],[2,1,0.4],[3,0,0.2]]"), "{}", steps[1]); - let expired = lines.iter().find(|l| l.contains("lease_expired")).unwrap(); - assert!(expired.contains("\"island_id\":3"), "{expired}"); - assert!(expired.contains("\"dropped_uncommitted\":{\"island_id\":3,\"outer_version\":2"), "{expired}"); - assert!(lines.iter().any(|l| l.contains("pool_join") && l.contains("\"island_id\":4") && l.contains("\"catch_up\":true"))); - // Round 2 may step on A+B (2/3) before D's delta arrives; if D is in - // the round its weight must be 0. - assert!(!steps[2].contains("[4,2,100.0]"), "{}", steps[2]); - assert!(!steps[2].contains("[3,"), "expired island must not be merged: {}", steps[2]); - let _ = std::fs::remove_dir_all(&dir); - } - - /// Bandwidth records: every ELASTIC_INIT / DELTA_TENSOR received and - /// ELASTIC_BASE sent lands on the event tape as kind "transfer" with the - /// island, the exact frame bytes and a finite duration. - #[tokio::test] - async fn tensor_frames_record_transfer_bytes_and_seconds() { - let dir = std::env::temp_dir().join(format!("elastic-xfer-{}", std::process::id())); - std::fs::create_dir_all(&dir).unwrap(); - let tape = dir.join("tape.jsonl"); - let _ = std::fs::remove_file(&tape); - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run( - listener, - ElasticServerConfig { - params: ElasticParams { quorum_theta: 0.6, carry_gamma: 0.5, soft_deadline_s: 60, q_min: 1, max_carry_lag: 2 }, - key: KEY.to_vec(), - syncer_epoch: 3, - lease_s: 30.0, - total_steps: 1, - event_tape: Some(tape.clone()), - tick: Duration::from_millis(20), - outer_lr: 1.0, - outer_momentum: 0.0, - outer_delayed_nesterov_period: None, - final_grace: Duration::ZERO, - checkpoint_path: None, - checkpoint_every: 0, - resume: false, - expected_backend_identity: None, - }, - )); - let mut a = Island::connect(port, 5, 3).await; - a.join().await; - a.init(vec![0.0; 4]).await; - assert_eq!(a.next_base().await.0, 0); - a.update(0, 1.0).await; - assert_eq!(a.next_base().await.0, 1); - wait_for(&tape, "\"msg\":\"elastic_base\"", 2).await; - drop(a); - let _ = tokio::time::timeout(Duration::from_secs(5), server).await; - - let frame_len = |m: ElasticMsg| m.encode(KEY).1.len() as u64 + 13; - let init_bytes = frame_len(ElasticMsg::ElasticInit { syncer_epoch: 3, island_id: 5, params: vec![0.0; 4] }); - let delta_bytes = frame_len(ElasticMsg::DeltaTensor { - syncer_epoch: 3, island_id: 5, base_version: 0, c_tokens: 10, c_steps: 1, update: vec![1.0; 4], - }); - let lines: Vec = tape_lines(&tape).into_iter().filter(|l| l.contains("\"kind\":\"transfer\"")).collect(); - let find = |dir: &str, msg: &str| -> Vec { - lines.iter().filter(|l| l.contains(&format!("\"direction\":\"{dir}\",\"msg\":\"{msg}\""))).cloned().collect() - }; - let init = find("recv", "elastic_init"); - assert_eq!(init.len(), 1, "{lines:#?}"); - assert!(init[0].contains(&format!("\"bytes\":{init_bytes},")), "{}", init[0]); - let delta = find("recv", "delta_tensor"); - assert_eq!(delta.len(), 1, "{lines:#?}"); - assert!(delta[0].contains(&format!("\"bytes\":{delta_bytes},")), "{}", delta[0]); - let bases = find("send", "elastic_base"); - assert!(bases.len() >= 2, "{lines:#?}"); - for l in init.iter().chain(&delta).chain(&bases) { - assert!(l.contains("\"island_id\":5,"), "{l}"); - assert!(l.contains("\"syncer_epoch\":3,"), "{l}"); - assert!(!l.contains("\"seconds\":null"), "{l}"); - } - let _ = std::fs::remove_dir_all(&dir); - } - - fn assert_close(got: &[f32], want: f32) { - for g in got { - assert!((g - want).abs() < 1e-5, "got {got:?}, want {want}"); - } - } - - /// 3 islands, 3 tensor rounds, outer_lr 1 / momentum 0 so the new base - /// is base + weighted mean of updates. Round 0: A=1, B=2 (C absent) - /// -> 1.5. Round 1: C late on base 0 (4, weight 100*0.5), A=1, B=3 -> - /// 1.5 + (100+300+200)/250 = 3.9. Round 2: C sends 100 on base 2 then - /// expires (never merged); D joins and sends 1000 (weight 0, first - /// round), then A=1 reaches capacity 2/3 >= 0.6 -> 3.9 + 1 = 4.9. - #[tokio::test] - async fn tensor_rounds_match_hand_computed_weights() { - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run( - listener, - ElasticServerConfig { - params: ElasticParams { - quorum_theta: 0.6, - carry_gamma: 0.5, - soft_deadline_s: 60, - q_min: 1, - max_carry_lag: 2, - }, - key: KEY.to_vec(), - syncer_epoch: 1, - lease_s: 0.6, - total_steps: 3, - event_tape: None, - tick: Duration::from_millis(20), - outer_lr: 1.0, - outer_momentum: 0.0, - outer_delayed_nesterov_period: None, - final_grace: Duration::ZERO, - checkpoint_path: None, - checkpoint_every: 0, - resume: false, - expected_backend_identity: None, - }, - )); - let (mut a, mut b, mut c) = ( - Island::connect(port, 1, 1).await, - Island::connect(port, 2, 1).await, - Island::connect(port, 3, 1).await, - ); - a.join().await; - a.init(vec![0.0; 4]).await; - assert_eq!(a.next_base().await, (0, vec![0.0; 4])); - b.join().await; - c.join().await; - assert_eq!(b.next_base().await.0, 0); - assert_eq!(c.next_base().await.0, 0); - - a.update(0, 1.0).await; - b.update(0, 2.0).await; - let (v, p) = a.next_base().await; - assert_eq!(v, 1); - assert_close(&p, 1.5); - assert_eq!(b.next_base().await.0, 1); - assert_eq!(c.next_base().await.0, 1, "absent member still receives the base"); - - c.update(0, 4.0).await; // late: lag 1, discounted by 0.5 - tokio::time::sleep(Duration::from_millis(50)).await; - a.update(1, 1.0).await; - b.update(1, 3.0).await; - let (v, p) = a.next_base().await; - assert_eq!(v, 2); - assert_close(&p, 3.9); - b.next_base().await; - c.next_base().await; - - c.update(2, 100.0).await; // then C goes silent and expires - for _ in 0..10 { - a.heartbeat().await; - b.heartbeat().await; - tokio::time::sleep(Duration::from_millis(100)).await; - } - let mut d = Island::connect(port, 4, 1).await; - assert!(matches!(d.join().await, ElasticMsg::JoinAck { catch_up: true, base_version: 2, .. })); - assert_eq!(d.next_base().await.0, 2); - d.update(2, 1000.0).await; // joiner's first round: weight 0 - tokio::time::sleep(Duration::from_millis(50)).await; - a.update(2, 1.0).await; - b.update(2, 1.0).await; - let (v, p) = a.next_base().await; - assert_eq!(v, 3); - // D at weight 0, C's 100 dropped with its lease: mean of A (and B) = 1. - assert_close(&p, 4.9); - let (final_version, params) = - tokio::time::timeout(Duration::from_secs(10), server).await.unwrap().unwrap().unwrap(); - assert_eq!(final_version, 3); - assert_close(¶ms.unwrap(), 4.9); - } - - /// rl-spot-cost-saving 5.2 (CPU part): anchor A (1) + droppable B (2). - /// Round 0: A=1, B=3 -> 2. B sends 100 on base 1 then LEAVEs (reclaim): - /// its delta is dropped_uncommitted. B's process is still alive and sends - /// 500 on base 1 after the LEAVE: refused (not a member), never merged. - /// A advances alone: 2+1=3, 3+1=4. B comes back on a new connection: - /// JOIN catch_up=true at base 3 and receives the syncer base (4.0); its - /// first-round delta (1000) has weight 0 -> 4+1=5. - /// rl-island-graceful-leave 4.3: an island leaving with any reason code - /// (and an unknown code) is removed at once; the other island keeps - /// producing rounds well before the 60 s soft deadline, and the tape names - /// each code. - #[tokio::test] - async fn leave_reasons_are_logged_by_name_and_the_other_island_continues() { - use crate::elastic::*; - let dir = std::env::temp_dir().join(format!("elastic-leave-reasons-{}", std::process::id())); - std::fs::create_dir_all(&dir).unwrap(); - let tape = dir.join("tape.jsonl"); - let _ = std::fs::remove_file(&tape); - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let reasons = [LEAVE_REASON_FINISHED, LEAVE_REASON_STOPPED, LEAVE_REASON_PREEMPTED, - LEAVE_REASON_INFRA_FAILED, LEAVE_REASON_TRAINING_FAILED, - LEAVE_REASON_CONFIG_MISMATCH, 42]; - let server = tokio::spawn(run( - listener, - ElasticServerConfig { - params: ElasticParams { quorum_theta: 0.75, carry_gamma: 0.5, soft_deadline_s: 60, q_min: 1, max_carry_lag: 2 }, - key: KEY.to_vec(), - syncer_epoch: 1, - lease_s: 30.0, - total_steps: reasons.len() as u64, - event_tape: Some(tape.clone()), - tick: Duration::from_millis(20), - outer_lr: 1.0, - outer_momentum: 0.0, - outer_delayed_nesterov_period: None, - final_grace: Duration::ZERO, - checkpoint_path: None, - checkpoint_every: 0, - resume: false, - expected_backend_identity: None, - }, - )); - let mut a = Island::connect(port, 1, 1).await; - a.join().await; - a.init(vec![0.0; 4]).await; - assert_eq!(a.next_base().await.0, 0); - for (version, reason) in reasons.iter().enumerate() { - let mut b = Island::connect(port, 2, 1).await; - b.join().await; - b.send(ElasticMsg::Leave { syncer_epoch: 1, island_id: 2, reason: *reason }).await; - wait_for(&tape, "pool_leave", version + 1).await; - a.update(version as u64, 1.0).await; - let (v, _) = tokio::time::timeout(Duration::from_secs(5), a.next_base()).await.unwrap(); - assert_eq!(v, version as u64 + 1); - } - let (final_version, _) = - tokio::time::timeout(Duration::from_secs(10), server).await.unwrap().unwrap().unwrap(); - assert_eq!(final_version, reasons.len() as u64); - let names: Vec = tape_lines(&tape) - .iter() - .filter(|l| l.contains("pool_leave")) - .map(|l| l.split("\"reason\":\"").nth(1).unwrap().split('"').next().unwrap().to_string()) - .collect(); - assert_eq!(names, ["finished", "stopped", "preempted", "infra_failed", "training_failed", - "config_mismatch", "unknown(42)"]); - } - - #[tokio::test] - async fn droppable_leave_rejects_late_delta_and_anchor_advances_alone() { - let dir = std::env::temp_dir().join(format!("elastic-droppable-{}", std::process::id())); - std::fs::create_dir_all(&dir).unwrap(); - let tape = dir.join("tape.jsonl"); - let _ = std::fs::remove_file(&tape); - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run( - listener, - ElasticServerConfig { - params: ElasticParams { quorum_theta: 0.75, carry_gamma: 0.5, soft_deadline_s: 60, q_min: 1, max_carry_lag: 2 }, - key: KEY.to_vec(), - syncer_epoch: 1, - lease_s: 30.0, - total_steps: 4, - event_tape: Some(tape.clone()), - tick: Duration::from_millis(20), - outer_lr: 1.0, - outer_momentum: 0.0, - outer_delayed_nesterov_period: None, - final_grace: Duration::ZERO, - checkpoint_path: None, - checkpoint_every: 0, - resume: false, - expected_backend_identity: None, - }, - )); - let (mut a, mut b) = (Island::connect(port, 1, 1).await, Island::connect(port, 2, 1).await); - a.join().await; - a.init(vec![0.0; 4]).await; - assert_eq!(a.next_base().await.0, 0); - b.join().await; - assert_eq!(b.next_base().await.0, 0); - a.update(0, 1.0).await; - b.update(0, 3.0).await; - let (v, p) = a.next_base().await; - assert_eq!(v, 1); - assert_close(&p, 2.0); - assert_eq!(b.next_base().await.0, 1); - - b.update(1, 100.0).await; - wait_for(&tape, "\"kind\":\"delta_accepted\",\"syncer_epoch\":1,\"island_id\":2", 2).await; - b.send(ElasticMsg::Leave { syncer_epoch: 1, island_id: 2, reason: crate::elastic::LEAVE_REASON_FINISHED }).await; - wait_for(&tape, "pool_leave", 1).await; - // Still-running process after LEAVE: the delta is refused, not stored. - b.update(1, 500.0).await; - let (t, msg) = tokio::time::timeout(Duration::from_secs(5), b.recv()).await.unwrap(); - assert_eq!(t, MSG_ERROR); - assert!(String::from_utf8_lossy(&msg).contains("not a member"), "{}", String::from_utf8_lossy(&msg)); - b.heartbeat().await; - let (t, msg) = tokio::time::timeout(Duration::from_secs(5), b.recv()).await.unwrap(); - assert_eq!(t, MSG_ERROR, "{}", String::from_utf8_lossy(&msg)); - - a.update(1, 1.0).await; - let (v, p) = a.next_base().await; - assert_eq!(v, 2); - assert_close(&p, 3.0); - a.update(2, 1.0).await; - let (v, p) = a.next_base().await; - assert_eq!(v, 3); - assert_close(&p, 4.0); - - let mut b2 = Island::connect(port, 2, 1).await; - assert!(matches!(b2.join().await, ElasticMsg::JoinAck { catch_up: true, base_version: 3, .. })); - let (v, p) = b2.next_base().await; - assert_eq!(v, 3); - assert_close(&p, 4.0); - b2.update(3, 1000.0).await; - tokio::time::sleep(Duration::from_millis(50)).await; - a.update(3, 1.0).await; - let (final_version, params) = - tokio::time::timeout(Duration::from_secs(10), server).await.unwrap().unwrap().unwrap(); - assert_eq!(final_version, 4); - assert_close(¶ms.unwrap(), 5.0); - - let lines = tape_lines(&tape); - let leave = lines.iter().find(|l| l.contains("pool_leave")).unwrap(); - assert!(leave.contains("\"reason\":\"finished\""), "{leave}"); - assert!(leave.contains("\"dropped_uncommitted\":{\"island_id\":2,\"outer_version\":1"), "{leave}"); - // Accepted from B: base 0, base 1 (then dropped), base 3 after rejoin; the - // post-LEAVE 500 never reaches the ledger. - let accepted_b = lines.iter().filter(|l| l.contains("\"kind\":\"delta_accepted\"") - && l.contains("\"island_id\":2")).count(); - assert_eq!(accepted_b, 3, "{lines:#?}"); - let steps: Vec<&String> = lines.iter().filter(|l| l.contains("\"kind\":\"outer_step\"")).collect(); - assert_eq!(steps.len(), 4, "{steps:#?}"); - for s in &steps[1..3] { - assert!(s.contains("\"weights\":[[1,"), "{s}"); - assert!(!s.contains("[2,"), "left island must not be merged: {s}"); - } - assert!(lines.iter().any(|l| l.contains("pool_join") && l.contains("\"island_id\":2") - && l.contains("\"catch_up\":true"))); - let _ = std::fs::remove_dir_all(&dir); - } - - /// D-S6: one tensor step with a checkpoint, then restart with --resume: - /// syncer_epoch increments, members/base/outer_version are restored, the - /// old epoch is fenced, and a further step continues from the base. - #[tokio::test] - async fn resume_restores_state_and_bumps_syncer_epoch() { - let dir = std::env::temp_dir().join(format!("elastic-ckpt-{}", std::process::id())); - std::fs::create_dir_all(&dir).unwrap(); - let ckpt = dir.join("elastic.ckpt"); - let _ = std::fs::remove_file(&ckpt); - let mk = |total_steps, resume| ElasticServerConfig { - params: ElasticParams { quorum_theta: 0.6, carry_gamma: 0.5, soft_deadline_s: 60, q_min: 1, max_carry_lag: 2 }, - key: KEY.to_vec(), - syncer_epoch: 0, - lease_s: 30.0, - total_steps, - event_tape: None, - tick: Duration::from_millis(20), - outer_lr: 1.0, - outer_momentum: 0.0, - outer_delayed_nesterov_period: None, - final_grace: Duration::ZERO, - checkpoint_path: Some(ckpt.clone()), - checkpoint_every: 1, - resume, - expected_backend_identity: None, - }; - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run(listener, mk(1, false))); - let (mut a, mut b) = (Island::connect(port, 1, 0).await, Island::connect(port, 2, 0).await); - a.join().await; - a.init(vec![0.0; 4]).await; - a.next_base().await; - b.join().await; - b.next_base().await; - a.update(0, 2.0).await; - b.update(0, 4.0).await; - let (_, p) = tokio::time::timeout(Duration::from_secs(10), server).await.unwrap().unwrap().unwrap(); - assert_close(&p.unwrap(), 3.0); - - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run(listener, mk(2, true))); - let mut stale = Island::connect(port, 1, 0).await; - stale.update(1, 1.0).await; - let (t, msg) = stale.recv().await; - assert_eq!(t, MSG_ERROR); - assert!(String::from_utf8_lossy(&msg).contains("fenced")); - // Restored members resend on the restored base with the new epoch. - let (mut a, mut b) = (Island::connect(port, 1, 1).await, Island::connect(port, 2, 1).await); - a.update(1, 1.0).await; - b.update(1, 3.0).await; - let (v, p) = tokio::time::timeout(Duration::from_secs(10), server).await.unwrap().unwrap().unwrap(); - assert_eq!(v, 2); - assert_close(&p.unwrap(), 5.0); - let saved = std::fs::read(&ckpt).unwrap(); - let mut sh = ElasticCoordinator::new(mk(0, false).params, 30.0, 0).unwrap(); - // magic 8 | identity flag 1 + 32 | coordinator length u64 | coordinator - assert_eq!((saved[8], &saved[9..41]), (1, &MILES_ID[..])); - let n = u64::from_le_bytes(saved[41..49].try_into().unwrap()) as usize; - sh = ElasticCoordinator::decode_state(sh.params, 30.0, &saved[49..49 + n], 0.0).unwrap(); - assert_eq!((sh.syncer_epoch, sh.outer_version), (1, 2)); - let _ = std::fs::remove_dir_all(&dir); - } - - /// P3 over the wire: a tensor run publishes sha256(base) per version; - /// SAMPLE_INDEX gets ACCEPT / ACCEPT_IS / REJECT replies, accepted - /// entries land in sample_index.jsonl, and nothing is forwarded. - #[tokio::test] - async fn sample_index_over_the_wire() { - use crate::elastic::{SampleIndexEntry, SAMPLE_INDEX_SCHEMA, VERDICT_ACCEPT, VERDICT_ACCEPT_IS, VERDICT_REJECT}; - let dir = std::env::temp_dir().join(format!("elastic-idx-{}", std::process::id())); - std::fs::create_dir_all(&dir).unwrap(); - let tape = dir.join("tape.jsonl"); - for f in ["tape.jsonl", "sample_index.jsonl", "status.json"] { - let _ = std::fs::remove_file(dir.join(f)); - } - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run( - listener, - ElasticServerConfig { - params: ElasticParams { quorum_theta: 0.6, carry_gamma: 0.5, soft_deadline_s: 60, q_min: 1, max_carry_lag: 2 }, - key: KEY.to_vec(), - syncer_epoch: 0, - lease_s: 30.0, - total_steps: 2, - event_tape: Some(tape.clone()), - tick: Duration::from_millis(20), - outer_lr: 1.0, - outer_momentum: 0.0, - outer_delayed_nesterov_period: None, - final_grace: Duration::ZERO, - checkpoint_path: None, - checkpoint_every: 0, - resume: false, - expected_backend_identity: None, - }, - )); - let mut a = Island::connect(port, 1, 0).await; - a.join().await; - a.init(vec![0.0; 4]).await; - let (_, base0) = a.next_base().await; - a.update(0, 1.0).await; - let (v, base1) = a.next_base().await; - assert_eq!(v, 1); - let entry = |ver: u64, base: &[f32], group: &str| SampleIndexEntry { - island_id: 1, - outer_version: ver, - inner_step: 0, - policy_hash: policy_hash(base), - group_id: group.into(), - prompt_id: "p".into(), - n: 8, - uri: "s3://b/g".into(), - size_bytes: 10, - sha256: "0".repeat(64), - has_behavior_logprob: true, - advantage_included: true, - created_at: 0.0, - schema: SAMPLE_INDEX_SCHEMA.into(), - }; - let mut verdicts = Vec::new(); - for e in [entry(1, &base1, "g1"), entry(0, &base0, "g0"), entry(0, &base1, "bad")] { - a.send(ElasticMsg::SampleIndex { syncer_epoch: 0, entry: e }).await; - let (t, p) = a.recv().await; - match ElasticMsg::decode(KEY, t, &p).unwrap() { - ElasticMsg::SampleVerdict { verdict, reason, .. } => verdicts.push((verdict, reason)), - m => panic!("{m:?}"), - } - } - assert_eq!( - verdicts, - vec![ - (VERDICT_ACCEPT, "same_version".to_string()), - (VERDICT_ACCEPT_IS, "stale_within_bound".to_string()), - (VERDICT_REJECT, "policy_hash_mismatch".to_string()), - ] - ); - a.update(1, 1.0).await; - tokio::time::timeout(Duration::from_secs(10), server).await.unwrap().unwrap().unwrap(); - let idx = std::fs::read_to_string(dir.join("sample_index.jsonl")).unwrap(); - assert_eq!(idx.lines().count(), 2, "{idx}"); - assert!(idx.contains("\"group_id\":\"g1\"") && idx.contains("\"verdict\":\"ACCEPT_IS\""), "{idx}"); - assert_eq!(tape_lines(&tape).iter().filter(|l| l.contains("sample_indexed")).count(), 3); - let status = std::fs::read_to_string(dir.join("status.json")).unwrap(); - assert!(status.contains("\"sample_index_count\":2"), "{status}"); - let _ = std::fs::remove_dir_all(&dir); - } - /// Final grace: after total_steps the server broadcasts FINISHED, answers - /// late DELTA_TENSOR / heartbeat / JOIN with FINISHED, and exits once - /// every member has left (before the grace window ends). - #[tokio::test] - async fn finished_grace_window_answers_late_islands_and_exits_on_leave() { - let dir = std::env::temp_dir().join(format!("elastic-fin-{}", std::process::id())); - std::fs::create_dir_all(&dir).unwrap(); - let tape = dir.join("tape.jsonl"); - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let started = Instant::now(); - let server = tokio::spawn(run( - listener, - ElasticServerConfig { - params: ElasticParams { quorum_theta: 0.5, carry_gamma: 0.5, soft_deadline_s: 60, q_min: 1, max_carry_lag: 2 }, - key: KEY.to_vec(), - syncer_epoch: 0, - lease_s: 30.0, - total_steps: 1, - event_tape: Some(tape), - tick: Duration::from_millis(20), - outer_lr: 1.0, - outer_momentum: 0.0, - outer_delayed_nesterov_period: None, - final_grace: Duration::from_secs(30), - checkpoint_path: None, - checkpoint_every: 0, - resume: false, - expected_backend_identity: None, - }, - )); - let (mut a, mut b) = (Island::connect(port, 1, 0).await, Island::connect(port, 2, 0).await); - a.join().await; - a.init(vec![0.0; 4]).await; - a.next_base().await; - b.join().await; - b.next_base().await; - a.update(0, 1.0).await; // 1/2 >= 0.5 -> step to v1 = total_steps - let expect_finished = |t: u8, p: &[u8]| match ElasticMsg::decode(KEY, t, p).unwrap() { - ElasticMsg::Finished { outer_version, policy_hash, .. } => { - assert_eq!(outer_version, 1); - assert_ne!(policy_hash, [0; 32]); - } - m => panic!("{m:?}"), - }; - // Broadcast reaches the slow island B unprompted (after BASE v1). - loop { - let (t, p) = b.recv().await; - if t == MSG_FINISHED { - expect_finished(t, &p); - break; - } - assert_eq!(t, MSG_ELASTIC_BASE); - } - // B's late delta, a heartbeat and a fresh JOIN all get FINISHED. - b.update(0, 5.0).await; - let (t, p) = b.recv().await; - expect_finished(t, &p); - b.heartbeat().await; - let (t, p) = b.recv().await; - expect_finished(t, &p); - let mut c = Island::connect(port, 3, 0).await; - c.send(ElasticMsg::Join { syncer_epoch: 0, island_id: 3, incarnation: 0, capacity: 1.0, backend_identity: MILES_ID, compat_group: String::new() }).await; - let (t, p) = c.recv().await; - expect_finished(t, &p); - let status = std::fs::read_to_string(dir.join("status.json")).unwrap(); - assert!(status.contains("\"state\":\"finished\""), "{status}"); - for i in [&mut a, &mut b] { - let id = i.id; - i.send(ElasticMsg::Leave { syncer_epoch: 0, island_id: id, reason: crate::elastic::LEAVE_REASON_FINISHED }).await; - } - let (v, p) = tokio::time::timeout(Duration::from_secs(10), server).await.unwrap().unwrap().unwrap(); - assert_eq!(v, 1); - assert_eq!(p.unwrap(), vec![1.0; 4], "late delta must not be merged"); - assert!(started.elapsed() < Duration::from_secs(20), "exited on all-left, not grace"); - let _ = std::fs::remove_dir_all(&dir); - } - - /// yeto-framework-decoupling 6.2a: the first JOIN pins the backend - /// identity; a JOIN with another identity (a "verl" island in a Miles - /// session) gets MSG_ERROR naming both hashes and is not admitted, a - /// second island with the same identity joins as before. The pin is in - /// the checkpoint (survives --resume) and a pre-identity YELSRV1 - /// checkpoint is refused. - #[tokio::test] - async fn backend_identity_pinned_by_first_join_refuses_other_backend() { - let dir = std::env::temp_dir().join(format!("elastic-ident-{}", std::process::id())); - std::fs::create_dir_all(&dir).unwrap(); - let ckpt = dir.join("elastic.ckpt"); - let _ = std::fs::remove_file(&ckpt); - let verl: [u8; 32] = [0x22; 32]; - let mk = |total_steps, resume| ElasticServerConfig { - params: ElasticParams { quorum_theta: 0.5, carry_gamma: 0.5, soft_deadline_s: 60, q_min: 1, max_carry_lag: 2 }, - key: KEY.to_vec(), - syncer_epoch: 0, - lease_s: 30.0, - total_steps, - event_tape: None, - tick: Duration::from_millis(20), - outer_lr: 1.0, - outer_momentum: 0.0, - outer_delayed_nesterov_period: None, - final_grace: Duration::ZERO, - checkpoint_path: Some(ckpt.clone()), - checkpoint_every: 1, - resume, - expected_backend_identity: None, - }; - let expect_refused = |t: u8, p: &[u8]| { - assert_eq!(t, MSG_ERROR); - let text = String::from_utf8_lossy(p).to_string(); - assert!(text.contains("backend identity mismatch, JOIN refused"), "{text}"); - assert!(text.contains(&"22".repeat(32)) && text.contains(&"11".repeat(32)), "{text}"); - assert!(text.contains("Miles vs verl"), "{text}"); - }; - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run(listener, mk(3, false))); - let mut a = Island::connect(port, 1, 0).await; - assert!(matches!(a.join().await, ElasticMsg::JoinAck { .. })); - a.init(vec![0.0; 4]).await; - a.next_base().await; - let mut v = Island::connect(port, 2, 0).await; - v.identity = verl; - v.send(ElasticMsg::Join { syncer_epoch: 0, island_id: 2, incarnation: 0, capacity: 1.0, backend_identity: verl, compat_group: "nvidia-h100".to_string() }).await; - let (t, p) = v.recv().await; - expect_refused(t, &p); - let mut b = Island::connect(port, 3, 0).await; - assert!(matches!(b.join().await, ElasticMsg::JoinAck { .. }), "same backend must join as before"); - b.next_base().await; - a.update(0, 1.0).await; // 1/2 >= 0.5 -> v1, checkpoint written before the broadcast - assert_eq!(a.next_base().await.0, 1); - server.abort(); - let bytes = std::fs::read(&ckpt).unwrap(); - assert_eq!(&bytes[..8], SERVER_STATE_MAGIC); - assert_eq!(bytes[8], 1); - assert_eq!(&bytes[9..41], &MILES_ID); - - // Resume keeps the pin: the verl island is still refused. - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run(listener, mk(3, true))); - let mut v = Island::connect(port, 2, 1).await; - v.send(ElasticMsg::Join { syncer_epoch: 1, island_id: 2, incarnation: 0, capacity: 1.0, backend_identity: verl, compat_group: "nvidia-h100".to_string() }).await; - let (t, p) = v.recv().await; - expect_refused(t, &p); - server.abort(); - - // A v1 (pre-identity) checkpoint is refused with a clear message. - let mut old = SERVER_STATE_MAGIC_V1.to_vec(); - old.extend_from_slice(&bytes[41..]); - std::fs::write(&ckpt, &old).unwrap(); - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let err = run(listener, mk(3, true)).await.unwrap_err(); - assert!(format!("{err:#}").contains("predates the backend identity field"), "{err:#}"); - let _ = std::fs::remove_dir_all(&dir); - } - - /// yeto-framework-decoupling 7.7c (user decision 2026-10-09): the first - /// JOIN pins compat_group; an H200 island in an H100 session and an Ascend - /// island are refused with both groups named, the syncer keeps running and - /// a second H100 island still joins. An H200 session admits H200. - #[tokio::test] - async fn compat_group_pinned_by_first_join_refuses_other_card_type() { - let mk = || ElasticServerConfig { - params: ElasticParams { quorum_theta: 0.5, carry_gamma: 0.5, soft_deadline_s: 60, q_min: 1, max_carry_lag: 2 }, - key: KEY.to_vec(), - syncer_epoch: 0, - lease_s: 30.0, - total_steps: 3, - event_tape: None, - tick: Duration::from_millis(20), - outer_lr: 1.0, - outer_momentum: 0.0, - outer_delayed_nesterov_period: None, - final_grace: Duration::ZERO, - checkpoint_path: None, - checkpoint_every: 1, - resume: false, - expected_backend_identity: None, - }; - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run(listener, mk())); - let mut a = Island::connect(port, 1, 0).await; - assert!(matches!(a.join().await, ElasticMsg::JoinAck { .. })); - a.init(vec![0.0; 4]).await; - a.next_base().await; - for (id, group) in [(2u32, "nvidia-h200"), (4, "ascend-910b")] { - let mut h = Island::connect(port, id, 0).await; - h.compat = group.to_string(); - h.identity = [0x33; 32]; // the group is inside the identity hash too - h.send(ElasticMsg::Join { - syncer_epoch: 0, island_id: id, incarnation: 0, capacity: 1.0, - backend_identity: h.identity, compat_group: h.compat.clone(), - }).await; - let (t, p) = h.recv().await; - assert_eq!(t, MSG_ERROR); - let text = String::from_utf8_lossy(&p).to_string(); - assert!(text.contains("compat_group mismatch, JOIN refused"), "{text}"); - assert!(text.contains(&format!("兼容组不同：nvidia-h100 对 {group}，容差未标定")), "{text}"); - } - let mut b = Island::connect(port, 3, 0).await; - assert!(matches!(b.join().await, ElasticMsg::JoinAck { .. }), "same card type must still join"); - server.abort(); - - // An H200 session admits a second H200 island. - let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(run(listener, mk())); - let mut a = Island::connect(port, 1, 0).await; - a.compat = "nvidia-h200".to_string(); - assert!(matches!(a.join().await, ElasticMsg::JoinAck { .. })); - a.init(vec![0.0; 4]).await; - a.next_base().await; - let mut b = Island::connect(port, 2, 0).await; - b.compat = "nvidia-h200".to_string(); - assert!(matches!(b.join().await, ElasticMsg::JoinAck { .. })); - server.abort(); - } -} diff --git a/syncer/src/hmac.rs b/syncer/src/hmac.rs new file mode 100644 index 000000000..ed4e748f1 --- /dev/null +++ b/syncer/src/hmac.rs @@ -0,0 +1,35 @@ +//! HMAC-SHA256 (RFC 2104) for the island HELLO MAC. + +use sha2::{Digest, Sha256}; + +pub const HMAC_LEN: usize = 32; + +pub fn hmac_sha256(key: &[u8], msg: &[u8]) -> [u8; 32] { + const B: usize = 64; + let mut k = [0u8; B]; + if key.len() > B { + k[..32].copy_from_slice(&Sha256::digest(key)); + } else { + k[..key.len()].copy_from_slice(key); + } + let mut inner = Sha256::new(); + inner.update(k.map(|b| b ^ 0x36)); + inner.update(msg); + let ih = inner.finalize(); + let mut outer = Sha256::new(); + outer.update(k.map(|b| b ^ 0x5c)); + outer.update(ih); + outer.finalize().into() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn matches_rfc4231_case_2() { + let mac = hmac_sha256(b"Jefe", b"what do ya want for nothing?"); + let hex: String = mac.iter().map(|b| format!("{b:02x}")).collect(); + assert_eq!(hex, "5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843"); + } +} diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 2c0a7989c..f66470d48 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -1,6 +1,5 @@ -mod elastic; -mod elastic_server; mod fold; +mod hmac; mod shard; #[cfg(test)] mod head_memory; @@ -145,27 +144,20 @@ struct Args { /// this is explicitly enabled. #[arg(long, default_value_t = false)] require_profile_binding: bool, - /// Inter-island scheduling mode: legacy (default, unchanged behavior) - /// or elastic (change rl-inter-island-scheduling). - #[arg(long, default_value = "legacy")] - island_scheduling_mode: String, /// Capacity fraction that must arrive before a merge (θ). Fragment /// rounds: unset keeps the count quorum K (--quorum); set replaces it. - /// elastic: defaults to 0.75. #[arg(long)] quorum_theta: Option, /// Weight discount γ^lag of a late delta. Fragment rounds: lag in /// fragment versions, unset = 1 (no discount). - /// elastic: carried-over deltas, defaults to 0.5. #[arg(long)] carry_gamma: Option, /// Soft deadline in seconds: merge what arrived even below quorum. - /// Fragment rounds: unset = off (also caps the grace window). - /// elastic: defaults to --quorum-timeout-s. + /// Unset = off (also caps the grace window). #[arg(long)] soft_deadline_s: Option, /// Minimum arrived islands for a capacity-quorum or soft-deadline merge. - #[arg(long, default_value_t = elastic::DEFAULT_Q_MIN)] + #[arg(long, default_value_t = 1)] q_min: u32, /// Fragment rounds: relative island capacity for --quorum-theta, as /// comma-separated id=capacity (unlisted islands have capacity 1). @@ -175,31 +167,26 @@ struct Args { /// first --learners). Unset = 0..--learners. #[arg(long)] max_learners: Option, - /// elastic: largest base lag still carried over; larger is rejected. - #[arg(long, default_value_t = elastic::DEFAULT_MAX_CARRY_LAG)] - max_carry_lag: u32, /// HMAC-SHA256 island key; falls back to the YETO_ISLAND_HMAC_KEY - /// environment variable. elastic: signs the elastic message types. - /// legacy: every HELLO must carry an HMAC under this key. + /// environment variable. Every HELLO must carry an HMAC under this key. #[arg(long)] island_hmac_key: Option, - /// legacy only: run without island authentication (no HMAC key). Only + /// Run without island authentication (no HMAC key). Only /// for local tests and benchmarks; also enabled by /// YETO_SYNCER_ALLOW_UNAUTHENTICATED=1. Never set by the launcher's RL path. #[arg(long, default_value_t = false)] allow_unauthenticated_islands: bool, /// Island contract (64 lowercase hex chars) pinned by the head's config: - /// legacy HELLO session contracts and elastic JOIN identities must match - /// it, so the first island to connect no longer decides the contract. + /// HELLO session contracts must match it, so the first island to + /// connect no longer decides the contract. #[arg(long)] expected_island_contract: Option, /// Membership lease in seconds (no heartbeat -> removed; the island - /// must re-JOIN). Fragment rounds: unset = no lease. elastic: defaults - /// to 150 s, as the launcher and `yeto --rl-island-lease-s`. + /// must re-JOIN). Unset = no lease. #[arg(long)] island_lease_s: Option, /// Fencing token of this syncer incarnation; bump on every restart. - /// Fragment rounds: a HELLO with a lower epoch is refused (a HELLO + /// A HELLO with a lower epoch is refused (a HELLO /// without one is admitted); persisted next to the checkpoint. #[arg(long, default_value_t = 0)] syncer_epoch: u64, @@ -213,8 +200,10 @@ struct Args { /// Fragment rounds: wait for one donor snapshot before asking the next. #[arg(long, default_value_t = 600)] recovery_donor_timeout_s: u64, - /// elastic: after total_steps keep answering FINISHED for this many - /// seconds (or until every member left); defaults to the soft deadline. + /// Dynamic membership: after total_steps wait this many seconds for the + /// islands still connected to acknowledge the final cut, then finish + /// without the missing ones. Unset = every island at the final cut must + /// acknowledge within --final-ack-timeout-s. #[arg(long)] final_grace_s: Option, /// Rounds the exact fold cannot reproduce (quorum below the roster or @@ -237,10 +226,8 @@ struct Args { shard_count: u32, } -const DEFAULT_ELASTIC_LEASE_S: f64 = 150.0; - /// Membership options of the fragment-round server; every unset flag keeps -/// the fixed-membership behavior. The elastic server reads its own defaults. +/// the fixed-membership behavior. fn parse_membership(args: &Args) -> anyhow::Result { let mut capacities = std::collections::BTreeMap::new(); for entry in &args.island_capacity { @@ -267,20 +254,10 @@ fn parse_membership(args: &Args) -> anyhow::Result carry_gamma: args.carry_gamma.unwrap_or(1.0), capacities, max_learners: args.max_learners, + final_grace: args.final_grace_s.map(std::time::Duration::from_secs), }) } -fn parse_island_scheduling(args: &Args) -> anyhow::Result { - elastic::IslandSchedulingMode::parse( - &args.island_scheduling_mode, - args.quorum_theta.unwrap_or(elastic::DEFAULT_QUORUM_THETA), - args.carry_gamma.unwrap_or(elastic::DEFAULT_CARRY_GAMMA), - args.soft_deadline_s.unwrap_or(args.quorum_timeout_s), - args.q_min, - args.max_carry_lag, - ) -} - impl Args { fn shard_plan(&self) -> anyhow::Result { shard::ShardPlan::new(self.shard_index, self.shard_count) @@ -316,13 +293,7 @@ fn main() -> anyhow::Result<()> { } }; let final_ack_timeout_s = args.resolved_final_ack_timeout_s(); - let island_scheduling = parse_island_scheduling(&args)?; - // The elastic server applies the same flags with its own defaults. - let membership = if args.island_scheduling_mode == elastic::MODE_ELASTIC { - membership::MembershipConfig::default() - } else { - parse_membership(&args)? - }; + let membership = parse_membership(&args)?; let recovery = recovery::RecoveryConfig { bound_steps: args.recovery_bound_steps, max_donor_tries: args.recovery_max_donor_tries, @@ -353,17 +324,12 @@ fn main() -> anyhow::Result<()> { .or_else(|| std::env::var("YETO_ISLAND_HMAC_KEY").ok()) .filter(|k| !k.is_empty()) .map(String::into_bytes); - if matches!(island_scheduling, elastic::IslandSchedulingMode::Elastic(_)) - && island_hmac_key.is_none() - { - anyhow::bail!("elastic mode requires --island-hmac-key or YETO_ISLAND_HMAC_KEY"); - } if island_hmac_key.is_none() { let allowed = args.allow_unauthenticated_islands || std::env::var("YETO_SYNCER_ALLOW_UNAUTHENTICATED").as_deref() == Ok("1"); if !allowed { anyhow::bail!( - "legacy mode requires an island HMAC key (--island-hmac-key or \ + "the syncer requires an island HMAC key (--island-hmac-key or \ YETO_ISLAND_HMAC_KEY); pass --allow-unauthenticated-islands only for local tests" ); } @@ -405,12 +371,9 @@ fn main() -> anyhow::Result<()> { max_base_lag: args.max_base_lag, learner_weight, require_profile_binding: args.require_profile_binding, - island_scheduling, island_hmac_key, expected_island_contract, - island_lease_s: args.island_lease_s.unwrap_or(DEFAULT_ELASTIC_LEASE_S), syncer_epoch: args.syncer_epoch, - elastic_final_grace_s: args.final_grace_s, membership, recovery, weighted_fold: args.weighted_fold.parse()?, @@ -490,43 +453,17 @@ mod tests { } #[test] - fn island_scheduling_defaults_to_legacy_and_soft_deadline_inherits() { + fn membership_is_static_by_default_and_the_retired_mode_flag_is_refused() { let a = Args::try_parse_from(["yeto-syncer", "--learners", "2", "--total-steps", "8"]) .unwrap(); - assert_eq!(a.island_scheduling_mode, "legacy"); - assert_eq!((a.quorum_theta, a.carry_gamma, a.q_min, a.max_carry_lag), (None, None, 1, 2)); + assert_eq!((a.quorum_theta, a.carry_gamma, a.q_min), (None, None, 1)); assert!(parse_membership(&a).unwrap().is_static(), "fragment rounds stay fixed by default"); - assert_eq!(a.soft_deadline_s, None); - let e = Args::try_parse_from([ - "yeto-syncer", "--learners", "2", "--total-steps", "8", - "--island-scheduling-mode", "elastic", "--quorum-theta", "0.6", - "--carry-gamma", "0.25", "--soft-deadline-s", "30", "--q-min", "2", - "--max-carry-lag", "3", - ]) - .unwrap(); - let mode = parse_island_scheduling(&e).unwrap(); - assert_eq!( - mode, - elastic::IslandSchedulingMode::Elastic(elastic::ElasticParams { - quorum_theta: 0.6, carry_gamma: 0.25, soft_deadline_s: 30, q_min: 2, max_carry_lag: 3, - }) - ); - } - - #[test] - fn elastic_flag_defaults_are_unchanged() { - let e = Args::try_parse_from([ - "yeto-syncer", "--learners", "2", "--total-steps", "8", - "--island-scheduling-mode", "elastic", - ]) - .unwrap(); - assert_eq!( - parse_island_scheduling(&e).unwrap(), - elastic::IslandSchedulingMode::Elastic(elastic::ElasticParams { - quorum_theta: 0.75, carry_gamma: 0.5, soft_deadline_s: 900, q_min: 1, max_carry_lag: 2, - }) - ); - assert_eq!(e.island_lease_s.unwrap_or(DEFAULT_ELASTIC_LEASE_S), 150.0); + assert_eq!((a.soft_deadline_s, a.final_grace_s), (None, None)); + for retired in [["--island-scheduling-mode", "elastic"], ["--max-carry-lag", "2"]] { + let mut argv = vec!["yeto-syncer", "--learners", "2", "--total-steps", "8"]; + argv.extend(retired); + assert!(Args::try_parse_from(argv).is_err(), "{retired:?} is no longer a flag"); + } } #[test] @@ -535,7 +472,7 @@ mod tests { "yeto-syncer", "--learners", "2", "--total-steps", "8", "--quorum-theta", "0.5", "--q-min", "1", "--island-capacity", "0=8,1=4", "--island-lease-s", "30", "--soft-deadline-s", "120", "--carry-gamma", "0.5", - "--max-learners", "4", + "--max-learners", "4", "--final-grace-s", "45", ]) .unwrap(); let m = parse_membership(&a).unwrap(); @@ -546,6 +483,7 @@ mod tests { assert_eq!(m.soft_deadline, Some(std::time::Duration::from_secs(120))); assert_eq!(m.carry_gamma, 0.5); assert_eq!(m.id_limit(2), 4); + assert_eq!(m.final_grace, Some(std::time::Duration::from_secs(45))); let dup = Args::try_parse_from([ "yeto-syncer", "--learners", "2", "--total-steps", "8", "--island-capacity", "0=1,0=2", ]) diff --git a/syncer/src/membership.rs b/syncer/src/membership.rs index e16ae3efd..282f6f86f 100644 --- a/syncer/src/membership.rs +++ b/syncer/src/membership.rs @@ -1,9 +1,8 @@ -//! Dynamic island membership for the sharded syncer (S27 E-1/E-2). +//! Dynamic island membership of the fragment-round server. //! -//! This ports the membership half of `elastic::ElasticCoordinator` (leases, -//! heartbeats, re-JOIN with catch-up, fencing, capacity quorum θ, soft -//! deadline, carry discount) into the fragment-round server in `server.rs`, -//! so one syncer has both the fragment-round data plane (fragments, K + +//! Leases, heartbeats, re-JOIN with catch-up, fencing, capacity quorum θ, +//! soft deadline, carry discount and the final grace window live next to +//! the fragment-round server in `server.rs`, so one syncer has both the fragment-round data plane (fragments, K + //! adaptive grace, tokens²/steps weights, RDA) and dynamic membership. //! //! The module is a pure state machine: it never touches sockets or tensors, @@ -83,6 +82,10 @@ pub struct MembershipConfig { /// Admit island ids in `0..max_learners` (scale-up past the initial /// roster). None: only `0..learners`. pub max_learners: Option, + /// After the last step, wait at most this long for the islands still + /// connected to acknowledge the final cut, then finish without the + /// missing ones. None: every island at the final cut must acknowledge. + pub final_grace: Option, } impl Default for MembershipConfig { @@ -94,6 +97,7 @@ impl Default for MembershipConfig { carry_gamma: 1.0, capacities: BTreeMap::new(), max_learners: None, + final_grace: None, } } } @@ -446,6 +450,18 @@ impl Membership { } } + /// Live islands as `(id, capacity, lease remaining)`, in id order; the + /// lease is None without one. + pub fn live_status(&self, now: Instant) -> Vec<(u32, f64, Option)> { + self.live + .iter() + .map(|(&id, m)| { + let remaining = self.cfg.lease.map(|lease| lease.saturating_sub(now.saturating_duration_since(m.last_seen))); + (id, self.cfg.capacity(id), remaining) + }) + .collect() + } + /// The connection of `member` closed. Its generation is retired; the /// island stays eligible for a plain reconnect (not a re-JOIN). pub fn depart(&mut self, member: Member) { diff --git a/syncer/src/merge.rs b/syncer/src/merge.rs index 17e21ccbd..7110e3b90 100644 --- a/syncer/src/merge.rs +++ b/syncer/src/merge.rs @@ -25,10 +25,9 @@ //! (`delayed_nesterov_step`) refreshes that momentum only every N steps. /// w_m = c_tokens² / c_steps ("quantity × quality"). The only definition of the -/// learner merge weight: strict pushes (server.rs) and elastic deltas -/// (elastic.rs) both call it; Python `island_ledger.merge_weight` is locked to -/// it by tests/test_rl_inter_island_rust_golden.py. `c_steps` is u64 so the -/// elastic wire value needs no narrowing (u32 strict values widen losslessly). +/// learner merge weight (server.rs); Python `island_ledger.merge_weight` is +/// locked to it by tests/test_rl_inter_island_rust_golden.py. `c_steps` is u64 +/// so every caller's value widens losslessly. pub fn learner_weight(c_tokens: u64, c_steps: u64) -> f64 { if c_steps == 0 { 0.0 diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index 88eb98ae8..0056a677f 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -23,28 +23,28 @@ pub const MSG_FINAL_ACK: u8 = 12; pub const MSG_FINAL_FRAGMENT: u8 = 13; /// Benchmark-only report that one learner has frozen at its exact budget. pub const MSG_BUDGET_DONE: u8 = 14; -// Elastic inter-island scheduling messages (change rl-inter-island-scheduling, -// D-S1/D-S2/D-S4). Only an elastic-mode server may accept them; the legacy -// learner loop keeps rejecting every type it does not know. Each payload -// carries syncer_epoch first and a trailing HMAC-SHA256 (see elastic.rs). -pub const MSG_JOIN: u8 = 15; -pub const MSG_JOIN_ACK: u8 = 16; -pub const MSG_LEAVE: u8 = 17; -pub const MSG_LEASE_HEARTBEAT: u8 = 18; -pub const MSG_SAMPLE_INDEX: u8 = 19; -/// Elastic metadata-only delta arrival (base version and merge-weight -/// inputs); tensor payloads are not carried by this message yet. -pub const MSG_DELTA_READY: u8 = 20; -/// Elastic: island seeds the initial flat f32 parameters (first one wins). -pub const MSG_ELASTIC_INIT: u8 = 21; -/// Elastic: island update θ − base (flat f32) plus merge-weight inputs. -pub const MSG_DELTA_TENSOR: u8 = 22; -/// Elastic: coordinator broadcasts the published base after each step. -pub const MSG_ELASTIC_BASE: u8 = 23; -/// Elastic: coordinator verdict for a SAMPLE_INDEX entry. -pub const MSG_SAMPLE_VERDICT: u8 = 24; -/// Elastic: run finished (final outer_version + base policy hash). -pub const MSG_FINISHED: u8 = 25; +// Message ids 15-25 belonged to a retired whole-vector island protocol; +// they stay unassigned so an old island gets a clear first-frame error. +/// An island leaves the pool: u32 learner_id | u8 reason (LEAVE_REASON_*). +/// Sent before the island closes its connection; the reason goes on the +/// `pool_leave` tape record. +pub const MSG_LEAVE: u8 = 27; +pub const LEAVE_REASON_NAMES: [&str; 7] = [ + "finished", + "lease_expired", + "stopped", + "preempted", + "infra_failed", + "training_failed", + "config_mismatch", +]; + +/// Tape name of a LEAVE reason byte. +pub fn leave_reason_name(reason: u8) -> String { + LEAVE_REASON_NAMES + .get(reason as usize) + .map_or_else(|| format!("unknown_{reason}"), |name| (*name).to_string()) +} /// S26 E3b: link-bandwidth probe. First and only frame of its own connection /// (never part of a learner session): the island uploads N bytes, the syncer /// acknowledges with the u64 byte count, then streams M bytes back. @@ -835,11 +835,6 @@ mod tests { assert_eq!(MSG_FINAL_ACK, 12); assert_eq!(MSG_FINAL_FRAGMENT, 13); assert_eq!(MSG_BUDGET_DONE, 14); - assert_eq!( - [MSG_JOIN, MSG_JOIN_ACK, MSG_LEAVE, MSG_LEASE_HEARTBEAT, MSG_SAMPLE_INDEX, MSG_DELTA_READY, - MSG_ELASTIC_INIT, MSG_DELTA_TENSOR, MSG_ELASTIC_BASE, MSG_SAMPLE_VERDICT, MSG_FINISHED], - [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25] - ); } #[tokio::test] diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 9149b127f..5e35a6427 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -130,28 +130,19 @@ pub struct Config { /// Fail closed unless HELLO carries and matches this Config's canonical /// semantic profile hash. Disabled by default for legacy/non-SAO clients. pub require_profile_binding: bool, - /// Inter-island scheduling mode. Legacy (default) adds nothing to the - /// semantic profile hash and changes no server behavior. - pub island_scheduling: crate::elastic::IslandSchedulingMode, - /// HMAC-SHA256 key (never hashed into the profile). elastic: signs the - /// elastic message types. legacy (secret-handling-hardening): every HELLO - /// must carry a trailing HMAC under this key. + /// HMAC-SHA256 key (never hashed into the profile): every HELLO must + /// carry a trailing HMAC under this key. pub island_hmac_key: Option>, /// secret-handling-hardening: island contract (backend identity bound to - /// the LR schedule, 32 bytes) pinned by the head's launch config. legacy: - /// a HELLO whose session_contract_hash is not + /// the LR schedule, 32 bytes) pinned by the head's launch config. + /// A HELLO whose session_contract_hash is not /// sha256(domain || layout_fingerprint || pinned) is refused, so the first - /// island to connect no longer decides the contract. elastic: pins the - /// JOIN backend identity before the first JOIN. None keeps first-wins. + /// island to connect no longer decides the contract. None keeps first-wins. pub expected_island_contract: Option<[u8; 32]>, - /// elastic: membership lease (seconds without heartbeat before removal). - pub island_lease_s: f64, - /// Fencing token of this syncer incarnation. elastic: every message. - /// Fragment rounds: a HELLO that carries a lower epoch is refused; the - /// epoch is persisted next to the checkpoint and never lowered on resume. + /// Fencing token of this syncer incarnation: a HELLO that carries a + /// lower epoch is refused; the epoch is persisted next to the checkpoint + /// and never lowered on resume. pub syncer_epoch: u64, - /// elastic: final grace window after total_steps; None = soft deadline. - pub elastic_final_grace_s: Option, /// Dynamic membership of the fragment-round server (leases, re-JOIN, /// fencing, capacity quorum, soft deadline, carry discount). The default /// keeps the fixed-membership behavior. @@ -213,7 +204,6 @@ impl Config { LearnerWeight::Equal => 1, }); encoded.push(u8::from(self.require_profile_binding)); - self.island_scheduling.encode_contract(&mut encoded); // Appended only when set, so plain-Nesterov hashes stay unchanged. if let Some(period) = self.outer_delayed_nesterov_period { encoded.extend_from_slice(DELAYED_NESTEROV_PROFILE_DOMAIN); @@ -497,6 +487,11 @@ enum Event { Disconnected { member: Member, }, + /// The island announced why it leaves; its disconnect follows. + Leave { + member: Member, + reason: u8, + }, /// A peer-recovery control message (`recovery.rs`). Recovery { member: Member, @@ -898,11 +893,6 @@ fn validate_config(cfg: &Config) -> Result<()> { if cfg.learners == 0 { bail!("--learners must be positive"); } - if cfg.expert_token_weights - && matches!(cfg.island_scheduling, crate::elastic::IslandSchedulingMode::Elastic(_)) - { - bail!("--expert-token-weights needs fragment rounds; the whole-vector elastic loop has no per-tensor weights"); - } if cfg.quorum == 0 { bail!("--quorum must be positive"); } @@ -998,8 +988,7 @@ fn validate_config(cfg: &Config) -> Result<()> { } /// A sharded head runs independent generic fragment rounds per shard. Modes -/// whose semantics span all fragments (strict policy sweeps, the elastic -/// whole-model loop) or all shards (learner-budget cutoff) are refused. +/// whose semantics span all fragments (strict policy sweeps) or all shards (learner-budget cutoff) are refused. fn validate_shard_config(cfg: &Config) -> Result<()> { if !cfg.shard.is_sharded() { return Ok(()); @@ -1007,9 +996,6 @@ fn validate_shard_config(cfg: &Config) -> Result<()> { if cfg.policy_sweep_fragments.is_some() { bail!("--shard-count > 1 does not support --policy-sweep-fragments"); } - if cfg.island_scheduling != crate::elastic::IslandSchedulingMode::Legacy { - bail!("--shard-count > 1 requires the fragment-round scheduler (--island-scheduling-mode legacy)"); - } if cfg.learner_budget_steps.is_some() { bail!("--shard-count > 1 does not support --learner-budget-steps"); } @@ -1155,52 +1141,7 @@ fn check_resume_checkpoint(cfg: &Config) -> Result<()> { } pub async fn run(cfg: Config) -> Result<()> { - run_with_listener(cfg, None).await -} - -/// `run` on a listener the caller already bound (tests bind port 0 so no -/// other test can take the port between choosing and binding it). Only the -/// fragment-round loop uses it; `None` binds `cfg.port`. -async fn run_with_listener(cfg: Config, bound: Option) -> Result<()> { validate_config(&cfg)?; - if let crate::elastic::IslandSchedulingMode::Elastic(params) = cfg.island_scheduling { - // Whole-vector service loop (elastic_server.rs): membership, - // capacity-weighted stepping, soft deadline, leases, carry, tensor - // merge with the outer optimizer, checkpoint and tape. The - // fragment-round loop below never reaches it. - let listener = TcpListener::bind(("0.0.0.0", cfg.port)) - .await - .with_context(|| format!("bind elastic syncer port {}", cfg.port))?; - let version = crate::elastic_server::run( - listener, - crate::elastic_server::ElasticServerConfig { - params, - expected_backend_identity: cfg.expected_island_contract, - key: cfg - .island_hmac_key - .clone() - .context("elastic mode requires an island HMAC key")?, - syncer_epoch: cfg.syncer_epoch, - lease_s: cfg.island_lease_s, - total_steps: cfg.total_steps, - event_tape: cfg.event_tape.clone(), - tick: std::time::Duration::from_millis(50), - outer_lr: cfg.outer_lr, - outer_momentum: cfg.outer_momentum, - outer_delayed_nesterov_period: cfg.outer_delayed_nesterov_period, - final_grace: std::time::Duration::from_secs( - cfg.elastic_final_grace_s.unwrap_or(params.soft_deadline_s), - ), - checkpoint_path: cfg.checkpoint_path.clone(), - checkpoint_every: cfg.checkpoint_every, - resume: cfg.resume, - }, - ) - .await? - .0; - info!(outer_version = version, "elastic syncer finished"); - return Ok(()); - } check_resume_checkpoint(&cfg)?; let semantic_profile_hash = cfg.semantic_profile_hash(); if let Some(steps) = cfg.learner_budget_steps { @@ -1214,12 +1155,9 @@ async fn run_with_listener(cfg: Config, bound: Option) -> Result<() // before accepting any learner so every early failure stays recovery-only. remove_final_marker(cfg.checkpoint_path.as_ref().unwrap())?; } - let listener = match bound { - Some(listener) => listener, - None => TcpListener::bind(("0.0.0.0", cfg.port)) - .await - .with_context(|| format!("bind port {}", cfg.port))?, - }; + let listener = TcpListener::bind(("0.0.0.0", cfg.port)) + .await + .with_context(|| format!("bind port {}", cfg.port))?; let semantic_profile_sha256: String = semantic_profile_hash .iter() .map(|byte| format!("{byte:02x}")) @@ -1254,7 +1192,7 @@ async fn run_with_listener(cfg: Config, bound: Option) -> Result<() expected_contract: cfg.expected_island_contract, }); if hello_auth.key.is_none() { - warn!("legacy syncer runs WITHOUT island authentication (HELLO is not HMAC-checked)"); + warn!("syncer runs WITHOUT island authentication (HELLO is not HMAC-checked)"); } tokio::spawn(async move { loop { @@ -1331,7 +1269,7 @@ pub struct HelloAuth { pub expected_contract: Option<[u8; 32]>, } -/// Domain of the legacy HELLO MAC: HMAC-SHA256(key, domain || hello_body). +/// Domain of the HELLO MAC: HMAC-SHA256(key, domain || hello_body). pub const HELLO_MAC_DOMAIN: &[u8] = b"yeto-hello-mac-v1\0"; /// Domain of the RL session contract (yeto.rl.engine.backend_identity). pub const SESSION_CONTRACT_DOMAIN: &[u8] = b"yeto-rl-session-contract-v2\0"; @@ -1341,14 +1279,13 @@ pub const SESSION_CONTRACT_DOMAIN: &[u8] = b"yeto-rl-session-contract-v2\0"; fn open_hello<'a>(key: Option<&[u8]>, payload: &'a [u8]) -> Result<&'a [u8]> { let Some(key) = key else { return Ok(payload) }; let mut framed = Vec::with_capacity(payload.len() + 1); - // elastic::open MACs msg_type || body; HELLO uses its own domain instead. - if payload.len() < crate::elastic::HMAC_LEN { + if payload.len() < crate::hmac::HMAC_LEN { bail!("HELLO shorter than its HMAC (island has no YETO_ISLAND_HMAC_KEY?)"); } - let (body, mac) = payload.split_at(payload.len() - crate::elastic::HMAC_LEN); + let (body, mac) = payload.split_at(payload.len() - crate::hmac::HMAC_LEN); framed.extend_from_slice(HELLO_MAC_DOMAIN); framed.extend_from_slice(body); - let expected = crate::elastic::hmac_sha256(key, &framed); + let expected = crate::hmac::hmac_sha256(key, &framed); let diff = expected.iter().zip(mac).fold(0u8, |acc, (x, y)| acc | (x ^ y)); if diff != 0 { bail!("HELLO HMAC mismatch (wrong or missing island HMAC key)"); @@ -1636,7 +1573,7 @@ async fn handle_connection( syncer_profile_hash, compat_group, }; - // Strict and elastic alike: a HELLO whose contract differs from the + // A HELLO whose contract differs from the // established session is refused on this connection only (MSG_ERROR // naming both hashes). The syncer and the islands already admitted // keep running; one mis-wired island (e.g. a Miles island in a verl @@ -1929,6 +1866,7 @@ async fn read_loop( MSG_INIT_PARAMS => Ok(group.max_init_payload), MSG_PUSH_FRAGMENT => Ok(group.max_push_payload), MSG_HEARTBEAT => Ok(12), + MSG_LEAVE => Ok(5), MSG_FINAL_ACK => Ok(10), MSG_BUDGET_DONE => Ok(8), MSG_CHUNK => Ok(CHUNK_HEADER_SIZE + CHUNK_SIZE as u64), @@ -2210,6 +2148,27 @@ async fn dispatch_inner( .await .ok(); } + MSG_LEAVE => { + let mut r = Reader(payload); + let learner_id = r.u32()?; + let reason = r.u8()?; + if learner_id != group.member.learner_id { + bail!( + "LEAVE learner id {learner_id} does not match connected group {}", + group.member.learner_id + ); + } + if r.remaining() != 0 { + bail!("trailing bytes in LEAVE"); + } + event_tx + .send(Event::Leave { + member: group.member, + reason, + }) + .await + .ok(); + } MSG_FINAL_ACK => { let global_step = decode_final_ack(payload)?; event_tx @@ -2408,6 +2367,11 @@ async fn scheduler( let mut budget_reports: HashSet = HashSet::new(); let strict_exact = cfg.max_base_lag == Some(0); let fixed_roster = strict_exact && cfg.quorum == cfg.learners && cfg.grace_ms == 0; + // Dynamic membership writes island-pool records next to the merge records. + let dynamic = !cfg.membership.is_static(); + let mut lease_expired: HashSet = HashSet::new(); + // LEAVE reasons of islands whose disconnect has not arrived yet. + let mut leave_reasons: HashMap = HashMap::new(); let checkpoint_each_round = checkpoint_before_broadcast(&cfg); // Phase 1: wait until every fragment is initialized (via INIT_PARAMS or @@ -2442,6 +2406,17 @@ async fn scheduler( } Event::Hello { group } => { admit_recovering_outside_training(®istry, &group, "initialization").await; + if dynamic { + append_pool_record( + cfg.event_tape.as_deref(), + "pool_join", + cfg.syncer_epoch, + &format!( + "\"island_id\":{},\"generation\":{},\"catch_up\":false,\"base_version\":0", + group.member.learner_id, group.member.generation + ), + ); + } if state.is_none() { // Layout comes from the HELLO of the first learner. // (All learners must build identical layouts.) @@ -2512,7 +2487,7 @@ async fn scheduler( } } Event::Push { .. } => warn!("push before initialization; dropped"), - Event::Heartbeat { .. } => {} + Event::Heartbeat { .. } | Event::Leave { .. } => {} Event::FinalAck { member, .. } => warn!( learner_id = member.learner_id, generation = member.generation, @@ -2619,6 +2594,9 @@ async fn scheduler( let mut cutoff_established = false; let mut first_budget_report: Option<(Member, u64)> = None; let mut recovery = PeerRecovery::new(cfg.recovery.clone()); + if dynamic { + write_pool_status(&cfg, ®istry, st.global_step, "running"); + } 'outer: while next_t <= last_step || !inflight.is_empty() || !computing.is_empty() @@ -2630,6 +2608,18 @@ async fn scheduler( break 'outer; } let expired = expire_member_leases(®istry, &mut inflight, &mut step_rates).await; + for member in &expired { + lease_expired.insert(*member); + append_pool_record( + cfg.event_tape.as_deref(), + "pool_leave", + cfg.syncer_epoch, + &format!( + "\"island_id\":{},\"generation\":{},\"reason\":\"lease_expired\"", + member.learner_id, member.generation + ), + ); + } if !recovery.is_empty() { let now = Instant::now(); let candidates = recovery_candidates(®istry, &inflight); @@ -2700,6 +2690,9 @@ async fn scheduler( } } next_commit_t = schedule.step_after(next_commit_t); + if dynamic { + write_pool_status(&cfg, ®istry, st.global_step, "running"); + } } // Keep the pipeline full (throttled by min_round_interval_ms). @@ -3018,6 +3011,14 @@ async fn scheduler( push.base_version, ) }; + let pool = |kind: &str, fields: String| { + if dynamic { + append_pool_record(cfg.event_tape.as_deref(), kind, cfg.syncer_epoch, &fields); + } + }; + let rejected = |reason: &str| { + pool("delta_rejected", format!("\"island_id\":{learner_id},\"reason\":\"{reason}\"")); + }; if let Err(reason) = fenced { warn!( learner_id, @@ -3026,13 +3027,31 @@ async fn scheduler( fragment = fragment_id, "push fenced: {reason:#}" ); + rejected("fenced"); continue; } + let base_version = push.base_version; + let round_base = inflight + .iter() + .find(|round| round.t == global_step && round.p == fragment_id as usize) + .map(|round| round.base_version); let disposition = route_push(&mut inflight, member, push, cfg.max_base_lag, fixed_roster); match disposition { PushDisposition::Accepted => { step_rates.note(member, local_step, Instant::now()); + let lag = round_base.map_or(0, |base| base.saturating_sub(base_version)); + if lag == 0 { + pool("delta_accepted", format!("\"island_id\":{learner_id}")); + } else { + pool( + "delta_carried_over", + format!( + "\"island_id\":{learner_id},\"base\":{base_version},\"lag\":{lag},\"discount\":{:?}", + cfg.membership.carry_discount(lag) + ), + ); + } } PushDisposition::Duplicate => warn!( learner_id, @@ -3041,7 +3060,9 @@ async fn scheduler( fragment = fragment_id, "duplicate push rejected" ), - PushDisposition::StaleBase if strict_exact => { + // Fixed roster: a stale base is a strict failure. In an island + // pool a late island is only refused (delta_rejected record). + PushDisposition::StaleBase if strict_exact && !dynamic => { let metric = "rejected_stale_updates"; let message = format!( "learner {learner_id} generation {generation} step {global_step} fragment {fragment_id}: {disposition:?}" @@ -3049,16 +3070,33 @@ async fn scheduler( append_strict_failure(cfg.event_tape.as_deref(), metric, &message); bail!("RL strict failure {metric}: {message}"); } - disposition => warn!( - learner_id, - generation, - step = global_step, - fragment = fragment_id, - reason = ?disposition, - "push rejected" - ), + disposition => { + warn!( + learner_id, + generation, + step = global_step, + fragment = fragment_id, + reason = ?disposition, + "push rejected" + ); + rejected(match disposition { + PushDisposition::StaleBase => "carry_lag_exceeded", + PushDisposition::OutOfRound => "out_of_round", + PushDisposition::FutureBase => "future_base", + _ => "unexpected_member", + }); + } } } + Event::Leave { member, reason } => { + info!( + learner_id = member.learner_id, + generation = member.generation, + reason = %leave_reason_name(reason), + "island leaves the pool" + ); + leave_reasons.insert(member, reason); + } Event::Heartbeat { member, local_step } => { registry.lock().unwrap().membership.renew(member, Instant::now()); if is_current_member(®istry, member) { @@ -3131,6 +3169,18 @@ async fn scheduler( .unwrap() .membership .record_catch_up(group.member, &st.versions); + if dynamic { + append_pool_record( + cfg.event_tape.as_deref(), + "pool_join", + cfg.syncer_epoch, + &format!( + "\"island_id\":{},\"generation\":{},\"catch_up\":true,\"base_version\":{}", + group.member.learner_id, group.member.generation, st.global_step + ), + ); + write_pool_status(&cfg, ®istry, st.global_step, "running"); + } if !repull_rounds_for(&group, ®istry, &mut inflight, fixed_roster, &budget_cutoff).await { cutoff_established = true; break 'outer; @@ -3252,6 +3302,27 @@ async fn scheduler( generation = member.generation, "learner connection generation disconnected" ); + if dynamic && !lease_expired.contains(&member) { + // Only a departure without a newer generation leaves the + // pool; an expired lease already wrote its pool_leave. + let replaced = registry.lock().unwrap().current.contains_key(&member.learner_id); + if !replaced { + append_pool_record( + cfg.event_tape.as_deref(), + "pool_leave", + cfg.syncer_epoch, + &format!( + "\"island_id\":{},\"generation\":{},\"reason\":\"{}\"", + member.learner_id, + member.generation, + leave_reasons + .get(&member) + .map_or_else(|| "disconnected".to_string(), |r| leave_reason_name(*r)) + ), + ); + } + } + leave_reasons.remove(&member); step_rates.remove(member); let candidates = recovery_candidates(®istry, &inflight); let out = recovery.on_gone(member, &candidates, Instant::now()); @@ -3436,7 +3507,13 @@ async fn scheduler( .map(|group| group.member.learner_id) .collect() }; + if dynamic { + write_pool_status(&cfg, ®istry, st.global_step, "finishing"); + } finalize_learners(&cfg, &st, &mut events, ®istry, &final_members).await?; + if dynamic { + write_pool_status(&cfg, ®istry, st.global_step, "finished"); + } info!("training complete after {} outer steps", cfg.total_steps); // Give writer tasks a moment to flush the final control frames. tokio::time::sleep(Duration::from_secs(2)).await; @@ -3636,7 +3713,7 @@ async fn collect_budget_reports( generation = member.generation, "disconnected while waiting for learner-budget cutoff" ), - Event::Init { .. } | Event::FinalAck { .. } => {} + Event::Init { .. } | Event::FinalAck { .. } | Event::Leave { .. } => {} Event::Recovery { member, .. } => ignore_recovery_message(member, "the budget cutoff"), } } @@ -4062,6 +4139,9 @@ fn prepare_round_compute( tensor_scales.push(scales); } } + if !cfg.membership.is_static() { + append_outer_step_record(cfg, round, &ordered_pushes, &weights); + } let Some(fold) = round.fold.take() else { let tensor_scales = cfg.expert_token_weights.then_some(tensor_scales.as_slice()); return st.prepare_merge_scaled( @@ -4789,6 +4869,12 @@ async fn finalize_learners( if expected.is_empty() { bail!("cannot finalize: no live learner groups at the final cut"); } + // Dynamic membership: the final cut goes to every island that is (or + // becomes) connected, the wait is bounded by the final grace window and + // an island that left or never answered does not fail the run. + let final_grace = cfg.membership.final_grace; + let mut expected = expected.clone(); + let expected = &mut expected; let mut eligible: HashMap> = HashMap::new(); for group in current_groups(registry) { if expected.contains(&group.member.learner_id) { @@ -4811,28 +4897,45 @@ async fn finalize_learners( } } - let deadline = Instant::now() + Duration::from_secs(cfg.final_ack_timeout_s); + let wait = final_grace.unwrap_or(Duration::from_secs(cfg.final_ack_timeout_s)); + let deadline = Instant::now() + wait; let mut acknowledged: HashMap = HashMap::new(); // With dynamic membership, an island whose lease runs out during // finalization is absent: the final cut no longer waits for its ACK. let mut absence = FinalAbsence::new(cfg, expected, registry); - while !absence.pending(&acknowledged).is_empty() { + while !absence.pending(expected, &acknowledged).is_empty() { let now = Instant::now(); let wake = absence.next_deadline(registry).map_or(deadline, |at| at.min(deadline)); let event = match tokio::time::timeout(wake.saturating_duration_since(now), events.recv()).await { Ok(event) => event.context("event channel closed during finalization")?, Err(_) if Instant::now() < deadline => { - for learner_id in absence.drop_expired(registry, Instant::now()) { + for learner_id in absence.drop_expired(registry, expected, Instant::now()) { warn!(learner_id, "lease expired during finalization; final ACK no longer required"); } - if absence.pending(&acknowledged).is_empty() && acknowledged.is_empty() { + if absence.pending(expected, &acknowledged).is_empty() && acknowledged.is_empty() { bail!("cannot finalize: every learner's lease expired before it acknowledged the final cut"); } continue; } Err(_) => { - let mut missing: Vec = absence.pending(&acknowledged).into_iter().collect(); + let mut missing: Vec = absence.pending(expected, &acknowledged).into_iter().collect(); missing.sort_unstable(); + if final_grace.is_some() { + warn!(?missing, "final grace window ended; finishing without these islands"); + let acked: Vec = acknowledged.keys().copied().collect(); + append_pool_record( + cfg.event_tape.as_deref(), + "final_grace_expired", + cfg.syncer_epoch, + &format!( + "\"final_step\":{},\"acknowledged\":{},\"missing\":{}", + st.global_step, + json_ids(&acked), + json_ids(&missing) + ), + ); + break; + } bail!( "finalization timed out after {}s waiting for learner acknowledgements: {:?}", cfg.final_ack_timeout_s, @@ -4852,7 +4955,7 @@ async fn finalize_learners( } => { let learner_id = member.learner_id; match record_final_ack( - expected, + &*expected, &eligible, &mut acknowledged, member, @@ -4889,6 +4992,11 @@ async fn finalize_learners( Event::Hello { group } => { admit_recovering_outside_training(registry, &group, "finalization").await; let member = group.member; + if final_grace.is_some() && is_current_member(registry, member) { + // A late (re-)JOIN still gets the final cut and may + // acknowledge it within the grace window. + expected.insert(member.learner_id); + } if expected.contains(&member.learner_id) && !acknowledged.contains_key(&member.learner_id) && is_current_member(registry, member) @@ -4918,6 +5026,20 @@ async fn finalize_learners( } } Event::Disconnected { member } => { + if final_grace.is_some() + && !acknowledged.contains_key(&member.learner_id) + && !current_groups(registry) + .iter() + .any(|group| group.member.learner_id == member.learner_id) + { + // The island left: stop waiting for it unless it re-JOINs. + expected.remove(&member.learner_id); + info!( + learner_id = member.learner_id, + "island left during the final grace window" + ); + continue; + } warn!( learner_id = member.learner_id, generation = member.generation, @@ -4926,6 +5048,7 @@ async fn finalize_learners( } Event::Push { .. } | Event::Heartbeat { .. } + | Event::Leave { .. } | Event::Init { .. } | Event::Recovery { .. } | Event::BudgetDone { .. } => { @@ -4959,7 +5082,6 @@ async fn finalize_learners( /// from the wait when its lease expires, or when it stays disconnected for one /// lease without a new generation saying HELLO. struct FinalAbsence { - expected: HashSet, lease: Option, disconnected_since: HashMap, dropped: HashSet, @@ -4976,11 +5098,12 @@ impl FinalAbsence { disconnected_since.insert(*id, now); } } - Self { expected: expected.clone(), lease, disconnected_since, dropped: HashSet::new() } + Self { lease, disconnected_since, dropped: HashSet::new() } } - fn pending(&self, acknowledged: &HashMap) -> HashSet { - self.expected + /// `expected` is the finalization's live set (a late JOIN adds to it). + fn pending(&self, expected: &HashSet, acknowledged: &HashMap) -> HashSet { + expected .iter() .copied() .filter(|id| !acknowledged.contains_key(id) && !self.dropped.contains(id)) @@ -5017,7 +5140,7 @@ impl FinalAbsence { } /// Drop every pending learner whose lease expired; returns them in id order. - fn drop_expired(&mut self, registry: &Registry, now: Instant) -> Vec { + fn drop_expired(&mut self, registry: &Registry, expected: &HashSet, now: Instant) -> Vec { let Some(lease) = self.lease else { return Vec::new() }; let mut gone: Vec = registry .lock() @@ -5032,7 +5155,7 @@ impl FinalAbsence { .filter(|(_, since)| now.saturating_duration_since(**since) >= lease) .map(|(id, _)| *id), ); - gone.retain(|id| self.expected.contains(id) && self.dropped.insert(*id)); + gone.retain(|id| expected.contains(id) && self.dropped.insert(*id)); for id in &gone { self.disconnected_since.remove(id); } @@ -5110,6 +5233,119 @@ async fn send_fragment( } /// One JSONL record per merge: the event tape. +/// One island-pool record on the event tape: `{"kind":..,"syncer_epoch":.., +/// ,"time_unix":..}`. Written only with dynamic membership; the +/// kinds (`pool_join`, `pool_leave`, `delta_carried_over`, `delta_rejected`, +/// `outer_step`, `final_grace_expired`) are what the merge-metrics folder +/// and the launcher's pool reader consume. +fn append_pool_record(path: Option<&std::path::Path>, kind: &str, syncer_epoch: u64, fields: &str) { + use std::io::Write; + + let Some(path) = path else { + return; + }; + let sep = if fields.is_empty() { "" } else { "," }; + let line = format!( + "{{\"kind\":\"{kind}\",\"syncer_epoch\":{syncer_epoch}{sep}{fields},\"time_unix\":{:?}}}\n", + unix_now() + ); + let result = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path) + .and_then(|mut file| file.write_all(line.as_bytes())); + if let Err(error) = result { + warn!("event tape write failed: {error}"); + } +} + +/// The `outer_step` pool record of one merged round: who arrived, each +/// push's normalized merge weight (`[island, base_version, weight]`), and +/// which expected islands are absent. `timed_out` marks a soft-deadline +/// merge below quorum. +fn append_outer_step_record( + cfg: &Config, + round: &Round, + ordered_pushes: &[(&Member, &Push)], + weights: &[f64], +) { + let capacity = |members: &mut dyn Iterator| -> f64 { + members.map(|member| cfg.membership.capacity(member.learner_id)).sum() + }; + let cap_arrived = capacity(&mut round.pushes.keys()); + let cap_total = capacity(&mut round.expected_members.iter()); + let total: f64 = weights.iter().sum(); + let entry = |(member, push): &(&Member, &Push), weight: f64| { + format!("[{},{},{:?}]", member.learner_id, push.base_version, weight) + }; + let raw: Vec = ordered_pushes.iter().zip(weights).map(|(p, w)| entry(p, *w)).collect(); + let normalized: Vec = ordered_pushes + .iter() + .zip(weights) + .map(|(p, w)| entry(p, if total > 0.0 { w / total } else { 0.0 })) + .collect(); + let absent: Vec = round + .expected_members + .iter() + .filter(|member| !round.pushes.contains_key(member)) + .map(|member| member.learner_id) + .collect(); + append_pool_record( + cfg.event_tape.as_deref(), + "outer_step", + cfg.syncer_epoch, + &format!( + "\"step\":{},\"fragment\":{},\"base_version\":{},\"arrived\":{},\"cap_arrived\":{:?},\"cap_total\":{:?},\"timed_out\":{},\"raw_weights\":[{}],\"weights\":[{}],\"absent\":{},\"step_seconds\":{:?}", + round.t, + round.p, + round.base_version, + round.pushes.len(), + cap_arrived, + cap_total, + !round.quorum_reached(), + raw.join(","), + normalized.join(","), + json_ids(&absent), + round.started.elapsed().as_secs_f64(), + ), + ); +} + +/// Schema of `status.json`, written next to the event tape with dynamic +/// membership (read by `yeto.rl.engine.island_status`). +const POOL_STATUS_SCHEMA: &str = "yeto.syncer.elastic-status/v1"; + +/// Atomically rewrite `status.json` beside the event tape: run state, syncer +/// epoch, committed step and the live islands. Best effort (a write failure +/// is logged, never fatal). +fn write_pool_status(cfg: &Config, registry: &Registry, outer_version: u64, state: &str) { + let Some(tape) = cfg.event_tape.as_deref() else { + return; + }; + let (syncer_epoch, membership_epoch, live) = { + let registry = registry.lock().unwrap(); + let membership = ®istry.membership; + (membership.syncer_epoch(), membership.epoch, membership.live_status(Instant::now())) + }; + let cap_total: f64 = live.iter().map(|(_, capacity, _)| capacity).sum(); + let islands: Vec = live + .iter() + .map(|(id, capacity, lease)| { + let lease = lease.map_or_else(|| "null".to_string(), |d| format!("{:?}", d.as_secs_f64())); + format!("\"{id}\":{{\"capacity\":{capacity:?},\"lease_remaining_s\":{lease}}}") + }) + .collect(); + let body = format!( + "{{\"schema\":\"{POOL_STATUS_SCHEMA}\",\"state\":\"{state}\",\"syncer_epoch\":{syncer_epoch},\"membership_epoch\":{membership_epoch},\"outer_version\":{outer_version},\"cap_total\":{cap_total:?},\"islands\":{{{}}}}}", + islands.join(",") + ); + let path = tape.with_file_name("status.json"); + let temporary = tape.with_file_name(".status.json.tmp"); + if let Err(error) = std::fs::write(&temporary, body).and_then(|()| std::fs::rename(&temporary, &path)) { + warn!(path = %path.display(), "status write failed: {error}"); + } +} + fn append_strict_failure(path: Option<&std::path::Path>, metric: &str, message: &str) { use std::io::Write; @@ -5398,7 +5634,7 @@ mod hello_auth_tests { let mut input = HELLO_MAC_DOMAIN.to_vec(); input.extend_from_slice(body); let mut out = body.to_vec(); - out.extend_from_slice(&crate::elastic::hmac_sha256(key, &input)); + out.extend_from_slice(&crate::hmac::hmac_sha256(key, &input)); out } @@ -5777,12 +6013,9 @@ mod tests { max_base_lag: Some(0), learner_weight: LearnerWeight::Equal, require_profile_binding: false, - island_scheduling: Default::default(), island_hmac_key: None, expected_island_contract: None, - island_lease_s: 30.0, syncer_epoch: 0, - elastic_final_grace_s: None, membership: Default::default(), recovery: Default::default(), weighted_fold: Default::default(), @@ -5837,46 +6070,6 @@ mod tests { assert_eq!(digest, "f51de8a095c0fa428c5483c7921da805bb8b3cd2736fe19fdd03f47e4f733913"); } - #[test] - fn elastic_profile_hash_matches_the_python_canonical_vector() { - // Same inputs as semantic_profile_hash_matches_the_python_canonical_vector - // plus elastic defaults; Python: tests/test_rl_inter_island_contract.py. - let mut config = round_test_config(8); - config.learners = 2; - config.quorum = 2; - config.grace_gamma = 0.8; - config.grace_tau = 2.0; - config.pipeline = 2; - config.sync_interval_steps = 24.0; - config.quorum_timeout_s = 900; - config.final_ack_timeout_s = 900; - config.outer_lr = 0.7; - config.outer_momentum = 0.9; - config.checkpoint_path = Some(std::path::PathBuf::from("/ignored/actor.ckpt")); - config.checkpoint_every = 1; - config.resume = true; - config.mark_final_checkpoint = true; - config.require_profile_binding = true; - config.island_scheduling = crate::elastic::IslandSchedulingMode::Elastic( - crate::elastic::ElasticParams { - quorum_theta: 0.75, - carry_gamma: 0.5, - soft_deadline_s: 900, - q_min: 1, - max_carry_lag: 2, - }, - ); - let digest = config - .semantic_profile_hash() - .iter() - .map(|byte| format!("{byte:02x}")) - .collect::(); - assert_eq!( - digest, - "c677c127b1ad3e1da179863ee02b6b2e1bc232edb03b088bad62251e5d8d9ba5" - ); - } - #[test] fn delayed_nesterov_profile_hash_matches_the_python_canonical_vector() { // Python: tests/test_rl_inter_island_contract.py (DELAYED_NESTEROV_4). @@ -5908,53 +6101,8 @@ mod tests { } #[test] - fn island_scheduling_mode_is_bound_into_the_profile_hash() { - use crate::elastic::{ElasticParams, IslandSchedulingMode}; - let mut config = round_test_config(8); - let legacy = config.semantic_profile_hash(); - // An explicit legacy mode and an HMAC key leave the hash untouched. - config.island_scheduling = IslandSchedulingMode::Legacy; - config.island_hmac_key = Some(b"secret".to_vec()); - assert_eq!(config.semantic_profile_hash(), legacy); - let base = ElasticParams { - quorum_theta: 0.75, - carry_gamma: 0.5, - soft_deadline_s: 900, - q_min: 1, - max_carry_lag: 2, - }; - config.island_scheduling = IslandSchedulingMode::Elastic(base); - let elastic = config.semantic_profile_hash(); - assert_ne!(elastic, legacy, "legacy HELLO must not match an elastic server"); - let variants = [ - ElasticParams { quorum_theta: 0.5, ..base }, - ElasticParams { carry_gamma: 0.25, ..base }, - ElasticParams { soft_deadline_s: 60, ..base }, - ElasticParams { q_min: 2, ..base }, - ElasticParams { max_carry_lag: 3, ..base }, - ]; - for v in variants { - config.island_scheduling = IslandSchedulingMode::Elastic(v); - assert_ne!(config.semantic_profile_hash(), elastic, "{v:?}"); - } - } - - // --- 0.8a: legacy equivalence (strict synchronous rounds) --------------- - - #[test] - fn legacy_contract_encoding_appends_nothing() { - let mut encoded = vec![1, 2, 3]; - crate::elastic::IslandSchedulingMode::Legacy.encode_contract(&mut encoded); - assert_eq!(encoded, vec![1, 2, 3]); - assert_eq!( - round_test_config(8).island_scheduling, - crate::elastic::IslandSchedulingMode::Legacy - ); - } - - #[test] - fn legacy_round_needs_every_learner_and_adds_no_extra_wait() { - // Strict legacy: quorum == learners, grace 0. + fn fixed_roster_round_needs_every_learner_and_adds_no_extra_wait() { + // Strict fixed roster: quorum == learners, grace 0. let first = member(0, 10); let second = member(1, 20); let mut round = test_round(vec![first, second]); @@ -5969,7 +6117,7 @@ mod tests { } #[test] - fn legacy_round_timeout_never_merges_a_partial_quorum() { + fn fixed_roster_round_timeout_never_merges_a_partial_quorum() { let first = member(0, 10); let mut round = test_round(vec![first, member(1, 20)]); round.quorum_size = 2; @@ -5980,7 +6128,7 @@ mod tests { } #[test] - fn legacy_strict_round_rejects_a_late_delta_without_carrying_it() { + fn fixed_roster_round_rejects_a_late_delta_without_carrying_it() { let only = member(0, 10); let mut rounds = vec![test_round(vec![only])]; assert_eq!( @@ -5990,72 +6138,6 @@ mod tests { assert!(rounds[0].pushes.is_empty()); } - #[tokio::test] - async fn legacy_server_rejects_every_elastic_frame_type() { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - let key = b"k".to_vec(); - for msg_type in [MSG_JOIN, MSG_LEAVE, MSG_LEASE_HEARTBEAT, MSG_SAMPLE_INDEX, - MSG_DELTA_READY, MSG_ELASTIC_INIT, MSG_DELTA_TENSOR] { - // Bind here and hand the listener over: choosing a free port and - // binding it later lets a parallel test take it in between. - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let mut config = round_test_config(1); - config.port = port; - config.island_hmac_key = Some(key.clone()); - let server = tokio::spawn(run_with_listener(config, Some(listener))); - let mut stream = None; - for _ in 0..100 { - if let Ok(s) = tokio::net::TcpStream::connect(("127.0.0.1", port)).await { - stream = Some(s); - break; - } - tokio::time::sleep(Duration::from_millis(20)).await; - } - let mut stream = stream.expect("legacy syncer did not start"); - // A correctly sealed elastic payload (epoch 0) still must not be - // accepted: legacy only knows HELLO/DATA_HELLO/BW_PROBE as a first frame. - let payload = crate::elastic::seal(&key, msg_type, 0u64.to_le_bytes().to_vec()); - let mut frame = MAGIC.to_le_bytes().to_vec(); - frame.push(msg_type); - frame.extend_from_slice(&(payload.len() as u64).to_le_bytes()); - frame.extend_from_slice(&payload); - stream.write_all(&frame).await.unwrap(); - // Legacy answers with one MSG_ERROR naming the bad first frame - // and then closes the connection. - let mut reply = Vec::new(); - tokio::time::timeout(Duration::from_secs(5), stream.read_to_end(&mut reply)) - .await - .expect("legacy syncer kept the connection open") - .ok(); - assert_eq!(reply.get(4), Some(&MSG_ERROR), "type {msg_type}"); - let text = String::from_utf8_lossy(&reply[13..]); - assert!( - text.contains(&format!("first frame must be HELLO/DATA_HELLO/BW_PROBE, got {msg_type}")), - "type {msg_type}: {text}" - ); - assert!(!server.is_finished(), "type {msg_type}: server exited"); - server.abort(); - } - } - - #[tokio::test] - async fn elastic_mode_requires_an_hmac_key() { - let mut config = round_test_config(1); - config.port = 0; - config.island_scheduling = crate::elastic::IslandSchedulingMode::Elastic( - crate::elastic::ElasticParams { - quorum_theta: 0.75, - carry_gamma: 0.5, - soft_deadline_s: 900, - q_min: 1, - max_carry_lag: 2, - }, - ); - let err = run(config).await.unwrap_err(); - assert!(format!("{err:#}").contains("HMAC key")); - } - fn sweep_test_config(fragments: u32, total_steps: u64) -> Config { let mut config = round_test_config(total_steps); config.policy_sweep_fragments = Some(fragments); @@ -7413,6 +7495,54 @@ mod tests { assert!(started.elapsed() < Duration::from_secs(3)); } + #[tokio::test] + async fn final_grace_finishes_without_a_silent_island_and_tapes_it() { + let state = broadcast_test_state(DTYPE_F32); + let (group, _frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group.clone()); + let expected = HashSet::from([group.member.learner_id]); + let (_event_sender, mut events) = mpsc::channel(1); + let dir = std::env::temp_dir().join(format!("yeto-final-grace-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let tape_path = dir.join("tape.jsonl"); + let _ = std::fs::remove_file(&tape_path); + let mut cfg = round_test_config(0); + cfg.final_ack_timeout_s = 30; + cfg.membership.final_grace = Some(Duration::from_millis(300)); + cfg.event_tape = Some(tape_path.clone()); + let started = Instant::now(); + + finalize_learners(&cfg, &state, &mut events, ®istry, &expected) + .await + .expect("the grace window ends the run without the silent island"); + + assert!(started.elapsed() >= Duration::from_millis(250)); + assert!(started.elapsed() < Duration::from_secs(5), "bounded by the grace, not the ack timeout"); + let tape = std::fs::read_to_string(&tape_path).unwrap(); + let _ = std::fs::remove_dir_all(&dir); + assert!(tape.contains("\"kind\":\"final_grace_expired\""), "{tape}"); + assert!(tape.contains("\"missing\":[0]") && tape.contains("\"acknowledged\":[]"), "{tape}"); + } + + #[tokio::test] + async fn final_grace_stops_waiting_for_an_island_that_left() { + let state = broadcast_test_state(DTYPE_F32); + let (group, _frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group.clone()); + let expected = HashSet::from([group.member.learner_id]); + let (event_sender, mut events) = mpsc::channel(1); + let member = group.member; + registry.lock().unwrap().unregister(member); + event_sender.send(Event::Disconnected { member }).await.unwrap(); + let mut cfg = round_test_config(0); + cfg.membership.final_grace = Some(Duration::from_secs(30)); + let started = Instant::now(); + + finalize_learners(&cfg, &state, &mut events, ®istry, &expected).await.unwrap(); + + assert!(started.elapsed() < Duration::from_secs(5), "a departed island is not waited for"); + } + fn final_leased_registry(group: Arc, lease: Duration) -> (Registry, Config) { let mut cfg = round_test_config(0); cfg.learners = 2; diff --git a/syncer/tests/fixtures/elastic_ledger_golden.txt b/syncer/tests/fixtures/elastic_ledger_golden.txt deleted file mode 100644 index ffa637334..000000000 --- a/syncer/tests/fixtures/elastic_ledger_golden.txt +++ /dev/null @@ -1,25 +0,0 @@ -# generated by tests/test_rl_inter_island_rust_golden.py; do not edit -params 0.75 0.5 1 2 30.0 -join 1 3.0 0.0 => join false 0 -join 2 1.0 0.0 => join false 0 -join 3 1.0 0.0 => join false 0 -submit 1 0 100 4 => delta_accepted -advance 0 => no_step 0 -submit 2 0 60 3 => delta_accepted -advance 0 => outer_step 0 2 4.0 5.0 3 1:0:0.6756756756756757,2:0:0.32432432432432434 -submit 3 0 50 5 => delta_carried_over 1 -submit 1 1 120 4 => delta_accepted -advance 1 => outer_step 1 1 3.0 5.0 2,3 1:1:0.935064935064935,3:0:0.06493506493506493 -join 4 1.0 10.0 => join true 2 -submit 4 2 80 2 => delta_accepted -submit 1 2 90 3 => delta_accepted -submit 2 2 70 7 => delta_accepted -advance 0 => outer_step 2 3 5.0 6.0 3 1:2:0.7941176470588235,2:2:0.20588235294117646,4:2:0.0 -submit 3 0 50 5 => delta_rejected carry_lag_exceeded -submit 2 3 40 2 => delta_accepted -leave 2 => leave 3 40 2 -advance 1 => no_step 3 -submit 4 3 30 3 => delta_accepted -advance 1 => outer_step 3 1 1.0 5.0 1,3 4:3:1.0 -submit 1 3 90 3 => delta_carried_over 1 -advance 0 => no_step 4 diff --git a/tests/pool_fakes.py b/tests/pool_fakes.py new file mode 100644 index 000000000..d4eb3ab17 --- /dev/null +++ b/tests/pool_fakes.py @@ -0,0 +1,105 @@ +"""In-process stand-ins for an island-pool syncer at the session level. + +:class:`FakePoolSyncer` keeps one global LoRA vector and its outer version; +:meth:`FakePoolSyncer.client` returns a :class:`FakePoolClient` with the +interface :class:`yeto.rl.engine.bridges.PoolAvgSync` uses from +:class:`yeto.rl.pool_client.PoolIslandClient`. Each delta is merged at once +(base += update, version += 1). ``final`` ends the run at that version: the +next base is the final cut and a later push raises ``PoolFinished``. +""" + +from __future__ import annotations + +import threading +import time + +import torch + +from yeto.rl.pool_client import JoinAck, PoolBase, PoolFinished + + +class FakePoolSyncer: + def __init__(self, base=None, *, version: int = 0, final: int | None = None, + catch_up: bool | None = None, reply_delay_s: float = 0.0, + base_delay_s: float = 0.0) -> None: + self.base = None if base is None else torch.tensor(base, dtype=torch.float32) + self.version = version + self.final = final + self.catch_up = catch_up + self.reply_delay_s = reply_delay_s + self.base_delay_s = base_delay_s # a merged base becomes visible this long after the push + self.deltas: list[dict] = [] + self.leaves: list[int] = [] + self.inits = 0 + self.lock = threading.Lock() + + def client(self, island_id: int = 0) -> "FakePoolClient": + return FakePoolClient(self, island_id) + + +class FakePoolClient: + def __init__(self, syncer: FakePoolSyncer, island_id: int) -> None: + self.syncer = syncer + self.island_id = island_id + self.latest_base: PoolBase | None = None + self.latest_base_timing: tuple[float, float] | None = None + self.errors: list[str] = [] + self.inner_step = 0 + self.final_outer_version: int | None = None + self.paused: list[float] = [] + self.closed = False + self._on_event = lambda ev: None + + def _publish(self) -> None: + s = self.syncer + if s.base is not None: + final = s.final is not None and s.version >= s.final + self.latest_base = PoolBase(s.base.clone(), s.version, final=final) + self.latest_base_timing = (time.time(), 0.0) + + def join(self) -> JoinAck: + s = self.syncer + catch_up = s.version > 0 if s.catch_up is None else s.catch_up + return JoinAck(base_version=s.version, catch_up=catch_up) + + def wait_base(self, *, newer_than, timeout_s): + if newer_than is not None and self.syncer.base_delay_s: + time.sleep(self.syncer.base_delay_s) + self._publish() + base = self.latest_base + if base is not None and (newer_than is None or base.outer_version > newer_than or base.final): + return base + return None + + def elastic_init(self, params) -> None: + with self.syncer.lock: + self.syncer.inits += 1 + if self.syncer.base is None: + self.syncer.base = params.clone() + + def raise_if_finished(self) -> None: + s = self.syncer + if s.final is not None and s.version >= s.final: + self._publish() + raise PoolFinished(s.final) + + def delta_tensor(self, *, base_version, c_tokens, c_steps, update) -> int: + s = self.syncer + if s.reply_delay_s: + time.sleep(s.reply_delay_s) + with s.lock: + s.deltas.append({"island": self.island_id, "base_version": base_version, + "c_tokens": c_tokens, "c_steps": c_steps, "update": update.clone()}) + s.base = s.base + update + s.version = s.final if s.final is not None else s.version + 1 + return s.version + + def pause_link(self, seconds: float) -> None: + self.paused.append(float(seconds)) + self._on_event({"event": "elastic_debug_pause", "pause_s": float(seconds)}) + + def leave(self, reason: int = 0) -> None: + self.syncer.leaves.append(int(reason)) + + def close(self) -> None: + self.closed = True diff --git a/tests/test_launch_preflight_guards.py b/tests/test_launch_preflight_guards.py index ab9be3986..54696e1bb 100644 --- a/tests/test_launch_preflight_guards.py +++ b/tests/test_launch_preflight_guards.py @@ -406,7 +406,7 @@ def factory(args, spec, m, n, addr): t = launcher._island_task(factory, args, None, 0, 2, "x", {}) assert t.envs == {"A": "1"} and seen[0] is args args = SimpleNamespace(rl_negative_test_run=True, rl_lr_schedule="auto", rl_max_policy_age=0, - rl_island_scheduling="legacy", training_mode="rl", rl_engine="ports") + training_mode="rl", rl_engine="ports") ov = {1: {"identity_test_salt": "s1"}} t0 = launcher._island_task(factory, args, None, 0, 2, "x", ov) t1 = launcher._island_task(factory, args, None, 1, 2, "x", ov) @@ -426,7 +426,7 @@ def test_override_validation(): assert io_.overrides_of(a, 2) == {1: {"rl_lr_schedule": "constant"}} for item, msg in [("2:rl_lr_schedule=constant", "out of range"), ("1:model_revision=abc", "allowed"), - ("1:rl_max_carry_lag=3", "syncer"), + ("1:rl_sync_max_base_lag=3", "syncer"), ("1:rl_lr_schedule=cosine", "must be one of"), ("x:rl_lr_schedule=constant", "ISLAND:KEY=VALUE"), ("1:rl_max_policy_age=-1", ">= 0")]: @@ -455,7 +455,7 @@ def test_identity_salt_changes_contract_only_when_given(): # --- 3.4 manifest fields and island tape --------------------------------------------- def test_manifest_fields_and_tape_event(tmp_path, capsys): args = SimpleNamespace(rl_negative_test_run=True, rl_lr_schedule="auto", - rl_island_scheduling="elastic", rl_checkpoint_store=None) + rl_sync_island_lease_s=30, rl_checkpoint_store=None) ov = {1: {"rl_lr_schedule": "linear"}} m = io_.manifest_fields(args, ov) assert m["negative_test"] is True diff --git a/tests/test_peer_recovery_launch.py b/tests/test_peer_recovery_launch.py index 4edefe565..b2cec9e27 100644 --- a/tests/test_peer_recovery_launch.py +++ b/tests/test_peer_recovery_launch.py @@ -59,7 +59,7 @@ def test_without_peer_recovery_commands_are_unchanged(): (ASYNC + ("--rl-recovery-bound-steps", "4"), "need --rl-peer-recovery"), (ASYNC + ("--rl-peer-recovery-store", "/s"), "need --rl-peer-recovery"), (RECOVERY + ("--rl-recovery-max-donor-tries", "0"), "must be >= 1"), - (ASYNC + ("--rl-syncer-epoch", "2"), "need --rl-island-scheduling elastic"), + (ASYNC + ("--rl-syncer-epoch", "2"), "needs island pool options"), (RECOVERY + ("--syncer-shards", "2"), "one syncer process"), ], ) diff --git a/tests/test_rl_backend_identity.py b/tests/test_rl_backend_identity.py index 3cf28772d..a85a4b079 100644 --- a/tests/test_rl_backend_identity.py +++ b/tests/test_rl_backend_identity.py @@ -152,13 +152,14 @@ def test_island_contract_binds_algorithm_spec(): island_contract_sha256(sha, lr, "not-a-hash") -def test_bridge_configs_bind_lr_schedule_into_hello_and_join(monkeypatch): +def test_bridge_configs_bind_lr_schedule_into_hello(monkeypatch): from types import SimpleNamespace - import yeto.rl.elastic_client as ec + import yeto.rl.pool_client as pc from yeto.rl import bridge, decoupled + from yeto.rl.core import CanonicalTensorSpec, build_avg_layout from yeto.rl.engine.backend_identity import island_contract_sha256, lr_schedule_sha256 - from yeto.rl.engine.bridges import ElasticAvgSync + from yeto.rl.engine.bridges import PoolAvgSync for cls in (bridge.BridgeConfig, decoupled.DecoupledBridgeConfig): assert {"lr_schedule_sha256", "algorithm_semantic_sha256", "code_pack_digest"} <= { @@ -168,14 +169,18 @@ def test_bridge_configs_bind_lr_schedule_into_hello_and_join(monkeypatch): seen = [] class Spy: - def __init__(self, config, key, **kw): - seen.append(config.backend_identity_sha256) + def __init__(self, config, layout, **kw): + seen.append((config.session_contract_hash, layout)) - monkeypatch.setattr(ec, "ElasticIslandClient", Spy) - monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k") + monkeypatch.setattr(pc, "PoolIslandClient", Spy) + spec = CanonicalTensorSpec("base_model.model.layer.lora_A.weight", (1, 2), "float32", 2) cfg = SimpleNamespace(syncer_addr=("x", 0), learner_id=1, backend_identity_sha256=sha, lr_schedule_sha256=lr, algorithm_semantic_sha256="e" * 64, - code_pack_digest="f" * 64, - expected_specs=(), global_rounds=1, local_optimizer_steps=1) - ElasticAvgSync(cfg)._client() - assert seen == [island_contract_sha256(sha, lr, "e" * 64, "f" * 64)] + code_pack_digest="f" * 64, wan_streams=1, + expected_specs=(spec,), global_rounds=1, local_optimizer_steps=1) + sync = PoolAvgSync(cfg) + sync.specs = (spec,) + sync._client() + layout = build_avg_layout((spec,)) + contract = island_contract_sha256(sha, lr, "e" * 64, "f" * 64) + assert seen == [(bridge._session_contract(layout, contract), seen[0][1])] diff --git a/tests/test_rl_convertible_remote.py b/tests/test_rl_convertible_remote.py index 94c0a21eb..b35a4c800 100644 --- a/tests/test_rl_convertible_remote.py +++ b/tests/test_rl_convertible_remote.py @@ -8,7 +8,7 @@ import pytest -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine.convertible import ( CONVERSION_DECISION_KIND, ConvertibleNode, diff --git a/tests/test_rl_decoupled_pool_e2e.py b/tests/test_rl_decoupled_pool_e2e.py new file mode 100644 index 000000000..2634e629f --- /dev/null +++ b/tests/test_rl_decoupled_pool_e2e.py @@ -0,0 +1,190 @@ +"""Decoupled islands in an island pool against the real Rust syncer. CPU, no Ray. + +Two decoupled bridges (``yeto.rl.decoupled``) on a syncer with membership +options: a soft deadline merges rounds without the island that did not push, +the heartbeat keeps an idle island's lease, a silent island's lease expires +and it re-JOINs (catch-up skips the versions it missed), and both islands +then finish the run on the final cut. +""" + +from __future__ import annotations + +import socket +import subprocess +import time +from pathlib import Path + +import pytest +import torch + +from yeto.rl.core import canonical_state +from yeto.rl.decoupled import DecoupledBridgeConfig, DecoupledRlBridge + +ROOT = Path(__file__).resolve().parent.parent +MODEL_REVISION = "a" * 40 +LORA_CONFIG_HASH = "b" * 64 +LEASE_S = 1.5 +SOFT_DEADLINE_S = 1 +TOTAL_STEPS = 6 # three sweeps of two fragments + + +@pytest.fixture(scope="module") +def syncer_binary(): + subprocess.run(["cargo", "build", "-q"], cwd=ROOT / "syncer", check=True) + return ROOT / "syncer/target/debug/yeto-syncer" + + +def _port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _state(version, tensors): + return canonical_state(version, tensors, base_model_revision=MODEL_REVISION, + lora_config_hash=LORA_CONFIG_HASH) + + +def _start_pool_syncer(binary, port, tmp_path, log): + """The decoupled role syncer as the launcher builds it for a pool with + --rl-sync-quorum-theta 1 --rl-sync-soft-deadline-s 1 --rl-sync-island-lease-s 1.5.""" + argv = [str(binary), "--port", str(port), "--learners", "2", "--quorum", "1", + "--grace-ms", "0", "--pipeline", "2", "--sync-interval-steps", "2", + "--delta-correction", "none", "--total-steps", str(TOTAL_STEPS), + "--outer-lr", "0.7", "--outer-momentum", "0.9", "--quorum-timeout-s", "30", + "--quorum-theta", "1.0", "--soft-deadline-s", str(SOFT_DEADLINE_S), + "--island-lease-s", str(LEASE_S), "--syncer-epoch", "0", + "--final-grace-s", "10", "--max-base-lag", "2", "--learner-weight", "equal", + "--checkpoint-path", str(tmp_path / "pool.ckpt"), "--checkpoint-every", "1", + "--event-tape", str(tmp_path / "tape.jsonl")] + return subprocess.Popen(argv, stdout=log, stderr=subprocess.STDOUT) + + +def _bridge(initial, port, learner_id): + return DecoupledRlBridge(initial, DecoupledBridgeConfig( + syncer_addr=("127.0.0.1", port), learner_id=learner_id, total_fragment_steps=TOTAL_STEPS, + num_fragments=2, pipeline=2, local_horizon=2, expected_specs=initial.specs, + base_model_revision=MODEL_REVISION, lora_config_hash=LORA_CONFIG_HASH, + canonical_layout_hash=initial.layout_hash, wan_streams=0, + island_lease_s=LEASE_S, syncer_epoch=0)) + + +def _wait(predicate, timeout=20.0, what="condition"): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + value = predicate() + if value: + return value + time.sleep(0.05) + raise AssertionError(f"timed out waiting for {what}") + + +class _Island: + """One decoupled island: its local state and progress counters.""" + + def __init__(self, bridge): + self.bridge = bridge + self.state = None + self.steps = 0 + self.tokens = 0 + + def take_initial_cut(self): + cut = self.bridge.wait_for_initial_cut(optimizer_steps=0, action_tokens=0) + self.bridge.commit_initial_cut(cut, optimizer_steps=0, action_tokens=0) + self.state = cut.state + + def drain(self): + batch = self.bridge.drain_broadcasts(self.state, optimizer_steps=self.steps, + action_tokens=self.tokens) + self.bridge.commit_broadcasts(batch, optimizer_steps=self.steps, action_tokens=self.tokens) + self.state = batch.state + return self.bridge.fragment_versions + + def train(self, offset): + """One local round of H=2 optimizer steps: every value moves by ``offset``.""" + self.steps += 2 + self.tokens += 20 + self.state = _state(self.steps, {name: value + offset + for name, value in self.state.tensors.items()}) + + def submit(self): + return self.bridge.submit_ready(self.state, optimizer_steps=self.steps, + action_tokens=self.tokens) + + def push_both_fragments(self): + fragments = set() + _wait(lambda: fragments.update(s.fragment_id for s in self.submit()) or fragments == {0, 1}, + what=f"island {self.bridge.config.learner_id} pushes both fragments") + + +def _tape_kinds(path: Path) -> list[str]: + import json + + if not path.exists(): + return [] + return [json.loads(line).get("kind") for line in path.read_text().splitlines() if line.strip()] + + +def test_decoupled_pool_soft_deadline_lease_and_rejoin(syncer_binary, tmp_path): + initial = _state(0, {f"base_model.model.layer{i}.lora_A.weight": torch.zeros(1, i + 1) + for i in range(4)}) + port = _port() + log = open(tmp_path / "syncer.log", "w") + process = _start_pool_syncer(syncer_binary, port, tmp_path, log) + islands = [_Island(_bridge(initial, port, learner_id)) for learner_id in range(2)] + try: + for island in islands: + island.bridge.start() + for island in islands: + island.take_initial_cut() + fast, slow = islands + + # Sweep 1: only island 0 pushes; the soft deadline merges both + # fragment rounds without island 1, which keeps its lease while idle. + fast.train(1.0) + fast.push_both_fragments() + time.sleep(3 * LEASE_S) + _wait(lambda: slow.drain() == (1, 2), what="island 1 sees sweep 1") + assert fast.drain() == (1, 2) + assert slow.bridge.client.lease_expiries == 0, "an idle island lost its lease" + assert fast.bridge.client.lease_expiries == 0 + + # Island 1 goes silent: its lease expires, the client redials (re-JOIN, + # outside the reconnect budget of the fixed roster) and is caught up. + generation = slow.bridge.client.connection_generation + slow.bridge._heartbeat._paused = lambda: True + _wait(lambda: slow.bridge.client.lease_expiries >= 1, what="lease expiry") + slow.bridge._heartbeat._paused = lambda: False + _wait(lambda: slow.bridge.client.connection_generation != generation, what="re-JOIN") + + # Sweep 2 merges without island 1 again; it then sees versions it + # skipped and answers the next sweep's PULLs on its new connection. + fast.drain() + fast.train(1.0) + fast.push_both_fragments() + _wait(lambda: slow.drain() == (3, 4), what="island 1 sees sweep 2") + assert fast.drain() == (3, 4) + + # Sweep 3: both islands push and finish on the final cut. + slow.train(3.0) + fast.train(1.0) + slow.push_both_fragments() + fast.push_both_fragments() + finals = [] + for island in islands: + manifest, final = island.bridge.wait_for_final_cut(policy_version=island.steps) + island.bridge.acknowledge_finalization(manifest) + finals.append(final) + assert process.wait(timeout=30) == 0 + assert finals[0].tensors.keys() == finals[1].tensors.keys() + for name in finals[0].tensors: + assert torch.equal(finals[0].tensors[name], finals[1].tensors[name]) + kinds = _tape_kinds(tmp_path / "tape.jsonl") + assert "pool_leave" in kinds and kinds.count("pool_join") >= 3 # two joins + one re-JOIN + finally: + for island in islands: + island.bridge.close() + if process.poll() is None: + process.kill() + process.wait() + log.close() diff --git a/tests/test_rl_elastic_backend_identity.py b/tests/test_rl_elastic_backend_identity.py deleted file mode 100644 index 9a4fc6f96..000000000 --- a/tests/test_rl_elastic_backend_identity.py +++ /dev/null @@ -1,183 +0,0 @@ -"""Elastic JOIN carries the backend identity hash (yeto-framework-decoupling 6.2a). No Ray. - -The JOIN golden frame is the syncer's own encode output -(syncer/src/elastic.rs ``join_frame_with_backend_identity_golden_and_old_format_refused``); -the real-syncer test needs YETO_TEST_ELASTIC_SYNCER (a built ``yeto-syncer``). -""" - -from __future__ import annotations - -import os -import socket -import struct -import subprocess -import time -from types import SimpleNamespace - -import pytest - -from yeto.rl.adapters.miles.identity import backend_identity -from yeto.rl.elastic_client import ( - MSG_JOIN, ElasticClientConfig, ElasticIslandClient, ElasticProtocolError, Join, decode, encode, - seal, -) -from yeto.rl.engine.backend_identity import BackendIdentity, param_map_sha256 - -JOIN_GOLDEN = ("03000000000000000100000007000000000000000000000000000440" - + "ab" * 32 - + "0b0000006e76696469612d68313030d76ccd494a9ee9bb15d43cdcafb727ab8d0c651b9b81c6481e80ff4084314b4c") # s19-compat: + compat_group "nvidia-h100" - -MILES = backend_identity() -VERL = BackendIdentity("verl", "0" * 40, "nvidia", param_map_sha256({"verl": "map"}), "nvidia-h100") - - -def test_join_golden_frame_matches_rust(): - msg = Join(3, 1, 7, 2.5, bytes([0xab]) * 32, "nvidia-h100") - t, payload = encode(msg, b"k1") - assert t == MSG_JOIN and payload.hex() == JOIN_GOLDEN - assert decode(b"k1", t, payload) == msg - tampered = bytearray(payload) - tampered[28] ^= 1 # first identity byte is under the HMAC - with pytest.raises(ElasticProtocolError): - decode(b"k1", t, bytes(tampered)) - - -def test_old_join_frame_without_identity_is_refused(): - old = seal(b"k1", MSG_JOIN, struct.pack(" int: - with socket.socket() as s: - s.bind(("127.0.0.1", 0)) - return s.getsockname()[1] - - -@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), - reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") -def test_real_syncer_refuses_verl_island_in_miles_session(tmp_path): - """Real Rust syncer, three fake islands: Miles island 1 joins and seeds the base, a - verl island 2 is refused with both hashes in the error, Miles island 3 joins as - before and one outer step merges islands 1 and 3.""" - port, key = _free_port(), "ident-e2e" - syncer = subprocess.Popen( - [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", - "--total-steps", "1", "--event-tape", str(tmp_path / "tape.jsonl"), - "--island-scheduling-mode", "elastic", "--quorum-theta", "1.0", "--carry-gamma", "0.5", - "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", - "--island-lease-s", "6", "--syncer-epoch", "0", "--outer-lr", "1.0", "--outer-momentum", "0.0"], - env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), - stdout=open(tmp_path / "syncer.log", "w"), stderr=subprocess.STDOUT) - - def client(island, ident): - return ElasticIslandClient(ElasticClientConfig( - ("127.0.0.1", port), island, lease_s=6.0, join_timeout_s=5.0, - backend_identity_sha256=ident.sha256()), key.encode()) - - a, v, b = client(1, MILES), client(2, VERL), client(3, MILES) - try: - for _ in range(50): - try: - a.join() - break - except OSError: - time.sleep(0.1) - a.elastic_init([0.0] * 4) - base0 = a.wait_base(newer_than=None, timeout_s=10.0) - with pytest.raises(ElasticProtocolError, match="backend identity mismatch") as exc: - v.join() - text = str(exc.value) - assert VERL.sha256() in text and MILES.sha256() in text and "Miles vs verl" in text - assert b.join() is not None # same backend: unchanged behaviour - b.wait_base(newer_than=None, timeout_s=10.0) - a.delta_tensor(base_version=base0.outer_version, c_tokens=10, c_steps=1, update=[1.0] * 4) - b.delta_tensor(base_version=base0.outer_version, c_tokens=10, c_steps=1, update=[3.0] * 4) - base1 = a.wait_base(newer_than=base0.outer_version, timeout_s=10.0) - assert base1 is not None and base1.outer_version == 1, a.errors - assert list(base1.params) == pytest.approx([2.0] * 4) # equal weights, outer lr 1 - finally: - for c in (a, v, b): - try: - c.close() - except Exception: # noqa: BLE001 - pass - syncer.kill() - syncer.wait(5) - - -@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), - reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") -def test_real_syncer_refuses_same_backend_with_other_lr_schedule(tmp_path): - """S17 N17: two Miles islands whose LR schedules differ declare different bound - identities (island_contract_sha256); the elastic syncer refuses the second JOIN.""" - from yeto.rl.engine.backend_identity import island_contract_sha256, lr_schedule_sha256 - - port, key = _free_port(), "lr-e2e" - syncer = subprocess.Popen( - [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", - "--total-steps", "1", "--event-tape", str(tmp_path / "tape.jsonl"), - "--island-scheduling-mode", "elastic", "--quorum-theta", "1.0", "--carry-gamma", "0.5", - "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", - "--island-lease-s", "6", "--syncer-epoch", "0", "--outer-lr", "1.0", "--outer-momentum", "0.0"], - env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), - stdout=open(tmp_path / "syncer.log", "w"), stderr=subprocess.STDOUT) - const = island_contract_sha256(MILES.sha256(), lr_schedule_sha256("constant", 6, 1e-5)) - other = island_contract_sha256(MILES.sha256(), lr_schedule_sha256("constant", 6, 2e-5)) - - def client(island, ident): - return ElasticIslandClient(ElasticClientConfig( - ("127.0.0.1", port), island, lease_s=6.0, join_timeout_s=5.0, - backend_identity_sha256=ident), key.encode()) - - a, bad = client(1, const), client(2, other) - try: - for _ in range(50): - try: - a.join() - break - except OSError: - time.sleep(0.1) - a.elastic_init([0.0] * 4) - a.wait_base(newer_than=None, timeout_s=10.0) - with pytest.raises(ElasticProtocolError, match="backend identity mismatch") as exc: - bad.join() - assert const in str(exc.value) and other in str(exc.value) - finally: - for c in (a, bad): - try: - c.close() - except Exception: # noqa: BLE001 - pass - syncer.kill() - syncer.wait() diff --git a/tests/test_rl_elastic_debug_pause.py b/tests/test_rl_elastic_debug_pause.py deleted file mode 100644 index 76df5267b..000000000 --- a/tests/test_rl_elastic_debug_pause.py +++ /dev/null @@ -1,213 +0,0 @@ -"""Test switch --rl-elastic-debug-pause: the island's link to the syncer goes silent for S -seconds (no heartbeat, sends/receives held) while the process keeps running, to exercise -lease expiry -> "rejoin required" -> automatic re-JOIN on real hardware. CPU, no Ray. - -The end-to-end case runs the real Rust syncer when YETO_TEST_ELASTIC_SYNCER points to it.""" - -from __future__ import annotations - -import json -import os -import subprocess -import time -from pathlib import Path - -import pytest -import torch - -from test_rl_elastic_lease_rejoin import MultiConnSyncer, _client, _wait -from test_rl_engine_driver import _driver, _engine, _run_threads, _strict_config -from test_rl_inter_island_elastic_client import FakeSyncer -from test_rl_inter_island_ports import _free_port, _sync -from yeto.rl.elastic_client import ( - DeltaTensor, ElasticClientConfig, ElasticIslandClient, LeaseHeartbeat, parse_elastic_debug_pause, -) - -PAUSE_ENV = "YETO_RL_ELASTIC_DEBUG_PAUSE" - - -def test_parse_and_launcher_env(): - from test_rl_launcher import _args - from yeto import launcher - - assert parse_elastic_debug_pause("1:1:120") == {1: (1, 120.0)} - assert parse_elastic_debug_pause(" 0:2:5, 1:1:7.5 ") == {0: (2, 5.0), 1: (1, 7.5)} - assert parse_elastic_debug_pause("") == {} and parse_elastic_debug_pause(None) == {} - for bad in ("1:120", "1:1:0", "-1:1:5", "1:1:5:9", "a:1:5"): - with pytest.raises(ValueError): - parse_elastic_debug_pause(bad) - assert launcher.elastic_debug_pause_env(_args()) == {} - assert launcher.elastic_debug_pause_env(_args(("--rl-island-scheduling", "elastic"))) == {} - assert launcher.elastic_debug_pause_env(_args(("--rl-island-scheduling", "elastic", - "--rl-elastic-debug-pause", "1:1:120"))) == { - PAUSE_ENV: "1:1:120"} - with pytest.raises(ValueError, match="needs --rl-island-scheduling elastic"): - launcher.elastic_debug_pause_env(_args(("--rl-elastic-debug-pause", "1:1:120"))) - - -def test_default_off_island_task_envs_unchanged(monkeypatch, tmp_path): - """Without the flag the island task / Modal cfg carry no new env (same bytes as before).""" - from test_rl_inter_island_launcher import _fake_sky, _prov - from test_rl_launcher import _args, _prepare_rl_args - from yeto import launcher - from yeto.gpu_spec import parse_gpu_spec - - _fake_sky(monkeypatch) - monkeypatch.setattr("yeto.launch.rl_island_task.submitted_unix", lambda: 1.0) - monkeypatch.setenv("HOME", str(tmp_path)) - monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k") - base = ("--gpu", "modal:1xh100,modal:1xh100", "--rl-engine", "ports", "--rl-island-scheduling", "elastic") - - def build(extra): - args = _args(base + extra) - _prepare_rl_args(_prov(args)) - spec = parse_gpu_spec(args.gpu)[1] - task = launcher.make_miles_island_task(args, spec, 1, 2, "127.0.0.1:29400") - return task, launcher.build_modal_island_config(args, spec, 1, task, "1.2.3.4:5000") - - off, off_cfg = build(()) - on, on_cfg = build(("--rl-elastic-debug-pause", "1:1:120")) - assert PAUSE_ENV not in (off.envs or {}) and PAUSE_ENV not in off_cfg.envs - assert on.envs[PAUSE_ENV] == "1:1:120" and on_cfg.envs[PAUSE_ENV] == "1:1:120" - assert {k: v for k, v in on.envs.items() if k != PAUSE_ENV} == off.envs - assert on.run == off.run and on.setup == off.setup and on_cfg.run_script == off_cfg.run_script - - -def test_default_off_bridge_never_pauses(tmp_path, monkeypatch): - monkeypatch.delenv(PAUSE_ENV, raising=False) - fake = FakeSyncer() - engine = _engine(torch.tensor([1.0, 3.0])) - client = ElasticIslandClient(ElasticClientConfig(("x", 0), 0, lease_s=3.0, join_timeout_s=2), - fake.key, connect=fake.connect) - called = [] - client.pause_link = called.append - sync = _sync(tmp_path, engine, client) - _driver(engine, sync, tmp_path).run() - assert called == [] and client._pause_until == 0.0 and not client.link_paused() - assert "elastic_debug_pause" not in (tmp_path / "events.jsonl").read_text() - deltas = [m for m in fake.frames if isinstance(m, DeltaTensor)] - assert [(d.base_version, d.update) for d in deltas] == [(0, (1.0, 3.0)), (1, (1.0, 3.0))] - - -def test_bridge_pauses_only_the_named_island_once_after_after_v(tmp_path, monkeypatch): - monkeypatch.setenv(PAUSE_ENV, "0:1:0.05,5:0:9") - fake = FakeSyncer() - engine = _engine(torch.tensor([1.0, 3.0])) - client = ElasticIslandClient(ElasticClientConfig(("x", 0), 0, lease_s=3.0, join_timeout_s=2), - fake.key, connect=fake.connect) - sync = _sync(tmp_path, engine, client, rounds=3) - _driver(engine, sync, tmp_path).run() - ev = [json.loads(l) for l in (tmp_path / "events.jsonl").read_text().splitlines()] - armed = [e for e in ev if e.get("phase") == "elastic_debug_pause_armed" or e.get("event") == "elastic_debug_pause_armed"] - assert len(armed) == 1 and armed[0].get("base_version") == 1 and armed[0].get("pause_s") == 0.05 - assert [e["pause_s"] for e in ev if e.get("event") == "elastic_debug_pause"] == [0.05] - - -def test_pause_silences_heartbeats_and_holds_sends(): - fake, events = MultiConnSyncer(), [] - c = _client(fake, events) # lease 0.15 s -> heartbeat every 0.05 s - c.join() - hb = lambda: sum(isinstance(m, LeaseHeartbeat) for _, m in fake.frames) # noqa: E731 - assert _wait(lambda: hb() >= 2) - c.pause_link(0.6) - time.sleep(0.1) - n0 = hb() - t0 = time.monotonic() - c.delta_tensor(base_version=0, c_tokens=1, c_steps=1, update=(1.0, 1.0)) # held until the pause ends - assert time.monotonic() - t0 >= 0.4 - assert hb() == n0 or hb() <= n0 + 1 # at most one heartbeat raced the end of the pause - assert _wait(lambda: hb() >= n0 + 3) # renewal resumes afterwards - assert _wait(lambda: any(e.get("event") == "elastic_debug_pause_end" for e in events)) - c.close() - - -@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), - reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") -def test_link_pause_longer_than_lease_rejoins_against_real_rust_syncer(tmp_path, monkeypatch): - """Island 1 goes silent for 5 s (lease 2 s) after v1: the syncer expires its lease, - the island's heartbeat is answered "rejoin required", it re-JOINs (catch_up) and then - takes part in later outer steps; both islands end normally and the syncer exits 0.""" - port, key = _free_port(), "pause-e2e" - out = Path(os.environ.get("YETO_TEST_ELASTIC_OUT", tmp_path)) - out.mkdir(parents=True, exist_ok=True) - for old in out.glob("*.jsonl"): - old.unlink() - monkeypatch.setenv(PAUSE_ENV, "1:1:5") - monkeypatch.setenv("YETO_RL_ELASTIC_DEBUG_DELAY", "0:1.5,1:0.5") # keep island 0 slow - rounds = 8 - syncer = subprocess.Popen( - [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", - "--total-steps", str(rounds), "--event-tape", str(out / "syncer-tape.jsonl"), - "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", - "--soft-deadline-s", "2", "--q-min", "1", "--max-carry-lag", "2", - "--island-lease-s", "2", "--final-grace-s", "10", "--syncer-epoch", "0"], - env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), - stdout=open(out / "syncer.log", "w"), stderr=subprocess.STDOUT) - try: - time.sleep(1.0) - engines = (_engine(torch.tensor([1.0, 3.0])), _engine(torch.tensor([3.0, 5.0]))) - drivers = [] - for i, e in enumerate(engines): - c = ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), i, lease_s=2.0), - key.encode()) - cfg = _strict_config(out, e, learner_id=i, rounds=rounds, tape=f"island-{i}.jsonl") - drivers.append(_driver(e, ElasticAvgSync_(cfg, c), out, name=f"events-{i}.jsonl", learner_id=i)) - results, errors = _run_threads(drivers) - assert errors == {}, errors - assert syncer.wait(30) == 0 - rows = [json.loads(l) for l in (out / "syncer-tape.jsonl").read_text().splitlines()] - leaves = [r for r in rows if r.get("kind") == "pool_leave" and r.get("island_id") == 1 - and r.get("reason") == "lease_expired"] - assert len(leaves) == 1, [r for r in rows if r.get("kind", "").startswith("pool")] - joins = [r for r in rows if r.get("kind") == "pool_join" and r.get("island_id") == 1 - and r["membership_epoch"] > leaves[0]["membership_epoch"]] - assert joins and joins[0]["catch_up"] is True - assert "1 is not a member (rejoin required)" in (out / "syncer.log").read_text() - ev1 = [json.loads(l) for l in (out / "events-1.jsonl").read_text().splitlines()] - names = [e.get("event") for e in ev1] - assert names.count("elastic_debug_pause") == 1 and names.count("elastic_debug_pause_end") == 1 - rj = [e for e in ev1 if e.get("event") == "elastic_rejoin"] - assert len(rj) == 1 and rj[0]["catch_up"] is True - # after the re-JOIN island 1 trains and its delta is merged with weight > 0 - idx = rows.index(joins[0]) - later = [r for r in rows[idx:] if r.get("kind") == "outer_step"] - w1 = [[x[2] for x in (r.get("weights") or []) if x[0] == 1] for r in later] - assert any(v and v[0] > 0 for v in w1), later - # S17 M1: the re-JOINed island also takes part in the LAST outer step (it ends on the - # syncer's final version, not on its local round count) - assert any(x[0] == 1 for x in (later[-1].get("raw_weights") or [])), later[-1] - assert names.count("elastic_finished") <= 1 - finally: - if syncer.poll() is None: - syncer.kill() - - -def ElasticAvgSync_(cfg, client): # noqa: N802 - from yeto.rl.engine.bridges import ElasticAvgSync - - return ElasticAvgSync(cfg, client=client, base_wait_s=60.0) - - -def test_self_built_client_events_reach_the_tape_once(tmp_path, monkeypatch): - """1b-g: a client built by ElasticAvgSync._client() was hooked twice, so every - elastic_rejoin / elastic_debug_pause appeared twice on the island tape.""" - from yeto.rl.engine import bridges - - fake = FakeSyncer() - engine = _engine(torch.tensor([1.0, 3.0])) - sync = _sync(tmp_path, engine, None) - built = [] - - def make(*a, **kw): - c = ElasticIslandClient(ElasticClientConfig(("x", 0), 0, lease_s=3.0, join_timeout_s=2), - fake.key, connect=fake.connect, on_event=kw.get("on_event")) - built.append(c) - return c - import yeto.rl.elastic_client as ec - monkeypatch.setattr(ec, "ElasticIslandClient", make) - monkeypatch.setattr(ec, "hmac_key_from_env", lambda: fake.key) - monkeypatch.setenv(PAUSE_ENV, "0:1:0.05") - _driver(engine, sync, tmp_path).run() - assert built and bridges.ElasticAvgSync - ev = [json.loads(l) for l in (tmp_path / "events.jsonl").read_text().splitlines()] - assert [e.get("event") for e in ev].count("elastic_debug_pause") == 1 diff --git a/tests/test_rl_elastic_finished.py b/tests/test_rl_elastic_finished.py deleted file mode 100644 index e67b50bd3..000000000 --- a/tests/test_rl_elastic_finished.py +++ /dev/null @@ -1,124 +0,0 @@ -"""0.26: a late island ends normally when the syncer finished (FINISHED or closed -connection after the final version) instead of BrokenPipe / exit 1. CPU, no Ray.""" - -from __future__ import annotations - -import json -import socket -import threading - -import pytest -import torch - -from yeto.protocol import read_frame, write_frame -from yeto.rl.elastic_client import ( - MSG_FINISHED, DeltaTensor, ElasticBase, ElasticClientConfig, ElasticFinished, ElasticInit, - ElasticIslandClient, Finished, Join, JoinAck, Leave, decode, encode, -) - -from test_rl_engine_driver import _driver, _engine, _strict_config -from yeto.rl.engine.bridges import ElasticAvgSync - -KEY = b"k" - - -class FinishingSyncer: - """Other islands push the run to ``final`` with this island's first delta; the - syncer then (optionally) sends FINISHED and goes away.""" - - def __init__(self, *, final, send_finished): - self.a, self.b = socket.socketpair() - self.final, self.send_finished = final, send_finished - self.version, self.base = 0, None - self.leaves = 0 - threading.Thread(target=self._serve, daemon=True).start() - - def _base(self): - write_frame(self.b, *encode(ElasticBase(0, self.version, tuple(self.base)), KEY)) - - def _serve(self): - while True: - try: - t, p = read_frame(self.b, 1 << 20) - except Exception: - return - msg = decode(KEY, t, p) - if isinstance(msg, Join): - write_frame(self.b, *encode(JoinAck(0, msg.island_id, 1, 0, bytes(32), False), KEY)) - elif isinstance(msg, ElasticInit): - self.base = list(msg.params) - self._base() - elif isinstance(msg, DeltaTensor): - self.version = self.final - self.base = [b + u for b, u in zip(self.base, msg.update)] - self._base() - if self.send_finished: - write_frame(self.b, *encode(Finished(0, self.final, bytes([5]) * 32), KEY)) - continue # the grace window: the island answers with LEAVE - self.b.close() - return - elif isinstance(msg, Leave): - self.leaves += 1 - self.b.close() - return - - def connect(self, addr): - return self.a - - -def test_finished_frame_roundtrip(): - msg = Finished(3, 6, bytes([9]) * 32) - t, payload = encode(msg, KEY) - assert t == MSG_FINISHED == 25 and decode(KEY, t, payload) == msg - assert len(payload) == 8 + 8 + 32 + 32 - - -@pytest.mark.parametrize("send_finished", [True, False]) -def test_late_island_ends_normally_when_the_syncer_finished(tmp_path, send_finished): - fake = FinishingSyncer(final=3, send_finished=send_finished) - engine = _engine(torch.tensor([1.0, 3.0])) - client = ElasticIslandClient( - ElasticClientConfig(("x", 0), 0, lease_s=30.0, join_timeout_s=2, heartbeat_connection=False), - KEY, connect=fake.connect) - sync = ElasticAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=3), client=client, - base_wait_s=5.0) - _driver(engine, sync, tmp_path).run() # no exception: a normal end (exit 0) - tape = [json.loads(l) for l in (tmp_path / "events.jsonl").read_text().splitlines()] - done = [e for e in tape if e["event"] == "elastic_finished"] - # S17 M1: the island stops as soon as it applied the final outer version (here right after its - # first delta), so it no longer pushes one more delta into the finished syncer; the FINISHED - # path (elastic_finished) is then only taken when the push races the end (next test). - assert sync.base_version == 3 and done == [] - assert sum(e["event"] == "rl_local_round" for e in tape) == 1 - if send_finished: # LEAVE first so the syncer can exit before its grace window ends - assert fake.leaves == 1 - client.close() - - -def test_disconnect_before_the_final_version_is_still_an_error(): - c = ElasticIslandClient(ElasticClientConfig(("x", 0), 0), KEY) - c.final_outer_version = 6 - c.latest_base = ElasticBase(0, 4, (0.0,)) - c.disconnected = True - c.raise_if_finished() # not final yet: no ElasticFinished - c.latest_base = ElasticBase(0, 6, (0.0,)) - with pytest.raises(ElasticFinished): - c.raise_if_finished() - - -def test_push_racing_the_end_takes_the_finished_path(tmp_path): - """The base the island holds is not final yet, but the syncer finished meanwhile: - FINISHED arrives before the next push and the island ends via elastic_finished.""" - fake = FinishingSyncer(final=3, send_finished=True) - engine = _engine(torch.tensor([1.0, 3.0])) - client = ElasticIslandClient( - ElasticClientConfig(("x", 0), 0, lease_s=30.0, join_timeout_s=2, heartbeat_connection=False), - KEY, connect=fake.connect) - sync = ElasticAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=4), client=client, - base_wait_s=5.0) - _driver(engine, sync, tmp_path).run() - tape = [json.loads(l) for l in (tmp_path / "events.jsonl").read_text().splitlines()] - done = [e for e in tape if e["event"] == "elastic_finished"] - assert len(done) == 1 and done[0]["final_outer_version"] == 3 and done[0]["reason"] == "FINISHED" - assert fake.leaves == 1 - client.close() diff --git a/tests/test_rl_elastic_fleet_end.py b/tests/test_rl_elastic_fleet_end.py index 3071ce64c..644edcea4 100644 --- a/tests/test_rl_elastic_fleet_end.py +++ b/tests/test_rl_elastic_fleet_end.py @@ -89,9 +89,9 @@ def test_final_grace_flag_and_status_state(tmp_path, monkeypatch): from yeto import launcher monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") - cmd = launcher.syncer_command(_args(("--rl-island-scheduling", "elastic", "--rl-soft-deadline-s", "240")), 2) + cmd = launcher.syncer_command(_args(("--rl-sync-quorum-theta", "0.75", "--rl-sync-soft-deadline-s", "240")), 2) assert " --final-grace-s 240" in cmd - cmd = launcher.syncer_command(_args(("--rl-island-scheduling", "elastic", "--rl-final-grace-s", "60")), 2) + cmd = launcher.syncer_command(_args(("--rl-sync-quorum-theta", "0.75", "--rl-sync-final-grace-s", "60")), 2) assert " --final-grace-s 60" in cmd assert "--final-grace-s" not in launcher.syncer_command(_args(), 2) @@ -106,15 +106,15 @@ def test_final_grace_flag_and_status_state(tmp_path, monkeypatch): assert ls.finished() -# -- 10-09 ruling: elastic is not a fixed roster ----------------------------------- -def test_controller_fixed_roster_only_for_non_elastic_rl(): +# -- 10-09 ruling: an island pool is not a fixed roster ----------------------------------- +def test_controller_fixed_roster_only_without_island_pool(): from types import SimpleNamespace from yeto.launcher import controller_fixed_roster - assert controller_fixed_roster(SimpleNamespace(training_mode="rl", rl_island_scheduling="legacy")) - assert not controller_fixed_roster(SimpleNamespace(training_mode="rl", rl_island_scheduling="elastic")) - assert not controller_fixed_roster(SimpleNamespace(training_mode="sft", rl_island_scheduling="legacy")) + assert controller_fixed_roster(SimpleNamespace(training_mode="rl")) + assert not controller_fixed_roster(SimpleNamespace(training_mode="rl", rl_sync_island_lease_s=30)) + assert not controller_fixed_roster(SimpleNamespace(training_mode="sft")) def _refusing_ops(text): diff --git a/tests/test_rl_elastic_lease_rejoin.py b/tests/test_rl_elastic_lease_rejoin.py deleted file mode 100644 index e0c06eebb..000000000 --- a/tests/test_rl_elastic_lease_rejoin.py +++ /dev/null @@ -1,204 +0,0 @@ -"""0.22: lease renewal on its own connection, automatic re-JOIN, give-up after N. CPU, no Ray.""" - -from __future__ import annotations - -import socket -import threading -import time - -import pytest - -from yeto.protocol import MSG_ERROR, read_frame, write_frame -from yeto.rl.elastic_client import ( - ELASTIC_REJOIN_FAILED_EXIT, DeltaTensor, ElasticBase, ElasticClientConfig, ElasticIslandClient, - ElasticInit, ElasticRejoinError, Join, JoinAck, LeaseHeartbeat, decode, encode, -) - -KEY = b"k" - - -class MultiConnSyncer: - """Fake elastic server with one socketpair per connect(): membership by island id, - heartbeats from a non-member answered like elastic.rs (MSG_ERROR ... rejoin required).""" - - def __init__(self, *, refuse_rejoin=False): - self.members: set[int] = set() - self.version, self.base = 0, [0.0, 0.0] - self.frames: list[tuple[int, object]] = [] # (connection index, message) - self.refuse_rejoin = refuse_rejoin - self.joins = 0 - self.conns = [] - - def connect(self, addr): - a, b = socket.socketpair() - idx = len(self.conns) - self.conns.append(b) - threading.Thread(target=self._serve, args=(b, idx), daemon=True).start() - return a - - def _base(self, sock): - write_frame(sock, *encode(ElasticBase(0, self.version, tuple(self.base)), KEY)) - - def _serve(self, sock, idx): - while True: - try: - t, p = read_frame(sock, 1 << 20) - except Exception: - return - msg = decode(KEY, t, p) - self.frames.append((idx, msg)) - if isinstance(msg, Join): - self.joins += 1 - if self.joins > 1 and self.refuse_rejoin: - write_frame(sock, MSG_ERROR, b"join refused (test)") - continue - self.members.add(msg.island_id) - write_frame(sock, *encode(JoinAck(0, msg.island_id, self.joins, self.version, - bytes(32), self.joins > 1), KEY)) - self._base(sock) - elif isinstance(msg, LeaseHeartbeat) and msg.island_id not in self.members: - write_frame(sock, MSG_ERROR, f"{msg.island_id} is not a member (rejoin required)".encode()) - elif isinstance(msg, ElasticInit): - self._base(sock) - elif isinstance(msg, DeltaTensor): - if msg.island_id not in self.members: - write_frame(sock, MSG_ERROR, f"{msg.island_id} is not a member".encode()) - continue - self.base = [b + u for b, u in zip(self.base, msg.update)] - self.version += 1 - self._base(sock) - - def expire(self, island_id): - self.members.discard(island_id) - - -def _client(fake, events, **kw): - cfg = ElasticClientConfig(("x", 0), 0, lease_s=0.15, join_timeout_s=1.0, **kw) - return ElasticIslandClient(cfg, KEY, connect=fake.connect, on_event=events.append) - - -def _wait(cond, timeout=3.0): - end = time.monotonic() + timeout - while time.monotonic() < end: - if cond(): - return True - time.sleep(0.01) - return False - - -def test_heartbeats_use_their_own_connection_while_main_send_is_blocked(): - fake, events = MultiConnSyncer(), [] - c = _client(fake, events) - c.join() - assert c.hb_sock is not None and len(fake.conns) == 2 - with c._send_lock: # a multi-MB DELTA_TENSOR in flight on the main connection - n0 = sum(isinstance(m, LeaseHeartbeat) for _, m in fake.frames) - assert _wait(lambda: sum(isinstance(m, LeaseHeartbeat) for _, m in fake.frames) >= n0 + 3) - assert {i for i, m in fake.frames if isinstance(m, LeaseHeartbeat)} == {1} - c.close() - - -def test_not_a_member_triggers_catch_up_rejoin_and_wait_returns_post_rejoin_base(): - fake, events = MultiConnSyncer(), [] - c = _client(fake, events) - c.join() - first = c.wait_base(newer_than=None, timeout_s=1.0) - fake.expire(0) # lease expired on the syncer while the island trained - c.delta_tensor(base_version=first.outer_version, c_tokens=1, c_steps=1, update=[1.0, 1.0]) - base = c.wait_base(newer_than=first.outer_version, timeout_s=3.0) # delta refused: no newer base - assert base is not None and base.outer_version == first.outer_version - rejoin = [e for e in events if e["event"] == "elastic_rejoin"] - assert rejoin and rejoin[0]["incarnation"] == 1 and rejoin[0]["catch_up"] is True - assert c.rejoins == 1 and c.ack.catch_up and 0 in fake.members - joins = [m for _, m in fake.frames if isinstance(m, Join)] - assert [j.incarnation for j in joins] == [0, 1] - c.close() - - -def test_gives_up_after_n_consecutive_rejoin_failures(): - fake, events = MultiConnSyncer(refuse_rejoin=True), [] - c = _client(fake, events, max_rejoin_failures=3) - c.config.join_timeout_s = 0.2 - c.join() - first = c.wait_base(newer_than=None, timeout_s=1.0) - fake.expire(0) - with pytest.raises(ElasticRejoinError, match="3 consecutive re-JOINs failed"): - c.wait_base(newer_than=first.outer_version, timeout_s=10.0) - assert [e["failures"] for e in events if e["event"] == "elastic_rejoin_failed"] == [1, 2, 3] - assert ELASTIC_REJOIN_FAILED_EXIT not in (4, 6) - c.close() - - -def test_learner_maps_rejoin_exhaustion_to_its_own_exit_code(): - from yeto.rl.adapters.miles.island_entry import _rejoin_exhausted - - try: - try: - raise ElasticRejoinError("x") - except ElasticRejoinError as inner: - raise RuntimeError("driver failed") from inner - except RuntimeError as outer: - assert _rejoin_exhausted(outer) - assert not _rejoin_exhausted(RuntimeError("other")) - - -def test_launcher_island_lease_default_and_flag(monkeypatch): - from test_rl_launcher import _args - from yeto import launcher - - monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") - cmd = launcher.syncer_command(_args(("--rl-island-scheduling", "elastic")), 2) - assert " --island-lease-s 150.0" in cmd - cmd = launcher.syncer_command(_args(("--rl-island-scheduling", "elastic", "--rl-island-lease-s", "450")), 2) - assert " --island-lease-s 450.0" in cmd - assert "--island-lease-s" not in launcher.syncer_command(_args(), 2) - with pytest.raises(ValueError, match="--rl-island-lease-s"): - launcher.syncer_command(_args(("--rl-island-lease-s", "450")), 2) # legacy refuses it - - -def test_rejoin_events_reach_the_island_tape(tmp_path): - import json - - import torch - - from test_rl_engine_driver import _driver, _engine, _strict_config - from yeto.rl.engine.bridges import ElasticAvgSync - - fake, events = MultiConnSyncer(), [] - c = _client(fake, events) - engine = _engine(torch.tensor([1.0, 3.0])) - sync = ElasticAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=3), client=c, - base_wait_s=5.0) - orig = engine.trainer.train_step - - def train_step(batch): # lease expires while round 1 trains - if fake.version == 1 and 0 in fake.members: - fake.expire(0) - _wait(lambda: c.rejoins >= 1) - return orig(batch) - engine.trainer.train_step = train_step - _driver(engine, sync, tmp_path).run() - tape = [json.loads(x) for x in (tmp_path / "events.jsonl").read_text().splitlines()] - assert [e["incarnation"] for e in tape if e["event"] == "elastic_rejoin"] == [1] - c.close() - - -def test_island_lease_default_is_the_same_in_syncer_cli_and_launcher(): - """rl-code-quality-cleanup 0.8 / rl-island-graceful-leave 3.3: a syncer started without --island-lease-s - uses the launcher's default.""" - import re - from pathlib import Path - - from yeto import launcher - - repo = Path(__file__).resolve().parents[1] - main_rs = (repo / "syncer/src/main.rs").read_text() - m = re.search(r"const DEFAULT_ELASTIC_LEASE_S: f64 = ([0-9.]+);", main_rs) - assert m, "island_lease_s default not found in syncer/src/main.rs" - from yeto.rl.elastic_client import DEFAULT_ISLAND_LEASE_S, ElasticClientConfig - from yeto.rl.engine.run_flags import rl_flags - - assert float(m.group(1)) == launcher.DEFAULT_ISLAND_LEASE_S == DEFAULT_ISLAND_LEASE_S == 150.0 - assert ElasticClientConfig(("x", 0), 0).lease_s == DEFAULT_ISLAND_LEASE_S # island side - lease = next(f for f in rl_flags() if f.option == "--rl-island-lease-s") - assert f"Default {launcher.DEFAULT_ISLAND_LEASE_S:g}" in lease.kwargs["help"] diff --git a/tests/test_rl_engine_selection.py b/tests/test_rl_engine_selection.py index 80c99a716..b7a653d4e 100644 --- a/tests/test_rl_engine_selection.py +++ b/tests/test_rl_engine_selection.py @@ -61,11 +61,9 @@ def test_ppo_is_not_routed_to_legacy(): @pytest.mark.parametrize( "kwargs, reason", [ - (dict(sync_preset="full-parameter", parameter_mode="full", island_scheduling="elastic"), "elastic"), (dict(sync_preset="sao-streaming-full", parameter_mode="full", - island_scheduling="elastic"), "elastic"), + island_scheduling="elastic"), "island pool"), (dict(expert_full_count=4), "expert-full"), - (dict(sync_preset="full-parameter", expert_full_count=4, island_scheduling="elastic"), "elastic"), (dict(use_critic=True), "critic"), (dict(extra_argv=("--use-critic",)), "critic"), (dict(extra_argv=("--advantage-estimator", "gspo")), "non-GRPO"), diff --git a/tests/test_rl_full_parameter_async.py b/tests/test_rl_full_parameter_async.py index 4401090f8..002f756f1 100644 --- a/tests/test_rl_full_parameter_async.py +++ b/tests/test_rl_full_parameter_async.py @@ -126,20 +126,52 @@ async def test_closed_round_is_not_uploaded(): @pytest.mark.asyncio -async def test_out_of_order_broadcasts_apply_in_version_order(): +async def test_in_order_broadcasts_of_one_drain_apply_in_version_order(): actor = _PairedRole("actor", [], []) pairing = _async_pairing(actor) - actor._staged_steps.append(3) # fragment 0, version 3: its base 1 is still on the wire + actor._staged_steps.extend([3, 2, 1]) committed = await pairing.commit_staged_broadcasts() - assert committed["actor"].applied_fragments == () - assert actor.fragment_versions == (0, 0) - actor._staged_steps.extend([2, 1]) - committed = await pairing.commit_staged_broadcasts() - assert actor.fragment_versions == (3, 2) + assert actor.fragment_versions == (3, 2) and actor.caught_up == [] assert set(committed["actor"].applied_fragments) == {0, 1} - assert pairing.overwritten_rounds == 3 + assert pairing.overwritten_rounds == 3 and pairing.caught_up_fragments == 0 + + +@pytest.mark.asyncio +async def test_a_version_that_skips_rounds_is_a_catch_up(): + # A re-JOIN catch-up (or a newer broadcast that overtook an older one, + # which the client then drops) skips versions: the island overwrites the + # fragment and releases the retained pairs of the skipped rounds. + actor = _PairedRole("actor", [_candidate("actor", 1)], [], release_broadcasts=False) + pairing = _async_pairing(actor) + assert _submitted(pairing.submit_ready()) == [1] + actor._queued_steps.clear() + actor._staged_steps.append(3) # fragment 0, version 3: version 1 never arrives + + committed = await pairing.commit_staged_broadcasts() + + assert committed["actor"].applied_fragments == (0,) + assert actor.fragment_versions == (3, 0) and actor.caught_up == [(0, 3)] + assert pairing.inflight == {} and pairing.caught_up_fragments == 1 + + actor._staged_steps.append(2) # fragment 1 continues in order + await pairing.commit_staged_broadcasts() + assert actor.fragment_versions == (3, 2) and pairing.caught_up_fragments == 1 + + +@pytest.mark.asyncio +async def test_a_catch_up_waits_until_every_role_staged_it(): + actor, critic = _PairedRole("actor", [], []), _PairedRole("critic", [], []) + pairing = _async_pairing(actor, critic) + actor._staged_steps.append(3) + + await pairing.commit_staged_broadcasts() + assert actor.fragment_versions == critic.fragment_versions == (0, 0) + + critic._staged_steps.append(3) + await pairing.commit_staged_broadcasts() + assert actor.fragment_versions == critic.fragment_versions == (3, 0) @pytest.mark.asyncio diff --git a/tests/test_rl_full_parameter_sync.py b/tests/test_rl_full_parameter_sync.py index 2bdb5c972..07efb2c09 100644 --- a/tests/test_rl_full_parameter_sync.py +++ b/tests/test_rl_full_parameter_sync.py @@ -7,7 +7,7 @@ import pytest from yeto.protocol import PartialMessageGenerationLost, PullRequest -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.adapters.miles.full_parameter.role_stream import ( RoleRoundResult, RoleSubmissionCandidate, @@ -70,6 +70,7 @@ def __init__( self.terminal_submitted = False self._queued_steps = [] self._staged_steps = [] + self.caught_up = [] self.local_boundaries = 0 self.closed = False self.final_acknowledged = False @@ -158,10 +159,14 @@ def staged_fragment_versions(self): def has_staged_broadcast(self, fragment_id, version): return (fragment_id, version) in self.staged_fragment_versions - async def apply_staged_broadcast(self, fragment_id, version): + async def apply_staged_broadcast(self, fragment_id, version, *, catch_up=False): versions = list(self.fragment_versions) assert (fragment_id, version) in self.staged_fragment_versions self._staged_steps.remove(version) + if catch_up: + self.caught_up.append((fragment_id, version)) + self._staged_steps[:] = [step for step in self._staged_steps + if (step - 1) % len(versions) != fragment_id or step > version] versions[fragment_id] = version self.fragment_versions = tuple(versions) return (fragment_id,) diff --git a/tests/test_rl_inter_island_contract.py b/tests/test_rl_inter_island_contract.py index bdf1bcd5e..c4c0408a7 100644 --- a/tests/test_rl_inter_island_contract.py +++ b/tests/test_rl_inter_island_contract.py @@ -1,4 +1,4 @@ -"""Syncer contract hash with the island-scheduling block (Python side of elastic.rs). No Ray.""" +"""Syncer semantic profile hash: island pool membership stays out of it. No Ray.""" from __future__ import annotations @@ -9,45 +9,24 @@ from yeto.syncer_profile import SyncerSemanticProfile LEGACY = "85a50d0e4fd5103b81a2295536fca0e5ad96cd50531fb1396208af8459be571e" -# Golden for the Rust side (elastic defaults 0.75 / 0.5 / 900 / 1 / 2 on the SAO profile). -ELASTIC_DEFAULT = "c677c127b1ad3e1da179863ee02b6b2e1bc232edb03b088bad62251e5d8d9ba5" # Golden for the Rust side (outer_delayed_nesterov_period=4 on the SAO profile). DELAYED_NESTEROV_4 = "2bf20f25b4d40d6d3db1d91e38d199e0347f9dace3fc481a2948fa7f1c99f126" -def _elastic(**kw): - raw = _syncer_profile() - raw.update(island_scheduling_mode="elastic", quorum_theta=0.75, carry_gamma=0.5, - soft_deadline_s=900, q_min=1, max_carry_lag=2) - raw.update(kw) - return raw - - -def test_legacy_hash_unchanged(): +def test_fixed_roster_hash_unchanged(): assert SyncerSemanticProfile.from_mapping(_syncer_profile()).sha256 == LEGACY - raw = _syncer_profile(); raw["island_scheduling_mode"] = "legacy" - assert SyncerSemanticProfile.from_mapping(raw).sha256 == LEGACY - -def test_elastic_default_hash_and_layout(): - p = SyncerSemanticProfile.from_mapping(_elastic()) - assert p.sha256 == ELASTIC_DEFAULT - tail = p.canonical_bytes()[-(33 + 1 + 22 + 1 + 7 + 8 + 8 + 8 + 4 + 4):] - assert tail.startswith(b"yeto-syncer-island-scheduling-v1\0\x16island_scheduling_mode\x07elastic") - assert SyncerSemanticProfile.from_mapping(_elastic(q_min=2)).sha256 != ELASTIC_DEFAULT - -def test_invalid_island_scheduling_profiles(): - raw = _syncer_profile(); raw["quorum_theta"] = 0.75 - with pytest.raises(ValueError): - SyncerSemanticProfile.from_mapping(raw) # parameter without elastic - raw = _elastic(); del raw["carry_gamma"] +@pytest.mark.parametrize("field, value", [ + ("island_scheduling_mode", "legacy"), ("island_scheduling_mode", "elastic"), + ("quorum_theta", 0.75), ("carry_gamma", 0.5), ("soft_deadline_s", 900), ("q_min", 1), + ("max_carry_lag", 2), +]) +def test_island_scheduling_fields_are_not_profile_fields(field, value): + raw = _syncer_profile() + raw[field] = value with pytest.raises(ValueError): SyncerSemanticProfile.from_mapping(raw) - with pytest.raises(ValueError): - SyncerSemanticProfile.from_mapping(_elastic(island_scheduling_mode="auto")) - with pytest.raises(ValueError): - SyncerSemanticProfile.from_mapping(_elastic(quorum_theta=0.0)) def test_contract_field_names_agree(): diff --git a/tests/test_rl_inter_island_elastic_client.py b/tests/test_rl_inter_island_elastic_client.py deleted file mode 100644 index 1a05584e8..000000000 --- a/tests/test_rl_inter_island_elastic_client.py +++ /dev/null @@ -1,277 +0,0 @@ -"""Island-side elastic client: frames byte-identical to syncer elastic.rs (tasks 0.14). No Ray.""" - -from __future__ import annotations - -import socket -import subprocess -import sys -import threading -import time -from types import SimpleNamespace - -import pytest -import torch - -from yeto.protocol import MSG_ERROR, read_frame, write_frame -from yeto.rl import bridge as bridge_mod -from yeto.rl.core import canonical_state -from yeto.rl.elastic_client import ( - MSG_JOIN, MSG_LEAVE, DeltaReady, DeltaTensor, ElasticBase, ElasticClientConfig, ElasticInit, - ElasticIslandClient, ElasticProtocolError, Join, JoinAck, LeaseHeartbeat, Leave, SampleIndex, - SampleVerdict, VERDICT_ACCEPT_IS, - decode, encode, hmac_key_from_env, hmac_sha256, -) - -# Produced by the syncer's own ElasticMsg::encode(b"k1") (elastic.rs at s15-interisland-rs -# 101a172c; 19 SAMPLE_INDEX and 24 SAMPLE_VERDICT re-dumped at 1be303e3), the messages of all_frames_roundtrip_and_reject_tampering_and_wrong_key), -# dumped from a scratch copy of that source; see progress.md. -RUST_GOLDEN = { - 15: "03000000000000000100000007000000000000000000000000000440abababababababababababababababababababababababababababababababab0b0000006e76696469612d68313030d76ccd494a9ee9bb15d43cdcafb727ab8d0c651b9b81c6481e80ff4084314b4c", # s17-elastic-identity: + backend identity; s19-compat: + compat_group - 16: "0300000000000000010000000400000000000000090000000000000007070707070707070707070707070707070707070707070707070707070707070193e467c96823a92f6a48b1fdc63e870e2cb17b0de0968b74790c97c7ad575810", - 17: "03000000000000000100000000d7833586bcbea7773fc51f18c8ded76ac70cb34571860c23d8ab8409dad86a27", - 18: "03000000000000000100000004000000000000000b00000000000000000000000000f43fa2b755da8b28ec797183f4a991453ba9ad27864e96abac41d01a6fcf3abcbd03", - 19: "030000000000000001000000090000000000000003000000000000000101010101010101010101010101010101010101010101010101010101010101020000006731010000007008000000000000001500000073333a2f2f6275636b65742f672e70617271756574000400000000000040000000616161616161616161616161616161616161616161616161616161616161616161616161616161616161616161616161616161616161616161616161616161610101000000000000f83f170000007965746f2e726c2e73616d706c652d696e6465782f7631295b5352ea4d70bf182ed6fa1a98b95658ae152abc6b0f56e218de674a595904", - 24: "030000000000000001120000007374616c655f77697468696e5f626f756e64010000000000000068215e6ef9d59311adef9dafaffce6c5424f9b88e9a6eb4ca254b270b5db2bfd", - 20: "03000000000000000100000009000000000000006400000000000000040000000000000022e1d4320b934424bedffe64ef9bf39a7176a2794b028031c953a164bd51904c", - 21: "03000000000000000100000002000000000000000000803f000000c0490f5332c1f0c2ed1b73ccb1d0e4e134f86b9a2b75ad0ca3df85691045b8c12c", - 22: "03000000000000000100000002000000000000000500000000000000010000000000000003000000000000000000003f0000803e000080bfbfa7704da9c65533c2208a2adf0ae3249de4ac6fef91d8d1b44da8ee328cff6e", - 23: "03000000000000000400000000000000010000000000000000004040ddcbb3df54cd41f060864a9f6390e53d4b9969546c5688e6019e6b49056c6a3e", -} -MSGS = { - 15: Join(3, 1, 7, 2.5, bytes([0xab]) * 32, "nvidia-h100"), - 16: JoinAck(3, 1, 4, 9, bytes([7]) * 32, True), - 17: Leave(3, 1, 0), - 18: LeaseHeartbeat(3, 1, 4, 11, 1.25), - 19: SampleIndex(3, 1, 9, 3, bytes([1]) * 32, "g1", "p", 8, "s3://bucket/g.parquet", 1024, - "a" * 64, True, True, 1.5, "yeto.rl.sample-index/v1"), - 24: SampleVerdict(3, VERDICT_ACCEPT_IS, "stale_within_bound", 1), - 20: DeltaReady(3, 1, 9, 100, 4), - 21: ElasticInit(3, 1, (1.0, -2.0)), - 22: DeltaTensor(3, 1, 2, 5, 1, (0.5, 0.25, -1.0)), - 23: ElasticBase(3, 4, (3.0,)), -} - - -def test_hmac_rfc4231_case_2(): - assert hmac_sha256(b"Jefe", b"what do ya want for nothing?").hex() == ( - "5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843") - - -@pytest.mark.parametrize("msg_type", sorted(MSGS)) -def test_frames_match_rust_bytes(msg_type): - t, payload = encode(MSGS[msg_type], b"k1") - assert t == msg_type and payload.hex() == RUST_GOLDEN[msg_type] - assert decode(b"k1", t, bytes.fromhex(RUST_GOLDEN[msg_type])) == MSGS[msg_type] - with pytest.raises(ElasticProtocolError): - decode(b"k2", t, payload) - bad = bytearray(payload); bad[0] ^= 1 - with pytest.raises(ElasticProtocolError): - decode(b"k1", t, bytes(bad)) - other = MSG_JOIN if t == MSG_LEAVE else MSG_LEAVE # the MAC binds the type byte - with pytest.raises(ElasticProtocolError): - decode(b"k1", other, payload) - - -def test_key_from_env(): - assert hmac_key_from_env({"YETO_ISLAND_HMAC_KEY": "abc"}) == b"abc" - with pytest.raises(ElasticProtocolError): - hmac_key_from_env({}) - - -class FakeSyncer: - """Socketpair stand-in for elastic_server.rs: JOIN->JOIN_ACK(+base), INIT->base, - each DELTA_TENSOR -> base' = base + update (version + 1).""" - - def __init__(self, key=b"k", catch_up=False, refuse=False, base=None): - self.a, self.b = socket.socketpair() - self.key, self.catch_up, self.refuse = key, catch_up, refuse - self.base, self.version = base, 0 - self.frames = [] - threading.Thread(target=self._serve, daemon=True).start() - - def _send_base(self): - write_frame(self.b, *encode(ElasticBase(0, self.version, tuple(self.base)), self.key)) - - def _serve(self): - while True: - try: - t, p = read_frame(self.b, 1 << 20) - except Exception: - return - msg = decode(self.key, t, p) - self.frames.append(msg) - if isinstance(msg, Join): - if self.refuse: - write_frame(self.b, MSG_ERROR, b"fenced: syncer_epoch 0 < 5") - continue - write_frame(self.b, *encode(JoinAck(0, msg.island_id, 1, self.version, - bytes(32), self.catch_up), self.key)) - if self.base is not None: - self._send_base() - elif isinstance(msg, ElasticInit): - if self.base is None: - self.base = list(msg.params) - self._send_base() - elif isinstance(msg, DeltaTensor): - self.base = [b + u for b, u in zip(self.base, msg.update)] - self.version += 1 - self._send_base() - - def connect(self, addr): - return self.a - - -def _client(fake, **kw): - cfg = ElasticClientConfig(("x", 0), island_id=3, capacity=2.0, lease_s=0.15, - join_timeout_s=2.0, **kw) - return ElasticIslandClient(cfg, fake.key, connect=fake.connect) - - -def test_join_heartbeat_delta_leave(): - fake = FakeSyncer(catch_up=True) - c = _client(fake) - ack = c.join() - assert ack.catch_up and c.heartbeat_period_s == pytest.approx(0.05) - time.sleep(0.25) - c.delta_ready(base_version=0, c_tokens=100, c_steps=4) - c.leave() - time.sleep(0.2) - c.close() - kinds = [type(m).__name__ for m in fake.frames] - assert kinds[0] == "Join" and fake.frames[0].capacity == 2.0 - assert kinds.count("LeaseHeartbeat") >= 2 - assert DeltaReady(0, 3, 0, 100, 4) in fake.frames - assert kinds[-1] == "Leave" # no heartbeat after LEAVE - - -def test_join_refused_fails_fast(): - c = _client(FakeSyncer(refuse=True)) - t0 = time.monotonic() - with pytest.raises(ElasticProtocolError, match="fenced"): - c.join() - assert time.monotonic() - t0 < 1.0 - c.close() - - -# rl-code-quality-cleanup 2.1: ElasticRlBridge / make_island_bridge (no production -# caller) were deleted with their tests; the elastic island path is -# yeto.rl.engine.bridges.ElasticAvgSync. - - -def test_legacy_never_imports_elastic_client(): - code = ("import sys; import yeto.rl.bridge, yeto.protocol; " - "assert 'yeto.rl.elastic_client' not in sys.modules") - subprocess.run([sys.executable, "-c", code], check=True) - - -# ----------------------------------------------------------------- 0.19 SAMPLE_INDEX wire alignment -def _entry(island="1", version=9, policy_hash="01" * 32, group="g1", has_lp=True): - from yeto.rl.engine.sample_pool import SampleIndexEntry - - return SampleIndexEntry(island, version, 3, policy_hash, group, "p", 8, "s3://bucket/g.parquet", - 1024, "a" * 64, has_lp, True, 1.5) - - -def test_sample_index_entry_roundtrip_matches_rust_frame(): - from yeto.rl.elastic_client import sample_index_from_entry, sample_index_to_entry - - msg = sample_index_from_entry(_entry(), syncer_epoch=3) - assert msg == MSGS[19] - assert encode(msg, b"k1")[1].hex() == RUST_GOLDEN[19] - assert sample_index_to_entry(msg) == _entry() - named = sample_index_from_entry(_entry(island="island-a"), syncer_epoch=3, island_ids={"island-a": 1}) - assert named == MSGS[19] - assert sample_index_to_entry(named, island_names={1: "island-a"}).island_id == "island-a" - - -def test_sample_index_conversion_rejects_bad_ids_and_hashes(): - from yeto.rl.elastic_client import sample_index_from_entry - - with pytest.raises(ElasticProtocolError, match="numeric id"): - sample_index_from_entry(_entry(island="island-a"), syncer_epoch=0) - with pytest.raises(ElasticProtocolError, match="u32"): - sample_index_from_entry(_entry(island=str(1 << 32)), syncer_epoch=0) - with pytest.raises(ElasticProtocolError, match="not hex"): - sample_index_from_entry(_entry(policy_hash="zz" * 32), syncer_epoch=0) - with pytest.raises(ElasticProtocolError, match="32 bytes"): - sample_index_from_entry(_entry(policy_hash="01" * 16), syncer_epoch=0) - - -def test_sample_verdict_names(): - assert MSGS[24].verdict_name == "ACCEPT_IS" - assert SampleVerdict(0, 9, "", 0).verdict_name == "UNKNOWN(9)" - - -def _f32_hash(values): - import hashlib - import struct as _struct - - return hashlib.sha256(_struct.pack(f"<{len(values)}f", *values)).hexdigest() - - -@pytest.mark.skipif(not __import__("os").environ.get("YETO_TEST_ELASTIC_SYNCER"), - reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") -def test_sample_index_three_verdicts_against_real_rust_syncer(tmp_path): - """One island, two outer versions: ACCEPT / ACCEPT_IS (needs IS correction) / REJECT, - then the status.json the syncer exports is parsed by the 0.18 reader.""" - import json - import os - - from yeto.rl.elastic_client import VERDICT_ACCEPT, VERDICT_REJECT - from yeto.rl.engine.island_status import read_syncer_status, scheduling_probe_from_status - - with socket.socket() as s: - s.bind(("127.0.0.1", 0)) - port = s.getsockname()[1] - key = "sample-e2e" - tape = tmp_path / "syncer-tape.jsonl" - syncer = subprocess.Popen( - [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "1", - "--total-steps", "3", "--event-tape", str(tape), - "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", - "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", - "--island-lease-s", "6", "--syncer-epoch", "2"], - env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), - stdout=open(tmp_path / "syncer.log", "w"), stderr=subprocess.STDOUT) - c = ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), 1, lease_s=6.0, syncer_epoch=2), - key.encode()) - try: - for _ in range(50): - try: - c.join() - break - except OSError: - time.sleep(0.1) - c.elastic_init([0.0] * 4) - base0 = c.wait_base(newer_than=None, timeout_s=10.0) - c.delta_tensor(base_version=base0.outer_version, c_tokens=10, c_steps=1, update=[1.0] * 4) - base1 = c.wait_base(newer_than=base0.outer_version, timeout_s=10.0) - assert base1 is not None, c.errors - h0, h1 = _f32_hash(base0.params), _f32_hash(base1.params) - verdicts = [c.sample_index(_entry(version=base1.outer_version, policy_hash=h1, group="g1")), - c.sample_index(_entry(version=base0.outer_version, policy_hash=h0, group="g0")), - c.sample_index(_entry(version=base0.outer_version, policy_hash=h1, group="bad"))] - assert [(v.verdict, v.reason, v.outer_lag) for v in verdicts] == [ - (VERDICT_ACCEPT, "same_version", 0), (VERDICT_ACCEPT_IS, "stale_within_bound", 1), - (VERDICT_REJECT, "policy_hash_mismatch", 1)] - status = None - for _ in range(50): - status = read_syncer_status(tmp_path) - if status is not None and status["sample_index_count"] == 2: - break - time.sleep(0.1) - assert status is not None and status["syncer_epoch"] == 2, status - assert status["policy_hash"] == h1 and status["outer_version"] == base1.outer_version - probe = scheduling_probe_from_status(tmp_path, island_id=1)() - assert probe["capacity"] == 1.0 and probe["syncer_epoch"] == 2 - idx = (tmp_path / "sample_index.jsonl").read_text().splitlines() - from yeto.rl.engine.sample_pool import SampleIndexEntry - - rows = [json.loads(l) for l in idx] - assert [(r["verdict"], r["entry"]["group_id"]) for r in rows] == [("ACCEPT", "g1"), ("ACCEPT_IS", "g0")] - assert SampleIndexEntry.from_json(rows[1]["entry"]) == _entry(version=base0.outer_version, - policy_hash=h0, group="g0") - finally: - c.close() - syncer.kill() - syncer.wait(5) diff --git a/tests/test_rl_inter_island_ports.py b/tests/test_rl_inter_island_ports.py deleted file mode 100644 index fa58486b6..000000000 --- a/tests/test_rl_inter_island_ports.py +++ /dev/null @@ -1,155 +0,0 @@ -"""ports engine + ElasticAvgSync (rl-inter-island-scheduling 0.15). CPU, no Ray. - -The end-to-end case runs the real Rust syncer binary when YETO_TEST_ELASTIC_SYNCER -points to one built from the syncer source with elastic support; otherwise skipped. -""" - -from __future__ import annotations - -import os -import socket -import subprocess -import time -from pathlib import Path -from types import SimpleNamespace - -import pytest -import torch - -from test_rl_engine_driver import NAME, _driver, _engine, _run_threads, _strict_config -from test_rl_inter_island_elastic_client import FakeSyncer -from yeto.rl.elastic_client import DeltaTensor, ElasticClientConfig, ElasticInit, ElasticIslandClient, Leave -from yeto.rl.engine.bridges import ElasticAvgSync, StrictAvgSync -from yeto.rl.adapters.miles import entry - - -def _sync(tmp_path, engine, client, *, learner_id=0, rounds=2): - return ElasticAvgSync(_strict_config(tmp_path, engine, learner_id=learner_id, rounds=rounds), - client=client, base_wait_s=10.0) - - -def test_ports_driver_with_elastic_sync_fake_syncer(tmp_path): - fake = FakeSyncer() - engine = _engine(torch.tensor([1.0, 3.0])) - client = ElasticIslandClient(ElasticClientConfig(("x", 0), 0, lease_s=3.0, join_timeout_s=2), - fake.key, connect=fake.connect) - result = _driver(engine, _sync(tmp_path, engine, client), tmp_path).run() - time.sleep(0.05) - inits = [m for m in fake.frames if isinstance(m, ElasticInit)] - deltas = [m for m in fake.frames if isinstance(m, DeltaTensor)] - assert inits == [ElasticInit(0, 0, (0.0, 0.0))] - assert [(d.base_version, d.update) for d in deltas] == [(0, (1.0, 3.0)), (1, (1.0, 3.0))] - assert isinstance(fake.frames[-1], Leave) - assert torch.equal(result.tensors[NAME], torch.tensor([[2.0, 6.0]])) - - -def test_entry_selects_by_mode(monkeypatch): - import yeto.rl.engine.bridges as b - monkeypatch.setattr(b, "StrictIslandProgress", lambda args: None) - base = dict(num_rollout=2, yeto_rl_bridge_config=SimpleNamespace()) - sync, _ = entry.build_sync(SimpleNamespace(**base), yeto_policy_sync=True) - assert type(sync) is StrictAvgSync # legacy: attribute absent - sync, _ = entry.build_sync(SimpleNamespace(**base, yeto_rl_island_scheduling="elastic", - yeto_rl_syncer_epoch=7), yeto_policy_sync=True) - assert type(sync) is ElasticAvgSync and sync.syncer_epoch == 7 - - -def _free_port(): - s = socket.socket(); s.bind(("127.0.0.1", 0)); p = s.getsockname()[1]; s.close(); return p - - -@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), - reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") -def test_two_ports_islands_against_real_rust_syncer(tmp_path): - port, key = _free_port(), "ports-e2e" - out = Path(os.environ.get("YETO_TEST_ELASTIC_OUT", tmp_path)) - out.mkdir(parents=True, exist_ok=True) - for old in out.glob("*.jsonl"): # the syncer tape appends - old.unlink() - syncer = subprocess.Popen( - [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", - "--total-steps", "2", "--event-tape", str(out / "syncer-tape.jsonl"), - "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", - "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", - "--island-lease-s", "6", "--syncer-epoch", "5"], - env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), - stdout=open(out / "syncer.log", "w"), stderr=subprocess.STDOUT) - try: - time.sleep(1.0) - engines = (_engine(torch.tensor([1.0, 3.0])), _engine(torch.tensor([3.0, 5.0]))) - drivers = [] - for i, e in enumerate(engines): - c = ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), i, lease_s=6.0, - syncer_epoch=5), key.encode()) - cfg = _strict_config(out, e, learner_id=i, rounds=2, tape=f"island-{i}.jsonl") - drivers.append(_driver(e, ElasticAvgSync(cfg, client=c, base_wait_s=30.0), out, - name=f"events-{i}.jsonl", learner_id=i)) - results, errors = _run_threads(drivers) - assert errors == {} - assert torch.equal(results[0].tensors[NAME], results[1].tensors[NAME]) - assert syncer.wait(15) == 0 - kinds = [l.split('"kind":"')[1].split('"')[0] - for l in (out / "syncer-tape.jsonl").read_text().splitlines()] - assert kinds.count("outer_step") == 2 and kinds.count("pool_join") == 2 - finally: - if syncer.poll() is None: - syncer.kill() - - -def test_elastic_emits_rl_local_round_like_strict(tmp_path): - """0.24: one rl_local_round per elastic round with reward/kl/update norm.""" - import inspect - import json - import math - - from yeto.rl import bridge as strict_bridge - - fake = FakeSyncer() - engine = _engine(torch.tensor([1.0, 3.0])) - client = ElasticIslandClient(ElasticClientConfig(("x", 0), 0, lease_s=3.0, join_timeout_s=2), - fake.key, connect=fake.connect) - _driver(engine, _sync(tmp_path, engine, client), tmp_path).run() - rounds = [json.loads(l) for l in (tmp_path / "events.jsonl").read_text().splitlines() - if '"rl_local_round"' in l] - assert len(rounds) == 2 - assert [r["rl/global_policy_version"] for r in rounds] == [0, 1] - assert all(math.isclose(r["rl/local_delta_norm"], math.sqrt(10.0), rel_tol=1e-6) for r in rounds) - for key in ("rl/reward_mean", "rl/current_vs_rollout_kl", "rl/local_delta_norm", "reward_mean"): - assert key in rounds[0] - strict_keys = set(__import__("re").findall(r'"(rl/[a-z_0-9]+|sync/bytes_sent)"', - inspect.getsource(strict_bridge.StrictRlBridge._record_submission))) - assert strict_keys <= set(rounds[0]) - - -def test_debug_delay_switch(tmp_path, monkeypatch): - """0.25: --rl-elastic-debug-delay-s -> env on the island -> sleep before each delta.""" - import json - - from test_rl_launcher import _args - from yeto import launcher - from yeto.rl.elastic_client import parse_elastic_debug_delay - - assert parse_elastic_debug_delay("0:90, 1:5") == {0: 90.0, 1: 5.0} - assert parse_elastic_debug_delay("30") == {-1: 30.0} - with pytest.raises(ValueError): - parse_elastic_debug_delay("0:-1") - assert launcher.elastic_debug_delay_env(_args()) == {} - assert launcher.elastic_debug_delay_env(_args(("--rl-island-scheduling", "elastic", - "--rl-elastic-debug-delay-s", "0:90"))) == { - "YETO_RL_ELASTIC_DEBUG_DELAY": "0:90"} - with pytest.raises(ValueError, match="needs --rl-island-scheduling elastic"): - launcher.elastic_debug_delay_env(_args(("--rl-elastic-debug-delay-s", "0:90"))) - - monkeypatch.setenv("YETO_RL_ELASTIC_DEBUG_DELAY", "0:7.5,1:2") - fake = FakeSyncer() - engine = _engine(torch.tensor([1.0, 3.0])) - client = ElasticIslandClient(ElasticClientConfig(("x", 0), 0, lease_s=3.0, join_timeout_s=2), - fake.key, connect=fake.connect) - sync = _sync(tmp_path, engine, client) - slept = [] - sync._sleep = slept.append - _driver(engine, sync, tmp_path).run() - assert slept == [7.5, 7.5] - phases = [json.loads(l) for l in (tmp_path / "events.jsonl").read_text().splitlines() - if '"elastic_debug_delay"' in l] - assert [p["delay_s"] for p in phases] == [7.5, 7.5] diff --git a/tests/test_rl_inter_island_rust_golden.py b/tests/test_rl_inter_island_rust_golden.py index 28b80ad7c..078342035 100644 --- a/tests/test_rl_inter_island_rust_golden.py +++ b/tests/test_rl_inter_island_rust_golden.py @@ -1,6 +1,6 @@ -"""Golden fixture: Python CrossIslandLedger (elastic) drives a fixed scenario; -the Rust ElasticCoordinator replays the same ops and must match -(syncer/src/elastic.rs::tests::python_ledger_golden_tape_matches). +"""Golden fixture: Python ``island_ledger.merge_weight`` (tokens^2 / steps) is +locked to the Rust ``merge::learner_weight`` +(syncer/src/merge.rs::tests::python_merge_weight_golden_matches). Regenerate: python tests/test_rl_inter_island_rust_golden.py --write """ @@ -9,133 +9,12 @@ import sys from pathlib import Path -import re - from yeto.rl.engine import island_ledger -from yeto.rl.engine.island_ledger import CrossIslandLedger, DeltaEntry REPO = Path(__file__).resolve().parents[1] -FIXTURE = REPO / "syncer/tests/fixtures/elastic_ledger_golden.txt" -# rl-code-quality-cleanup 2.5: Python merge_weight vs Rust merge::learner_weight. WEIGHT_FIXTURE = REPO / "syncer/tests/fixtures/merge_weight_golden.txt" WEIGHT_CASES = ((0, 0), (0, 4), (7, 0), (1, 1), (37, 4), (91, 4), (50, 5), (4096, 3), (1 << 20, 7), (123457, 13), (2**32 - 1, 1), (5, 2**32 + 9), (10**9, 999983)) -# Plain-text format (the syncer has no JSON dependency): one line per op, -# " => "; floats are Python repr (round-trip exact). -PARAMS = {"quorum_theta": 0.75, "carry_gamma": 0.5, "q_min": 1, "max_carry_lag": 2} -LEASE_S = 30.0 - -# Island ids are u32 on the wire; Python uses their decimal string. -OPS = [ - {"op": "join", "island": 1, "capacity": 3.0, "now": 0.0}, - {"op": "join", "island": 2, "capacity": 1.0, "now": 0.0}, - {"op": "join", "island": 3, "capacity": 1.0, "now": 0.0}, - # v0: big island alone is 3/5 < 0.75 -> wait; +island 2 -> 4/5 -> step. - {"op": "submit", "island": 1, "base": 0, "c_tokens": 100, "c_steps": 4}, - {"op": "advance", "timed_out": False}, - {"op": "submit", "island": 2, "base": 0, "c_tokens": 60, "c_steps": 3}, - {"op": "advance", "timed_out": False}, - # island 3 is late on base 0 (lag 1) -> carried with gamma^1. - {"op": "submit", "island": 3, "base": 0, "c_tokens": 50, "c_steps": 5}, - {"op": "submit", "island": 1, "base": 1, "c_tokens": 120, "c_steps": 4}, - {"op": "advance", "timed_out": True}, - # new island joins at v2 (catch-up -> weight 0 in its first round). - {"op": "join", "island": 4, "capacity": 1.0, "now": 10.0}, - {"op": "submit", "island": 4, "base": 2, "c_tokens": 80, "c_steps": 2}, - {"op": "submit", "island": 1, "base": 2, "c_tokens": 90, "c_steps": 3}, - {"op": "submit", "island": 2, "base": 2, "c_tokens": 70, "c_steps": 7}, - {"op": "advance", "timed_out": False}, - # island 3 submits base 0 at v3 (lag 3 > 2) -> rejected. - {"op": "submit", "island": 3, "base": 0, "c_tokens": 50, "c_steps": 5}, - # island 2 submits then leaves -> dropped uncommitted. - {"op": "submit", "island": 2, "base": 3, "c_tokens": 40, "c_steps": 2}, - {"op": "leave", "island": 2}, - {"op": "advance", "timed_out": True}, - {"op": "submit", "island": 4, "base": 3, "c_tokens": 30, "c_steps": 3}, - {"op": "advance", "timed_out": True}, - {"op": "submit", "island": 1, "base": 3, "c_tokens": 90, "c_steps": 3}, - {"op": "advance", "timed_out": False}, -] - - -def _parse_key(k: str, base: int) -> tuple[int, int]: - if "@" in k: - i, b = k.split("@") - return int(i), int(b) - return int(k), base - - -def generate() -> dict: - led = CrossIslandLedger(theta=PARAMS["quorum_theta"], quorum_min=PARAMS["q_min"], - gamma=PARAMS["carry_gamma"], max_carry_lag=PARAMS["max_carry_lag"], - lease_s=LEASE_S, mode="elastic") - expected = [] - for op in OPS: - k = op["op"] - if k == "join": - ev = led.join(str(op["island"]), now=op["now"], capacity=op["capacity"]) - expected.append({"kind": "join", "catch_up": ev["catch_up"], - "base_version": ev["base_version"]}) - elif k == "leave": - ev = led.leave(str(op["island"]), reason="left") - d = ev["dropped_uncommitted"] - expected.append({"kind": "leave", "dropped": None if d is None else - [d["outer_version"], d["c_tokens"], d["c_steps"]]}) - elif k == "submit": - ev = led.submit(DeltaEntry(str(op["island"]), op["base"], 0, - led.published[op["base"]], op["c_tokens"], op["c_steps"])) - out = {"kind": ev["kind"]} - if ev["kind"] == "delta_carried_over": - out["lag"] = ev["lag"] - if ev["kind"] == "delta_rejected": - out["reason"] = ev["reason"] - expected.append(out) - else: - base = led.outer_version - ev = led.try_advance(timed_out=op["timed_out"]) - if ev is None: - expected.append({"kind": "no_step", "outer_version": led.outer_version}) - else: - weights = sorted([*_parse_key(key, base), w] for key, w in ev["weights"].items()) - expected.append({"kind": "outer_step", "base_version": ev["base_version"], - "arrived": ev["arrived"], "cap_arrived": ev["cap_arrived"], - "cap_total": ev["cap_total"], "weights": weights, - "absent": sorted(int(a) for a in ev["absent"])}) - return {"params": PARAMS, "lease_s": LEASE_S, "ops": OPS, "expected": expected} - - -def render(g: dict) -> str: - p = g["params"] - lines = [f"# generated by tests/test_rl_inter_island_rust_golden.py; do not edit", - f"params {p['quorum_theta']!r} {p['carry_gamma']!r} {p['q_min']} " - f"{p['max_carry_lag']} {g['lease_s']!r}"] - for op, e in zip(g["ops"], g["expected"]): - k = op["op"] - if k == "join": - lhs = f"join {op['island']} {op['capacity']!r} {op['now']!r}" - rhs = f"join {str(e['catch_up']).lower()} {e['base_version']}" - elif k == "leave": - lhs = f"leave {op['island']}" - rhs = "leave " + ("none" if e["dropped"] is None else " ".join(map(str, e["dropped"]))) - elif k == "submit": - lhs = f"submit {op['island']} {op['base']} {op['c_tokens']} {op['c_steps']}" - rhs = e["kind"] + (f" {e['lag']}" if "lag" in e else "") + ( - f" {e['reason']}" if "reason" in e else "") - else: - lhs = f"advance {int(op['timed_out'])}" - if e["kind"] == "no_step": - rhs = f"no_step {e['outer_version']}" - else: - w = ",".join(f"{i}:{b}:{x!r}" for i, b, x in e["weights"]) - a = ",".join(map(str, e["absent"])) or "-" - rhs = (f"outer_step {e['base_version']} {e['arrived']} {e['cap_arrived']!r} " - f"{e['cap_total']!r} {a} {w}") - lines.append(f"{lhs} => {rhs}") - return "\n".join(lines) + "\n" - - -def test_fixture_is_current(): - assert FIXTURE.read_text() == render(generate()) def render_weights() -> str: @@ -151,33 +30,5 @@ def test_merge_weight_fixture_is_current(): assert WEIGHT_FIXTURE.read_text() == render_weights() -def _rust_const(name: str) -> int: - text = (REPO / "syncer/src/elastic.rs").read_text() - m = re.search(rf"pub const {name}: u\d+ = (\d+);", text) - assert m, f"{name} not found in syncer/src/elastic.rs" - return int(m.group(1)) - - -def test_lag_bounds_match_rust(): - pairs = { - "MAX_SAMPLE_OUTER_LAG vs StalenessPolicy.max_outer_lag": - (_rust_const("MAX_SAMPLE_OUTER_LAG"), island_ledger.StalenessPolicy().max_outer_lag), - "DEFAULT_MAX_CARRY_LAG vs island_ledger.DEFAULT_MAX_CARRY_LAG": - (_rust_const("DEFAULT_MAX_CARRY_LAG"), island_ledger.DEFAULT_MAX_CARRY_LAG), - "golden PARAMS max_carry_lag vs island_ledger.DEFAULT_MAX_CARRY_LAG": - (PARAMS["max_carry_lag"], island_ledger.DEFAULT_MAX_CARRY_LAG), - } - for what, (rust, python) in pairs.items(): - assert rust == python, f"{what}: Rust {rust} != Python {python}" - - -def test_fixture_has_a_lag_over_the_bound_rejected(): - lines = FIXTURE.read_text().splitlines() - assert "submit 3 0 50 5 => delta_rejected carry_lag_exceeded" in lines - # island 3 submits base 0 after advance to v3: lag 3 > max_carry_lag 2. - assert island_ledger.DEFAULT_MAX_CARRY_LAG == 2 - - if __name__ == "__main__" and "--write" in sys.argv: - FIXTURE.write_text(render(generate())) WEIGHT_FIXTURE.write_text(render_weights()) diff --git a/tests/test_rl_island_contract_v3.py b/tests/test_rl_island_contract_v3.py index d0b946d91..15b0fb078 100644 --- a/tests/test_rl_island_contract_v3.py +++ b/tests/test_rl_island_contract_v3.py @@ -17,7 +17,6 @@ import pytest from yeto.rl.adapters.miles.identity import backend_identity -from yeto.rl.elastic_client import ElasticClientConfig, ElasticIslandClient, ElasticProtocolError from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec, PluginRef from yeto.rl.engine.backend_identity import ( config_island_contract_sha256, @@ -90,44 +89,50 @@ def _free_port() -> int: @pytest.mark.parametrize("bad", ["other_pack", "v2_island"]) -def test_real_syncer_refuses_a_join_with_another_contract(tmp_path, syncer_binary, bad): # noqa: F811 - """Island 1 joins with contract v3 (pack A) and seeds the base. Island 2 then - declares pack B, or the v2 contract of the same run; the syncer refuses the - JOIN and names both contracts.""" +def test_real_syncer_refuses_a_pool_island_with_another_contract(tmp_path, syncer_binary, bad): # noqa: F811 + """Island 1 joins an island pool with contract v3 (pack A) and seeds the base. + Island 2 then declares pack B, or the v2 contract of the same run; the syncer + refuses its HELLO (only that connection) and names both session contracts.""" + import torch + + from yeto.rl.bridge import _session_contract + from yeto.rl.core import CanonicalTensorSpec, build_avg_layout + from yeto.rl.pool_client import PoolClientConfig, PoolIslandClient + port, key = _free_port(), "contract-v3" log = open(tmp_path / "syncer.log", "w") syncer = subprocess.Popen( - [str(syncer_binary), "--port", str(port), "--learners", "2", + [str(syncer_binary), "--port", str(port), "--learners", "1", "--total-steps", "1", "--event-tape", str(tmp_path / "tape.jsonl"), - "--island-scheduling-mode", "elastic", "--quorum-theta", "1.0", "--carry-gamma", "0.5", - "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", - "--island-lease-s", "6", "--syncer-epoch", "0", "--outer-lr", "1.0", - "--outer-momentum", "0.0"], + "--max-learners", "2", "--quorum", "1", "--quorum-theta", "1.0", "--max-base-lag", "2", + "--island-lease-s", "6", "--outer-lr", "1.0", "--outer-momentum", "0.0"], env={"YETO_ISLAND_HMAC_KEY": key, "PATH": "/usr/bin:/bin"}, stdout=log, stderr=subprocess.STDOUT) good = island_contract_sha256(IDENTITY, LR, SPEC.semantic_sha256(), PACK_A) other = (island_contract_sha256(IDENTITY, LR, SPEC.semantic_sha256(), PACK_B) if bad == "other_pack" else _v2_contract(IDENTITY, LR, SPEC.sha256())) + layout = build_avg_layout((CanonicalTensorSpec( + "base_model.model.layer.lora_A.weight", (2, 2), "float32", 4),)) - def client(island: int, contract: str) -> ElasticIslandClient: - return ElasticIslandClient(ElasticClientConfig( - ("127.0.0.1", port), island, lease_s=6.0, join_timeout_s=5.0, - backend_identity_sha256=contract), key.encode()) + def client(island: int, contract: str) -> PoolIslandClient: + return PoolIslandClient(PoolClientConfig( + ("127.0.0.1", port), island, lease_s=6.0, + session_contract_hash=_session_contract(layout, contract)), layout) - first, second = client(1, good), client(2, other) + import os + + os.environ["YETO_ISLAND_HMAC_KEY"] = key + first, second = client(0, good), client(1, other) try: - for _ in range(100): - try: - first.join() - break - except OSError: - time.sleep(0.1) - first.elastic_init([0.0] * 4) - first.wait_base(newer_than=None, timeout_s=10.0) - with pytest.raises(ElasticProtocolError, match="backend identity mismatch") as exc: + first.join() + first.elastic_init(torch.zeros(4)) + assert first.wait_base(newer_than=None, timeout_s=10.0) is not None + with pytest.raises(Exception, match="contract") as exc: second.join() - assert good in str(exc.value) and other in str(exc.value) + second.wait_base(newer_than=None, timeout_s=10.0) + assert _session_contract(layout, other).hex() in str(exc.value) finally: + os.environ.pop("YETO_ISLAND_HMAC_KEY", None) for c in (first, second): try: c.close() diff --git a/tests/test_rl_elastic_leave_reason.py b/tests/test_rl_island_leave_reason.py similarity index 72% rename from tests/test_rl_elastic_leave_reason.py rename to tests/test_rl_island_leave_reason.py index d380dbe3d..df0ab1ee8 100644 --- a/tests/test_rl_elastic_leave_reason.py +++ b/tests/test_rl_island_leave_reason.py @@ -1,7 +1,7 @@ """rl-island-graceful-leave: LEAVE reason codes (design D7). Parity with the Rust syncer, the exception -> code table, and the driver -closing an elastic bridge with the code of its departure.""" +closing an island-pool session with the code of its departure.""" from __future__ import annotations @@ -12,10 +12,10 @@ import pytest import torch -from test_rl_elastic_lease_rejoin import MultiConnSyncer, _client, _wait +from pool_fakes import FakePoolSyncer from test_rl_engine_driver import _driver, _engine, _strict_config -from yeto.rl.elastic_client import ElasticJoinRefused, ElasticRejoinError, Leave, LeaveReason -from yeto.rl.engine.bridges import ElasticAvgSync +from yeto.rl.engine.bridges import PoolAvgSync +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine.capabilities import CapabilityMismatch from yeto.rl.engine.controller import RecoveryRequired, TransactionFailed from yeto.rl.engine.cut import CutError @@ -30,9 +30,10 @@ def test_reason_codes_match_the_rust_syncer(): - src = (REPO / "syncer/src/elastic.rs").read_text() - rust = {m[0]: int(m[1]) for m in re.findall(r"pub const LEAVE_REASON_(\w+): u8 = (\d+);", src)} - assert rust == {r.name: r.value for r in LeaveReason} + src = (REPO / "syncer/src/protocol.rs").read_text() + block = re.search(r"LEAVE_REASON_NAMES: \[&str; \d+\] = \[(.*?)\];", src, re.S).group(1) + names = re.findall(r'"(\w+)"', block) + assert names == [r.name.lower() for r in sorted(LeaveReason, key=int)] @pytest.mark.parametrize("error, reason", [ @@ -43,7 +44,6 @@ def test_reason_codes_match_the_rust_syncer(): (RecoveryRequired("x"), LeaveReason.INFRA_FAILED), (RebuildNotStarted("x"), LeaveReason.INFRA_FAILED), (TransactionFailed.__new__(TransactionFailed), LeaveReason.INFRA_FAILED), - (ElasticRejoinError("x"), LeaveReason.INFRA_FAILED), (CutError("x"), LeaveReason.INFRA_FAILED), (EpochConflict("x"), LeaveReason.INFRA_FAILED), (RuntimeError("CUDA out of memory"), LeaveReason.INFRA_FAILED), # unmapped: D7 Q2 @@ -52,7 +52,6 @@ def test_reason_codes_match_the_rust_syncer(): (PolicyIdentityError("x"), LeaveReason.TRAINING_FAILED), (RoundFailedError("x"), LeaveReason.TRAINING_FAILED), (CapabilityMismatch("x"), LeaveReason.CONFIG_MISMATCH), - (ElasticJoinRefused("x"), LeaveReason.CONFIG_MISMATCH), ]) def test_exception_to_reason(error, reason): assert leave_reason_for(error) is reason @@ -71,22 +70,12 @@ def test_an_island_never_sends_lease_expired_or_unknown_codes(): assert sendable(LeaveReason.STOPPED) is LeaveReason.STOPPED -def _elastic(tmp_path, fake, *, rounds=2): - """The fake syncer cannot decode LEAVE (island-side decoder), so the - client's sent LEAVE frames are recorded on ``fake.leaves`` instead.""" +def _pool(tmp_path, fake, *, rounds=2): engine = _engine(torch.tensor([1.0, 3.0])) - c = _client(fake, []) - fake.leaves = [] - send = c._send - - def recording_send(msg): - send(msg) - if isinstance(msg, Leave): - fake.leaves.append(msg.reason) - c._send = recording_send - sync = ElasticAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=rounds), client=c, - base_wait_s=5.0) - return engine, sync, c + client = fake.client(0) + sync = PoolAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=rounds), client=client, + base_wait_s=5.0) + return engine, sync, client def _leaves(fake): @@ -98,8 +87,8 @@ def _tape(tmp_path): def test_a_finished_island_leaves_with_finished(tmp_path): - fake = MultiConnSyncer() - engine, sync, _ = _elastic(tmp_path, fake) + fake = FakePoolSyncer() + engine, sync, _ = _pool(tmp_path, fake) _driver(engine, sync, tmp_path).run() assert _leaves(fake) == [LeaveReason.FINISHED] assert not [e for e in _tape(tmp_path) if e["event"] == "island_leave"] # golden tapes unchanged @@ -113,8 +102,8 @@ def test_a_finished_island_leaves_with_finished(tmp_path): (RuntimeError("ray actor died"), LeaveReason.INFRA_FAILED), ]) def test_a_failing_island_leaves_with_the_reason_of_its_error(tmp_path, error, reason): - fake = MultiConnSyncer() - engine, sync, _ = _elastic(tmp_path, fake) + fake = FakePoolSyncer() + engine, sync, _ = _pool(tmp_path, fake) def train_step(batch): raise error @@ -127,8 +116,8 @@ def train_step(batch): def test_a_failed_leave_is_reported_and_the_original_error_still_propagates(tmp_path): - fake = MultiConnSyncer() - engine, sync, client = _elastic(tmp_path, fake) + fake = FakePoolSyncer() + engine, sync, client = _pool(tmp_path, fake) def train_step(batch): def broken(reason=0): @@ -144,8 +133,8 @@ def broken(reason=0): def test_a_reclaim_leave_is_sent_once_and_not_overwritten_by_close(tmp_path): - fake = MultiConnSyncer() - engine, sync, _ = _elastic(tmp_path, fake) + fake = FakePoolSyncer() + engine, sync, _ = _pool(tmp_path, fake) def train_step(batch): assert sync.leave(LeaveReason.PREEMPTED) # reclaim listener thread in production @@ -181,25 +170,25 @@ def test_the_miles_publication_error_is_a_training_failure(): def test_the_bridge_gives_the_client_the_syncer_lease(tmp_path): """The island heartbeats every lease/3 of the syncer's --island-lease-s, not the client's own default.""" + from yeto.rl.pool_client import DEFAULT_ISLAND_LEASE_S + engine = _engine(torch.tensor([1.0, 3.0])) config = _strict_config(tmp_path, engine, learner_id=0, rounds=2) - assert ElasticAvgSync(config, lease_s=45.0)._client_config().lease_s == 45.0 - from yeto.rl.elastic_client import DEFAULT_ISLAND_LEASE_S - - assert ElasticAvgSync(config)._client_config().lease_s == DEFAULT_ISLAND_LEASE_S + assert PoolAvgSync(config, lease_s=45.0).lease_s == 45.0 + assert PoolAvgSync(config).lease_s == DEFAULT_ISLAND_LEASE_S -def test_the_launcher_passes_the_lease_to_every_elastic_island(monkeypatch): +def test_the_launcher_passes_the_lease_to_every_pool_island(monkeypatch): from test_rl_launcher import _args from yeto import launcher from yeto.rl.adapters.miles.island_entry import parse_args monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") - flags = launcher._ports_infra_flags(_args(("--rl-island-scheduling", "elastic", - "--rl-island-lease-s", "90")))[1] - assert " --rl-island-lease-s 90.0" in flags - default = launcher._ports_infra_flags(_args(("--rl-island-scheduling", "elastic")))[1] - assert f" --rl-island-lease-s {launcher.DEFAULT_ISLAND_LEASE_S!r}" in default + flags = launcher._ports_infra_flags(_args(("--rl-sync-island-lease-s", "90")))[1] + assert " --rl-island-pool" in flags and " --rl-sync-island-lease-s 90.0" in flags + default = launcher._ports_infra_flags(_args(("--rl-sync-quorum-theta", "0.75")))[1] + assert f" --rl-sync-island-lease-s {launcher.DEFAULT_ISLAND_LEASE_S!r}" in default + assert " --rl-island-pool" not in launcher._ports_infra_flags(_args(()))[1] required = ["--model", "m", "--model-revision", "r", "--data", "d", "--learner-id", "0", "--reward-function", "f", "--reward-sha256", "s", "--source-sha256", "s", "--global-rounds", "1", "--groups-per-round", "1", "--samples-per-group", "1", @@ -208,33 +197,14 @@ def test_the_launcher_passes_the_lease_to_every_elastic_island(monkeypatch): "--event-tape", "t", "--actor-num-nodes", "1", "--actor-num-gpus-per-node", "1", "--inner-lr", "1e-5", "--seq-len", "1", "--seed", "0", "--miles-root", "m", "--lora-r", "8", "--lora-targets", "attention", "--rl-single-island-no-sync"] - assert parse_args(required + ["--rl-island-lease-s", "90"]).rl_island_lease_s == 90.0 - - -class _RefusingSyncer(MultiConnSyncer): - """Answers every JOIN with MSG_ERROR, like elastic_server.rs on a compat_group - or backend identity mismatch.""" - - def _serve(self, sock, idx): - from yeto.protocol import MSG_ERROR, read_frame, write_frame - - read_frame(sock, 1 << 20) - write_frame(sock, MSG_ERROR, b"compat_group mismatch, JOIN refused") - - -def test_a_refused_first_join_raises_elastic_join_refused(): - c = _client(_RefusingSyncer(), []) - with pytest.raises(ElasticJoinRefused, match="compat_group mismatch"): - c.join() - c.close() + assert parse_args(required + ["--rl-sync-island-lease-s", "90"]).rl_sync_island_lease_s == 90.0 def test_concurrent_leaves_send_one_frame(tmp_path): import threading - fake = MultiConnSyncer() - engine, sync, client = _elastic(tmp_path, fake) - client.join() + fake = FakePoolSyncer() + engine, sync, client = _pool(tmp_path, fake) barrier = threading.Barrier(8) results = [] @@ -251,7 +221,7 @@ def leave(reason): sync.close() -def test_the_verl_elastic_island_command_carries_the_lease(monkeypatch): +def test_the_verl_pool_island_command_carries_the_lease(monkeypatch): """verl islands are built by the same launcher task as Miles islands (make_miles_island_task with the verl ISLAND_ENTRY_MODULE).""" import sys @@ -264,11 +234,11 @@ def test_the_verl_elastic_island_command_carries_the_lease(monkeypatch): monkeypatch.setitem(sys.modules, "sky", types.SimpleNamespace( Task=_Task, Resources=_Resources, Storage=_Storage, StorageMode=_StorageMode)) monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") - args = _args(("--rl-backend", "verl", "--rl-engine", "ports", "--rl-island-scheduling", "elastic", - "--rl-island-lease-s", "90")) + args = _args(("--rl-backend", "verl", "--rl-engine", "ports", "--rl-sync-island-lease-s", "90", + "--rl-sync-max-base-lag", "1")) args.model_revision, args.data_revision = "a" * 40, "b" * 40 args.source_sha256, args.reward_sha256 = "c" * 64, "d" * 64 _prepare_rl_args(args) run = make_miles_island_task(args, parse_gpu_spec(args.gpu)[0], 0, 1, "127.0.0.1:1").run assert "yeto.rl.adapters.verl.island_entry" in run - assert " --rl-island-lease-s 90.0" in run + assert " --rl-sync-island-lease-s 90.0" in run and " --rl-island-pool" in run diff --git a/tests/test_rl_island_pool_e2e.py b/tests/test_rl_island_pool_e2e.py new file mode 100644 index 000000000..433c03c32 --- /dev/null +++ b/tests/test_rl_island_pool_e2e.py @@ -0,0 +1,223 @@ +"""Island pool end to end: PoolAvgSync islands (fake engines, IslandDriver) on the +real fragment-round syncer with membership options. CPU, no Ray. + +Covers what the retired whole-vector syncer was tested for: a run to the final +cut with pool records and status.json, a late island carried with a discount, +a stale delta refused, a lease expiry followed by an automatic re-JOIN, the +final grace window, and the LEAVE reason on the syncer tape. +""" + +from __future__ import annotations + +import json +import os +import socket +import subprocess +import threading +import time +from dataclasses import replace +from pathlib import Path + +import pytest +import torch + +from test_rl_engine_driver import _driver, _engine, _strict_config +from yeto.rl.engine.bridges import PoolAvgSync, RoundCutHooks +from yeto.rl.pool_client import LeaveReason + +ROOT = Path(__file__).resolve().parent.parent +KEY = "pool-test-key" + + +@pytest.fixture(scope="module") +def syncer_binary() -> Path: + subprocess.run(["cargo", "build", "-q"], cwd=ROOT / "syncer", check=True) + binary = ROOT / "syncer/target/debug/yeto-syncer" + assert binary.is_file() + return binary + + +@pytest.fixture(autouse=True) +def _island_key(monkeypatch): + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", KEY) + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +class Syncer: + def __init__(self, binary: Path, tmp_path: Path, *, learners: int, rounds: int, extra=()): + self.port = _free_port() + self.tape = tmp_path / "syncer" / "tape.jsonl" + self.tape.parent.mkdir() + argv = [str(binary), "--port", str(self.port), "--learners", str(learners), + "--total-steps", str(rounds), "--pipeline", "1", "--grace-ms", "0", + "--sync-interval-steps", "0", "--delta-correction", "none", + "--learner-weight", "tokens2-over-steps", "--event-tape", str(self.tape), + "--quorum-timeout-s", "30", "--final-ack-timeout-s", "30", *extra] + self.log = (tmp_path / "syncer.log").open("w") + self.proc = subprocess.Popen(argv, stdout=self.log, stderr=subprocess.STDOUT, + env={**os.environ, "YETO_ISLAND_HMAC_KEY": KEY}) + + def wait(self, timeout=60.0) -> int: + try: + return self.proc.wait(timeout) + finally: + self.log.close() + + def kill(self): + if self.proc.poll() is None: + self.proc.kill() + self.proc.wait() + + def records(self, kind=None): + rows = [json.loads(line) for line in self.tape.read_text().splitlines() if line.strip()] + rows = [r for r in rows if "kind" in r] + return rows if kind is None else [r for r in rows if r["kind"] == kind] + + +def _island(tmp_path, syncer, learner_id, rounds, *, delta=1.0, lease_s=30.0): + engine = _engine(torch.tensor([delta, 2 * delta])) + config = replace(_strict_config(tmp_path, engine, learner_id=learner_id, rounds=rounds, + tape=f"bridge-{learner_id}.jsonl"), + syncer_addr=("127.0.0.1", syncer.port), wan_streams=1) + sync = PoolAvgSync(config, lease_s=lease_s, base_wait_s=60.0) + driver = _driver(engine, sync, tmp_path, name=f"island-{learner_id}.jsonl", learner_id=learner_id) + return sync, driver + + +def _run(driver, errors): + def target(): + try: + driver.run() + except BaseException as error: # noqa: BLE001 - reported by the test + errors.append(error) + thread = threading.Thread(target=target, daemon=True) + thread.start() + return thread + + +def _tape(tmp_path, learner_id): + return [json.loads(line) for line in (tmp_path / f"island-{learner_id}.jsonl").read_text().splitlines()] + + +def _run_pool(tmp_path, syncer, islands): + errors: list[BaseException] = [] + threads = [_run(driver, errors) for _sync, driver in islands] + for thread in threads: + thread.join(120) + assert not any(thread.is_alive() for thread in threads), "an island did not finish" + assert errors == [] + return syncer.wait() + + +def test_two_islands_run_to_the_final_cut_with_pool_records(syncer_binary, tmp_path): + rounds = 3 + syncer = Syncer(syncer_binary, tmp_path, learners=2, rounds=rounds, + extra=["--quorum", "1", "--quorum-theta", "1.0", "--max-base-lag", "2", + "--island-lease-s", "30", "--final-grace-s", "20"]) + try: + islands = [_island(tmp_path, syncer, i, rounds) for i in range(2)] + assert _run_pool(tmp_path, syncer, islands) == 0 + finally: + syncer.kill() + for sync, _driver_ in islands: + assert sync.base_version == rounds # each island applied the final cut + steps = syncer.records("outer_step") + assert [s["step"] for s in steps] == [1, 2, 3] + assert all(s["arrived"] == 2 and not s["timed_out"] and s["absent"] == [] for s in steps) + assert all(abs(sum(w for _i, _b, w in s["weights"]) - 1.0) < 1e-9 for s in steps) + assert {r["island_id"] for r in syncer.records("pool_join")} == {0, 1} + assert len(syncer.records("delta_accepted")) == 2 * rounds + status = json.loads((syncer.tape.parent / "status.json").read_text()) + assert status["schema"] == "yeto.syncer.elastic-status/v1" + assert status["state"] == "finished" and status["outer_version"] == rounds + for learner_id in range(2): + tape = _tape(tmp_path, learner_id) + assert [e for e in tape if e["event"] == "rl_sync_timing"][-1]["outer_sync"] == "elastic" + assert sum(e["event"] == "rl_local_round" for e in tape) == rounds + + +def test_late_island_is_carried_with_a_discount_after_the_soft_deadline(syncer_binary, tmp_path): + rounds = 3 + syncer = Syncer(syncer_binary, tmp_path, learners=2, rounds=rounds, + extra=["--quorum", "2", "--max-base-lag", "2", "--soft-deadline-s", "1", + "--carry-gamma", "0.5", "--island-lease-s", "60", "--final-grace-s", "5"]) + try: + islands = [_island(tmp_path, syncer, i, rounds) for i in range(2)] + slow = islands[1][0] + slow.debug_delay_s = lambda: 2.5 # island 1 misses every soft deadline + assert _run_pool(tmp_path, syncer, islands) == 0 + finally: + syncer.kill() + steps = syncer.records("outer_step") + assert any(s["timed_out"] and s["absent"] == [1] for s in steps), steps + carried = syncer.records("delta_carried_over") + assert carried and all(c["island_id"] == 1 and c["discount"] == 0.5 ** c["lag"] for c in carried) + assert islands[0][0].base_version == rounds + # The dashboard folder reads the pool records as before. + from yeto.rl.engine.merge_metrics import MergeMetricsFolder + + folder = MergeMetricsFolder() + metrics = [m for m in map(folder.feed, syncer.records()) if m is not None] + assert len(metrics) == rounds + assert any(m["merge/carried_weight_share"] > 0 for m in metrics) + assert any(m["merge/island/0/weight_share"] == 1.0 and m["merge/timed_out"] for m in metrics) + + +def test_lease_expiry_rejoins_and_the_island_finishes(syncer_binary, tmp_path): + rounds = 8 + syncer = Syncer(syncer_binary, tmp_path, learners=2, rounds=rounds, + extra=["--quorum", "1", "--max-base-lag", "3", "--soft-deadline-s", "2", + "--island-lease-s", "1.5", "--final-grace-s", "10"]) + try: + islands = [_island(tmp_path, syncer, i, rounds, lease_s=1.5) for i in range(2)] + islands[0][0].debug_delay_s = lambda: 1.0 # keeps the run going while island 1 is silent + paused = islands[1][0] + paused.debug_pause = lambda: (1, 4.0) # after base 1: 4 s without heartbeats or pushes + assert _run_pool(tmp_path, syncer, islands) == 0 + finally: + syncer.kill() + leaves = syncer.records("pool_leave") + assert any(r["island_id"] == 1 and r["reason"] == "lease_expired" for r in leaves), leaves + rejoined = [r for r in syncer.records("pool_join") if r["island_id"] == 1 and r["catch_up"]] + assert rejoined, syncer.records("pool_join") + tape = _tape(tmp_path, 1) + assert any(e["event"] == "elastic_rejoin" for e in tape) + assert any(e["event"] == "elastic_debug_pause" for e in tape) + + +def test_final_grace_ends_the_run_without_an_island_that_left(syncer_binary, tmp_path): + rounds = 3 + syncer = Syncer(syncer_binary, tmp_path, learners=2, rounds=rounds, + extra=["--quorum", "1", "--max-base-lag", "2", "--island-lease-s", "30", + "--soft-deadline-s", "1", "--final-grace-s", "3"]) + try: + islands = [_island(tmp_path, syncer, i, rounds) for i in range(2)] + quitter = islands[1][0] + original = quitter.boundary + + def leave_after_first_round(driver, *, rollout_id, stats): + result = original(driver, rollout_id=rollout_id, stats=stats) + quitter.leave(LeaveReason.PREEMPTED) + raise KeyboardInterrupt("preempted") + quitter.boundary = leave_after_first_round + errors: list[BaseException] = [] + threads = [_run(driver, errors) for _sync, driver in islands] + for thread in threads: + thread.join(120) + assert syncer.wait() == 0 + finally: + syncer.kill() + assert islands[0][0].base_version == rounds + leaves = [r for r in syncer.records("pool_leave") if r["island_id"] == 1] + assert leaves and leaves[-1]["reason"] == "preempted", syncer.records("pool_leave") + assert any(isinstance(e, KeyboardInterrupt) for e in errors) + + +def test_pool_islands_take_round_cuts(): + """R-2: an island in a pool writes round cuts like a strict synced island.""" + assert issubclass(PoolAvgSync, RoundCutHooks) diff --git a/tests/test_rl_inter_island_launcher.py b/tests/test_rl_island_pool_launcher.py similarity index 59% rename from tests/test_rl_inter_island_launcher.py rename to tests/test_rl_island_pool_launcher.py index 8d709685f..c41a5aa3c 100644 --- a/tests/test_rl_inter_island_launcher.py +++ b/tests/test_rl_island_pool_launcher.py @@ -1,4 +1,4 @@ -"""--rl-island-scheduling launcher wiring (rl-inter-island-scheduling 0.13). No Ray.""" +"""Island pool launcher wiring: syncer membership flags and island flags. No Ray.""" from __future__ import annotations @@ -42,27 +42,30 @@ def _sft(): HMAC_CHECK = ': "${YETO_ISLAND_HMAC_KEY:?island HMAC key secret missing}" && ' -def test_legacy_syncer_command_is_byte_identical(head_launcher): - """secret-handling-hardening: the only change to the legacy RL line is the +POOL = ("--rl-sync-quorum-theta", "0.75", "--rl-sync-max-base-lag", "2") + + +def test_fixed_roster_syncer_command_is_byte_identical(head_launcher): + """secret-handling-hardening: the only change to the fixed-roster RL line is the island HMAC key check in front (the key itself is a secret env); SFT only gains the explicit --allow-unauthenticated-islands.""" - for extra in ((), ("--rl-island-scheduling", "legacy")): - args = _args(extra) - assert args.rl_island_scheduling == "legacy" - for n in (1, 2, 4): - assert launcher.syncer_command(args, n) == head_launcher.syncer_command(args, n).replace( - "mkdir -p ~/yeto-output && ", "mkdir -p ~/yeto-output && " + HMAC_CHECK, 1) - # P0-D8: one command per role; the HMAC check sits once in front. - commands = launcher.role_syncer_commands(args, n) - assert list(commands) == ["actor"] - assert launcher.syncer_command(args, n).endswith(HMAC_CHECK + commands["actor"]) + args = _args() + assert not launcher.island_pool(args) + for n in (1, 2, 4): + assert launcher.syncer_command(args, n) == head_launcher.syncer_command(args, n).replace( + "mkdir -p ~/yeto-output && ", "mkdir -p ~/yeto-output && " + HMAC_CHECK, 1) + # P0-D8: one command per role; the HMAC check sits once in front. + commands = launcher.role_syncer_commands(args, n) + assert list(commands) == ["actor"] + assert launcher.syncer_command(args, n).endswith(HMAC_CHECK + commands["actor"]) assert launcher.syncer_command(_sft(), 2) == ( head_launcher.syncer_command(_sft(), 2) + " --allow-unauthenticated-islands") -def test_legacy_learner_flags_unchanged(): +def test_fixed_roster_learner_flags_unchanged(): args = _args() - assert "--rl-island-scheduling" not in launcher._ports_infra_flags(args)[1] + flags = launcher._ports_infra_flags(args)[1] + assert "--rl-island-pool" not in flags and "--rl-sync-island-lease-s" not in flags @pytest.fixture(autouse=True) @@ -70,50 +73,56 @@ def _hmac_key(monkeypatch): monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") -def test_elastic_needs_hmac_key(monkeypatch): +def test_pool_and_fixed_roster_generate_a_key_when_none_is_set(monkeypatch): monkeypatch.delenv("YETO_ISLAND_HMAC_KEY") - with pytest.raises(ValueError, match="YETO_ISLAND_HMAC_KEY"): - launcher.syncer_command(_args(("--rl-island-scheduling", "elastic")), 2) - # secret-handling-hardening: legacy generates a per-run key instead. - assert launcher.syncer_command(_args(), 2).startswith("mkdir -p ~/yeto-output && " + HMAC_CHECK) + for extra in (POOL, ()): + assert launcher.syncer_command(_args(extra), 2).startswith( + "mkdir -p ~/yeto-output && " + HMAC_CHECK) -def test_elastic_syncer_flags_and_defaults(): - cmd = launcher.syncer_command(_args(("--rl-island-scheduling", "elastic")), 2) +def test_pool_syncer_flags_and_defaults(): + cmd = launcher.syncer_command(_args(POOL), 2) assert "k3y" not in cmd and "export YETO_ISLAND_HMAC_KEY" not in cmd # 0.20: secret, not argv assert cmd.startswith('mkdir -p ~/yeto-output && : "${YETO_ISLAND_HMAC_KEY:?') - cmd = launcher.syncer_command(_args(("--rl-island-scheduling", "elastic")), 2) - assert (" --island-scheduling-mode elastic --quorum-theta 0.75 --carry-gamma 0.5" - " --soft-deadline-s 900 --q-min 1 --max-carry-lag 2 --syncer-epoch 0") in cmd + # only the given options reach the syncer; lease and syncer epoch always do + assert " --quorum-theta 0.75 --island-lease-s 150.0 --syncer-epoch 0" in cmd + for flag in ("--carry-gamma", "--soft-deadline-s", "--q-min", "--final-grace-s", + "--island-scheduling-mode", "--max-carry-lag"): + assert flag not in cmd cmd = launcher.syncer_command(_args(( - "--rl-island-scheduling", "elastic", "--rl-quorum-theta", "0.6", "--rl-carry-gamma", - "0.25", "--rl-soft-deadline-s", "300", "--rl-q-min", "2", "--rl-max-carry-lag", "1", + "--rl-sync-quorum-theta", "0.6", "--rl-sync-carry-gamma", "0.25", "--rl-sync-soft-deadline-s", "300", + "--rl-sync-q-min", "2", "--rl-sync-island-lease-s", "60", "--rl-sync-max-base-lag", "1", "--rl-syncer-epoch", "4")), 3) - assert (" --island-scheduling-mode elastic --quorum-theta 0.6 --carry-gamma 0.25" - " --soft-deadline-s 300 --q-min 2 --max-carry-lag 1 --syncer-epoch 4") in cmd - assert "--island-scheduling-mode elastic" in launcher.role_syncer_commands( - _args(("--rl-island-scheduling", "elastic")), 2)["actor"] - assert " --rl-island-scheduling elastic --rl-syncer-epoch 0" in launcher._ports_infra_flags( - _args(("--rl-island-scheduling", "elastic")))[1] + assert (" --quorum-theta 0.6 --q-min 2 --carry-gamma 0.25 --soft-deadline-s 300" + " --island-lease-s 60.0 --syncer-epoch 4 --final-grace-s 300") in cmd + assert "--island-lease-s 150.0" in launcher.role_syncer_commands(_args(POOL), 2)["actor"] + assert launcher._ports_infra_flags(_args(POOL))[1].endswith( + " --rl-island-pool --rl-syncer-epoch 0 --rl-sync-island-lease-s 150.0") def test_invalid_combinations_refused(): - with pytest.raises(ValueError, match="need --rl-island-scheduling elastic"): - launcher.syncer_command(_args(("--rl-quorum-theta", "0.5")), 2) - with pytest.raises(ValueError, match="need --rl-island-scheduling elastic"): + with pytest.raises(ValueError, match="--rl-syncer-epoch needs island pool options"): launcher.syncer_command(_args(("--rl-syncer-epoch", "1")), 2) - with pytest.raises(ValueError): - launcher.syncer_command(_args(("--rl-island-scheduling", "elastic", - "--rl-quorum-theta", "1.5")), 2) - with pytest.raises(ValueError, match="only applies to --training-mode rl"): - launcher.syncer_command(SimpleNamespace(**vars(_sft()), rl_island_scheduling="elastic"), 2) - with pytest.raises(SystemExit): - _args(("--rl-island-scheduling", "auto")) + with pytest.raises(ValueError, match=r"quorum-theta must be in \(0, 1\]"): + launcher.syncer_command(_args(("--rl-sync-quorum-theta", "1.5")), 2) + with pytest.raises(ValueError, match="only apply to --training-mode rl"): + launcher.syncer_command(SimpleNamespace(**vars(_sft()), rl_sync_island_lease_s=30.0), 2) + for removed in ("--rl-island-scheduling", "--rl-max-carry-lag"): + with pytest.raises(SystemExit): + _args((removed, "1")) + + +def test_pool_needs_rounds_that_may_merge_without_an_island(): + from test_rl_launcher import _prepare_rl_args + + with pytest.raises(ValueError, match="merge without some island"): + _prepare_rl_args(_prov(_args(("--gpu", "modal:8xh100,modal:8xh100", "--rl-engine", "ports", + "--rl-sync-island-lease-s", "30")))) -def test_learner_passes_mode_to_elastic_config(): +def test_learner_passes_the_pool_switch_to_the_island_config(): from yeto.rl.adapters.miles import island_entry as learner - assert "island_scheduling" in Path(learner.__file__).read_text() + assert "rl_island_pool" in Path(learner.__file__).read_text() from yeto.rl.adapters.miles import elastic_wiring import inspect assert inspect.signature(elastic_wiring.build_elastic).parameters[ @@ -137,7 +146,7 @@ def test_syncer_epoch_read_from_status_json(tmp_path): reads = iter([None, None]) # the syncer is still starting for two polls real = launcher.syncer_status_reader(local_tape=str(tmp_path / "yeto-tape.jsonl")) sleep, clock = _no_sleep_clock() - args = _args(("--rl-island-scheduling", "elastic")) + args = _args(POOL) assert launcher.resolve_island_syncer_epoch( args, lambda: next(reads, None) or real(), sleep=sleep, clock=clock) == 6 assert args.rl_syncer_epoch == 6 @@ -149,21 +158,21 @@ def test_syncer_epoch_unreadable_falls_back_to_zero(tmp_path, capsys): (tmp_path / "status.json").write_text(json.dumps({"schema": "other", "syncer_epoch": 3})) sleep, clock = _no_sleep_clock() - args = _args(("--rl-island-scheduling", "elastic")) + args = _args(POOL) reader = launcher.syncer_status_reader(local_tape=str(tmp_path / "yeto-tape.jsonl")) assert launcher.resolve_island_syncer_epoch(args, reader, timeout_s=10, sleep=sleep, clock=clock) == 0 assert args.rl_syncer_epoch is None and "WARN" in capsys.readouterr().out assert " --rl-syncer-epoch 0" in launcher._ports_infra_flags(args)[1] -def test_syncer_epoch_explicit_or_legacy_never_reads(): +def test_syncer_epoch_explicit_or_fixed_roster_never_reads(): def boom(): raise AssertionError("status.json must not be read") - explicit = _args(("--rl-island-scheduling", "elastic", "--rl-syncer-epoch", "4")) + explicit = _args((*POOL, "--rl-syncer-epoch", "4")) assert launcher.resolve_island_syncer_epoch(explicit, boom) == 4 - legacy = _args() - assert launcher.resolve_island_syncer_epoch(legacy, boom) == 0 - assert "--rl-island-scheduling" not in launcher._ports_infra_flags(legacy)[1] + fixed = _args() + assert launcher.resolve_island_syncer_epoch(fixed, boom) == 0 + assert "--rl-island-pool" not in launcher._ports_infra_flags(fixed)[1] # ----------------------------------------------------------------- 0.20 HMAC key as a secret @@ -187,12 +196,12 @@ def test_hmac_key_is_secret_on_syncer_task(monkeypatch, tmp_path): import platform monkeypatch.setattr(platform, "system", lambda: "Linux") monkeypatch.setattr(platform, "machine", lambda: "x86_64") - task = launcher.make_syncer_task(_args(("--rl-island-scheduling", "elastic")), 2) + task = launcher.make_syncer_task(_args(POOL), 2) assert task.secrets == {"YETO_ISLAND_HMAC_KEY": "k3y"} assert "k3y" not in task.run and "k3y" not in json.dumps(task.envs or {}) - legacy = launcher.make_syncer_task(_args(), 2) - assert legacy.secrets == {"YETO_ISLAND_HMAC_KEY": "k3y"} # secret-handling-hardening - assert "k3y" not in legacy.run + fixed = launcher.make_syncer_task(_args(), 2) + assert fixed.secrets == {"YETO_ISLAND_HMAC_KEY": "k3y"} # secret-handling-hardening + assert "k3y" not in fixed.run def test_hmac_key_in_modal_cfg_envs_and_island_secrets(monkeypatch, tmp_path): @@ -201,7 +210,7 @@ def test_hmac_key_in_modal_cfg_envs_and_island_secrets(monkeypatch, tmp_path): _fake_sky(monkeypatch) monkeypatch.setenv("HOME", str(tmp_path)) - args = _args(("--gpu", "modal:8xh100", "--rl-engine", "ports", "--rl-island-scheduling", "elastic")) + args = _args(("--gpu", "modal:8xh100", "--rl-engine", "ports", *POOL)) _prepare_rl_args(_prov(args)) (spec,) = parse_gpu_spec(args.gpu) task = launcher.make_miles_island_task(args, spec, 0, 1, "127.0.0.1:29400") @@ -209,12 +218,12 @@ def test_hmac_key_in_modal_cfg_envs_and_island_secrets(monkeypatch, tmp_path): assert "k3y" not in task.run and "k3y" not in (task.setup or "") cfg = launcher.build_modal_island_config(args, spec, 0, task, "1.2.3.4:5000") assert cfg.envs["YETO_ISLAND_HMAC_KEY"] == "k3y" and "k3y" not in cfg.run_script - legacy = _args(("--gpu", "modal:8xh100", "--rl-engine", "ports")) - _prepare_rl_args(_prov(legacy)) - lt = launcher.make_miles_island_task(legacy, spec, 0, 1, "127.0.0.1:29400") - # secret-handling-hardening: legacy islands carry the key too. + fixed = _args(("--gpu", "modal:8xh100", "--rl-engine", "ports")) + _prepare_rl_args(_prov(fixed)) + lt = launcher.make_miles_island_task(fixed, spec, 0, 1, "127.0.0.1:29400") + # secret-handling-hardening: fixed-roster islands carry the key too. assert lt.secrets["YETO_ISLAND_HMAC_KEY"] == "k3y" - assert launcher.build_modal_island_config(legacy, spec, 0, lt, "1.2.3.4:5000").envs[ + assert launcher.build_modal_island_config(fixed, spec, 0, lt, "1.2.3.4:5000").envs[ "YETO_ISLAND_HMAC_KEY"] == "k3y" @@ -223,5 +232,5 @@ def test_head_job_ships_hmac_secret(): assert "secrets.update(launcher.island_hmac_secret(args))" in src assert launcher.island_hmac_secret(_args()) == {"YETO_ISLAND_HMAC_KEY": "k3y"} assert launcher.island_hmac_secret(_sft()) == {} - assert launcher.island_hmac_secret(_args(("--rl-island-scheduling", "elastic"))) == { + assert launcher.island_hmac_secret(_args(POOL)) == { "YETO_ISLAND_HMAC_KEY": "k3y"} diff --git a/tests/test_rl_island_pool_session.py b/tests/test_rl_island_pool_session.py new file mode 100644 index 000000000..2281e620e --- /dev/null +++ b/tests/test_rl_island_pool_session.py @@ -0,0 +1,147 @@ +"""PoolAvgSync at the session level against an in-process fake syncer: debug +switches, the finished race, island tape fields, and the launcher wiring. +CPU, no Ray. The real syncer is covered by test_rl_island_pool_e2e.py.""" + +from __future__ import annotations + +import json + +import pytest +import torch + +from pool_fakes import FakePoolSyncer +from test_rl_engine_driver import _driver, _engine, _strict_config +from yeto.rl.engine.bridges import PoolAvgSync +from yeto.rl.pool_client import ( + DEBUG_DELAY_ENV, DEBUG_PAUSE_ENV, parse_elastic_debug_delay, parse_elastic_debug_pause, +) + + +def _sync(tmp_path, engine, client, *, rounds=2, learner_id=0): + return PoolAvgSync(_strict_config(tmp_path, engine, learner_id=learner_id, rounds=rounds), + client=client, base_wait_s=5.0) + + +def _tape(tmp_path): + return [json.loads(line) for line in (tmp_path / "events.jsonl").read_text().splitlines()] + + +def test_debug_switches_parse_and_reach_the_island_env(): + from test_rl_launcher import _args + from yeto import launcher + + assert parse_elastic_debug_pause("1:1:120") == {1: (1, 120.0)} + assert parse_elastic_debug_pause(" 0:2:5, 1:1:7.5 ") == {0: (2, 5.0), 1: (1, 7.5)} + assert parse_elastic_debug_pause("") == {} and parse_elastic_debug_pause(None) == {} + for bad in ("1:120", "1:1:0", "-1:1:5", "1:1:5:9", "a:1:5"): + with pytest.raises(ValueError): + parse_elastic_debug_pause(bad) + assert parse_elastic_debug_delay("1:2.5,3") == {1: 2.5, -1: 3.0} + pool = ("--rl-sync-island-lease-s", "60") + assert launcher.elastic_debug_pause_env(_args(pool)) == {} + assert launcher.elastic_debug_pause_env(_args(pool + ("--rl-elastic-debug-pause", "1:1:120"))) == { + DEBUG_PAUSE_ENV: "1:1:120"} + assert launcher.elastic_debug_delay_env(_args(pool + ("--rl-elastic-debug-delay-s", "1:3"))) == { + DEBUG_DELAY_ENV: "1:3"} + with pytest.raises(ValueError, match="needs island pool options"): + launcher.elastic_debug_pause_env(_args(("--rl-elastic-debug-pause", "1:1:120"))) + + +def test_default_off_session_never_pauses_and_pushes_each_round(tmp_path, monkeypatch): + monkeypatch.delenv(DEBUG_PAUSE_ENV, raising=False) + fake = FakePoolSyncer() + engine = _engine(torch.tensor([1.0, 3.0])) + client = fake.client() + _driver(engine, _sync(tmp_path, engine, client), tmp_path).run() + assert client.paused == [] and fake.inits == 1 + assert [(d["base_version"], d["update"].tolist()) for d in fake.deltas] == [ + (0, [1.0, 3.0]), (1, [1.0, 3.0])] + assert "elastic_debug_pause" not in (tmp_path / "events.jsonl").read_text() + + +def test_session_pauses_only_the_named_island_once_after_after_v(tmp_path, monkeypatch): + monkeypatch.setenv(DEBUG_PAUSE_ENV, "0:1:0.05,5:0:9") + fake = FakePoolSyncer() + engine = _engine(torch.tensor([1.0, 3.0])) + client = fake.client() + _driver(engine, _sync(tmp_path, engine, client, rounds=3), tmp_path).run() + armed = [e for e in _tape(tmp_path) if e.get("phase") == "elastic_debug_pause_armed"] + assert len(armed) == 1 and armed[0]["base_version"] == 1 and armed[0]["pause_s"] == 0.05 + assert client.paused == [0.05] + assert [e["pause_s"] for e in _tape(tmp_path) if e.get("event") == "elastic_debug_pause"] == [0.05] + + +def test_an_island_that_applied_the_final_cut_stops_without_another_push(tmp_path): + """S17 M1: the island stops once it applied the final outer version.""" + fake = FakePoolSyncer(final=3) + engine = _engine(torch.tensor([1.0, 3.0])) + sync = _sync(tmp_path, engine, fake.client(), rounds=3) + _driver(engine, sync, tmp_path).run() + tape = _tape(tmp_path) + assert sync.base_version == 3 and not [e for e in tape if e["event"] == "elastic_finished"] + assert sum(e["event"] == "rl_local_round" for e in tape) == 1 + + +def test_push_racing_the_end_takes_the_finished_path(tmp_path): + """The island's base is not final yet, but the run ended meanwhile: the push is + refused and the island ends normally on the final cut (elastic_finished).""" + fake = FakePoolSyncer([0.0, 0.0], version=1, final=1) + engine = _engine(torch.tensor([1.0, 3.0])) + sync = _sync(tmp_path, engine, fake.client(), rounds=4) + fake.final = None # the start sees a running pool at version 1 ... + original = sync.start + + def start(driver): + started = original(driver) + fake.final = 1 # ... which finishes before the island's first push + return started + sync.start = start + _driver(engine, sync, tmp_path).run() + done = [e for e in _tape(tmp_path) if e["event"] == "elastic_finished"] + assert len(done) == 1 and done[0]["final_outer_version"] == 1 + assert fake.leaves == [0] and fake.deltas == [] + + +def test_pool_round_fields_match_the_strict_local_round(tmp_path): + fake = FakePoolSyncer() + engine = _engine(torch.tensor([3.0, 4.0])) + _driver(engine, _sync(tmp_path, engine, fake.client(), rounds=1), tmp_path).run() + [record] = [e for e in _tape(tmp_path) if e["event"] == "rl_local_round"] + assert record["rl/local_delta_norm"] == pytest.approx(5.0) + assert record["sync/base_outer_version"] == 0 and record["sync/bytes_sent"] == 8 + assert record["rl/global_policy_version"] == 0 + + +def test_the_island_records_its_sync_timing(tmp_path): + fake = FakePoolSyncer(reply_delay_s=0.2) + engine = _engine(torch.tensor([1.0, 3.0])) + _driver(engine, _sync(tmp_path, engine, fake.client(), rounds=1), tmp_path).run() + [timing] = [e for e in _tape(tmp_path) if e["event"] == "rl_sync_timing"] + assert timing["outer_sync"] == "elastic" and timing["base_version"] == 1 + assert timing["sync/global_ready_ts"] >= timing["sync/submit_ts"] + assert timing["sync/sync_wait_seconds"] >= 0.0 + + +def test_client_events_reach_the_tape_once(tmp_path, monkeypatch): + """1b-g: an injected client's tape hook is chained exactly once.""" + monkeypatch.setenv(DEBUG_PAUSE_ENV, "0:1:0.05") + fake = FakePoolSyncer() + engine = _engine(torch.tensor([1.0, 3.0])) + _driver(engine, _sync(tmp_path, engine, fake.client(), rounds=3), tmp_path).run() + assert [e.get("event") for e in _tape(tmp_path)].count("elastic_debug_pause") == 1 + + +def test_the_miles_entry_selects_the_pool_session(monkeypatch): + from types import SimpleNamespace + + from yeto.rl.adapters.miles import entry + from yeto.rl.engine.bridges import StrictAvgSync + + monkeypatch.setattr("yeto.rl.engine.bridges.StrictIslandProgress", lambda args: None) + base = dict(yeto_rl_bridge_config=SimpleNamespace(groups_per_round=2), rollout_batch_size=2) + sync, _ = entry.build_sync(SimpleNamespace(**base, yeto_rl_island_pool=True, + yeto_rl_syncer_epoch=3, yeto_rl_island_lease_s=45.0), + yeto_policy_sync=True) + assert isinstance(sync, PoolAvgSync) and sync.syncer_epoch == 3 and sync.lease_s == 45.0 + strict, _ = entry.build_sync(SimpleNamespace(**base), yeto_policy_sync=True) + assert isinstance(strict, StrictAvgSync) diff --git a/tests/test_rl_island_round_cuts.py b/tests/test_rl_island_round_cuts.py index dbba65e6f..9d98bed73 100644 --- a/tests/test_rl_island_round_cuts.py +++ b/tests/test_rl_island_round_cuts.py @@ -118,7 +118,8 @@ def emit(self, name, **fields): def test_synced_syncs_take_periodic_round_cuts_only(): strict = bridges.StrictAvgSync.__new__(bridges.StrictAvgSync) decoupled = bridges.DecoupledSync.__new__(bridges.DecoupledSync) - for sync in (strict, decoupled): + pool = bridges.PoolAvgSync.__new__(bridges.PoolAvgSync) # LoRA island pool (R-2) + for sync in (strict, decoupled, pool): assert isinstance(sync, bridges.RoundCutHooks) assert sync.round_cuts is None and sync.cut_due(4) is False # off by default driver = _Emitter() @@ -130,7 +131,6 @@ def test_synced_syncs_take_periodic_round_cuts_only(): sync.at_safe_point(driver, rollout_id=4) assert sync.round_cuts.calls == [("due", 4, False), ("prepare", 4, False), ("save", 4, False)] assert [name for name, _ in driver.events] == ["rl_round_cut", "rl_cut_saved"] - assert not isinstance(bridges.ElasticAvgSync.__new__(bridges.ElasticAvgSync), bridges.RoundCutHooks) def test_local_only_sync_still_cuts_at_the_end(): @@ -161,8 +161,9 @@ def test_resume_wiring_write_only_needs_a_synced_sync(tmp_path): assert full.cut_due(4) is False # no checkpoint wired yet: no cut with pytest.raises(ValueError, match="takes no island round cuts"): resume_wiring_for(miles_args, sync=bridges.LocalOnlySync(2), environ={}) - with pytest.raises(ValueError, match="takes no island round cuts"): - resume_wiring_for(miles_args, sync=bridges.ElasticAvgSync.__new__(bridges.ElasticAvgSync), environ={}) + pool = bridges.PoolAvgSync.__new__(bridges.PoolAvgSync) # R-2: pool islands cut too + pool_args = SimpleNamespace(yeto_rl_resume={**config, "state_dir": str(tmp_path / "pool-state")}) + assert resume_wiring_for(pool_args, sync=pool, environ={}).stop_after is None config.pop("write_only") with pytest.raises(ValueError, match="needs --rl-single-island-no-sync"): resume_wiring_for(miles_args, sync=strict, environ={}) @@ -190,6 +191,18 @@ def test_synced_launch_forwards_island_cut_store(): assert " --rl-cut-auto-mtbf-s 21600.0 --rl-cut-max-interval-s 1800.0" in flags +def test_pool_launch_forwards_island_cut_store(): + """R-2: LoRA islands in an island pool write round cuts like strict islands.""" + from yeto.launch import rl_ports + + args = _cli("--rl-checkpoint-store", "/mnt/store", "--rl-cut-every", "4", + "--rl-sync-quorum-theta", "0.75", "--rl-sync-max-base-lag", "2") + assert rl_ports._check_resume_launch_flags(args) is False + flags = rl_ports._ports_infra_flags(args)[1] + assert " --rl-island-pool" in flags + assert " --rl-island-cut-store /mnt/store/island-cuts" in flags and " --rl-cut-every 4" in flags + + @pytest.mark.parametrize("extra, match", [ (("--rl-checkpoint-store", "/s"), "needs --rl-single-island-no-sync, synced trainers, a cut interval"), (("--rl-checkpoint-store", "/s", "--rl-cut-every", "2", "--rl-stop-after-rounds", "3"), @@ -197,8 +210,6 @@ def test_synced_launch_forwards_island_cut_store(): (("--rl-checkpoint-store", "/s", "--rl-cut-keep", "2"), "a cut interval"), (("--rl-checkpoint-store", "/s", "--rl-cut-every", "2", "--rl-cut-auto-mtbf-s", "60"), "exclude each other"), - (("--rl-checkpoint-store", "/s", "--rl-cut-every", "2", "--rl-island-scheduling", "elastic"), - "need --rl-island-scheduling legacy"), (("--rl-cut-auto-mtbf-s", "60"), "needs --rl-checkpoint-store"), ]) def test_synced_launch_refusals(extra, match): diff --git a/tests/test_rl_island_tasks.py b/tests/test_rl_island_tasks.py index 2e63db89a..6e4876dcb 100644 --- a/tests/test_rl_island_tasks.py +++ b/tests/test_rl_island_tasks.py @@ -328,15 +328,16 @@ def test_non_ports_paths_refuse_island_tasks(factory): MIXED_GPUS = "aws:1xa100@us-east-1,aws:2xa100@us-west-2" MIXED_PLACEMENT = ("--rl-island-placement", "1:placement=fixed-partition", - "--rl-rollout-gpus", "1", "--rl-island-scheduling", "elastic") + "--rl-rollout-gpus", "1") +ISLAND_POOL = ("--rl-sync-quorum-theta", "0.75", "--rl-sync-max-base-lag", "2") -def _mixed_island_args(catalog, tmp_path, monkeypatch, spec, *extra): +def _mixed_island_args(catalog, tmp_path, monkeypatch, spec, *extra, pool=ISLAND_POOL): from yeto import runs monkeypatch.setattr(runs, "RUNS_DIR", tmp_path / "runs") - monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") # elastic scheduling needs one - args = _launcher_args("ports", (F, spec, *MIXED_PLACEMENT, *extra), gpu=MIXED_GPUS) + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") + args = _launcher_args("ports", (F, spec, *MIXED_PLACEMENT, *pool, *extra), gpu=MIXED_GPUS) args.data, args.data_revision, args.controller = str(catalog), None, "local" return args @@ -375,9 +376,8 @@ def test_island_tasks_with_mixed_placement_keep_both_refusals(catalog, tmp_path, "--rl-eval-samples-per-prompt", "1") with pytest.raises(ValueError, match="training islands do not run eval"): launcher._prepare_rl_args(args) - args = _mixed_island_args(catalog, tmp_path, monkeypatch, "l0=web1;l1=web2") - args.rl_island_scheduling = "legacy" - with pytest.raises(ValueError, match="needs --rl-island-scheduling elastic"): + args = _mixed_island_args(catalog, tmp_path, monkeypatch, "l0=web1;l1=web2", pool=()) + with pytest.raises(ValueError, match="needs island pool options"): launcher._prepare_rl_args(args) diff --git a/tests/test_rl_launcher_dense_full.py b/tests/test_rl_launcher_dense_full.py index 1ce3a062a..784ee37ff 100644 --- a/tests/test_rl_launcher_dense_full.py +++ b/tests/test_rl_launcher_dense_full.py @@ -88,7 +88,7 @@ def test_lora_island_command_is_unchanged(): (("--rl-placement", "colocated"), "fixed-partition"), (("--rl-parameter-layout-sha256", "nope"), "rl-parameter-layout-sha256"), (("--local-rl-rounds-per-sync", "2"), "H=1"), - (("--rl-island-scheduling", "elastic"), "elastic"), + (("--rl-sync-island-lease-s", "30"), "membership options need --rl-sync-full-parameter-rounds async"), ], ) def test_dense_full_refusals(extra, message, monkeypatch): diff --git a/tests/test_rl_launcher_sao_streaming.py b/tests/test_rl_launcher_sao_streaming.py index ddea7b4e8..98e70bd31 100644 --- a/tests/test_rl_launcher_sao_streaming.py +++ b/tests/test_rl_launcher_sao_streaming.py @@ -178,7 +178,7 @@ def test_island_entry_accepts_the_sao_island_command(tmp_path): @pytest.mark.parametrize( ("extra", "message"), ( - (("--rl-island-scheduling", "elastic"), "elastic"), + (("--rl-sync-island-lease-s", "30"), "island pool"), (("--experimental-rl-sync",), "experimental"), ), ) diff --git a/tests/test_rl_merge_metrics.py b/tests/test_rl_merge_metrics.py index 489a7d6d6..c9453cae6 100644 --- a/tests/test_rl_merge_metrics.py +++ b/tests/test_rl_merge_metrics.py @@ -3,10 +3,6 @@ from __future__ import annotations import json -import os -import socket -import subprocess -import time import pytest @@ -73,93 +69,3 @@ def test_wandb_routes_elastic_records_through_the_folder_only(): assert record_metrics(step, folder)["merge/island/1/weight_share"] == 1.0 strict = {"step": 3, "responders": []} assert record_metrics(strict, folder) == tape_metrics(strict) - - -def _wait_tape(tape, predicate, timeout=10.0): - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if tape.exists(): - rows = [json.loads(line) for line in tape.read_text().splitlines() if line.strip()] - if predicate(rows): - return rows - time.sleep(0.05) - raise AssertionError(f"tape condition not met: {tape.read_text() if tape.exists() else 'no tape'}") - - -@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), - reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") -def test_real_syncer_tape_merge_shares_and_clock(tmp_path): - """Real elastic syncer, two client islands (theta 0.5: one arrival steps). - - v0: island 1 alone -> share 1.0. Island 2 then sends on base 0 (lag 1, - discount 0.5, w = 20^2/2 * 0.5 = 100); island 1 on base 1 (w = 100) -> - v1 shares 0.5 / 0.5. Island 2 sends on base 7 (unknown_version), then on - base 2 -> v2 rejected fraction 1/2.""" - from yeto.rl.elastic_client import ElasticClientConfig, ElasticIslandClient - - with socket.socket() as s: - s.bind(("127.0.0.1", 0)) - port = s.getsockname()[1] - key = "merge-metrics" - tape = tmp_path / "syncer-tape.jsonl" - syncer = subprocess.Popen( - [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", - "--total-steps", "5", "--event-tape", str(tape), - "--island-scheduling-mode", "elastic", "--quorum-theta", "0.5", "--carry-gamma", "0.5", - "--soft-deadline-s", "60", "--q-min", "1", "--max-carry-lag", "2", - "--island-lease-s", "6", "--syncer-epoch", "1"], - env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), - stdout=open(tmp_path / "syncer.log", "w"), stderr=subprocess.STDOUT) - islands = {i: ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), i, lease_s=6.0, - syncer_epoch=1), key.encode()) - for i in (1, 2)} - started = time.time() - try: - for c in islands.values(): - for _ in range(50): - try: - c.join() - break - except OSError: - time.sleep(0.1) - islands[1].elastic_init([0.0] * 4) - assert islands[1].wait_base(newer_than=None, timeout_s=10.0) is not None - update = [0.1] * 4 - islands[1].delta_tensor(base_version=0, c_tokens=10, c_steps=1, update=update) - assert islands[2].wait_base(newer_than=0, timeout_s=10.0) is not None - islands[2].delta_tensor(base_version=0, c_tokens=20, c_steps=2, update=update) - _wait_tape(tape, lambda rows: any(r["kind"] == "delta_carried_over" for r in rows)) - islands[1].delta_tensor(base_version=1, c_tokens=10, c_steps=1, update=update) - assert islands[2].wait_base(newer_than=1, timeout_s=10.0) is not None - islands[2].delta_tensor(base_version=7, c_tokens=20, c_steps=2, update=update) - _wait_tape(tape, lambda rows: any(r["kind"] == "delta_rejected" for r in rows)) - islands[2].delta_tensor(base_version=2, c_tokens=20, c_steps=2, update=update) - rows = _wait_tape(tape, lambda rows: sum(r["kind"] == "outer_step" for r in rows) >= 3) - finally: - for c in islands.values(): - c.close() - syncer.kill() - syncer.wait(5) - - steps = _fold(rows) - assert [m["merge/base_version"] for m in steps[:3]] == [0, 1, 2] - v0, v1, v2 = steps[:3] - assert v0["merge/island/1/weight_share"] == 1.0 and v0["merge/submitted"] == 1 - assert v1["merge/island/1/weight_share"] == pytest.approx(0.5) - assert v1["merge/island/2/weight_share"] == pytest.approx(0.5) - assert v1["merge/carried_weight_share"] == pytest.approx(0.5) - assert v1["merge/carried_fraction"] == 0.5 and v1["merge/carried_discount_mean"] == 0.5 - assert v2["merge/rejected_fraction/unknown_version"] == 0.5 - assert v2["merge/island/2/weight_share"] == 1.0 - - # L5.7 clock: every line has time_unix within the test window; deltas - # carry their arrival offset; outer steps their duration and last island. - ended = time.time() - assert all(started - 1 <= r["time_unix"] <= ended + 1 for r in rows) - for r in rows: - if r["kind"] in ("delta_accepted", "delta_carried_over"): - assert r["arrival_offset_s"] >= 0.0 - outer = [r for r in rows if r["kind"] == "outer_step"][:3] - assert [r["last_arrival_island"] for r in outer] == [1, 1, 2] - assert all(r["step_seconds"] >= 0.0 for r in outer) - assert v1["merge/step_seconds"] == outer[1]["step_seconds"] diff --git a/tests/test_rl_merged_release.py b/tests/test_rl_merged_release.py index 0539f894f..8d6518ddf 100644 --- a/tests/test_rl_merged_release.py +++ b/tests/test_rl_merged_release.py @@ -294,7 +294,7 @@ def test_buffer_lane_of_island_args(): # ---------------------------------------------------------------- launch and island flags SYNCED = ("--rl-sample-role", "trainer", "--rl-max-policy-age", "2", "--rl-checkpoint-store", "modal-volume://yeto-c3/run-a", "--rl-sample-wait-s", "900", - "--rl-island-scheduling", "elastic", "--cluster-prefix", "c3-synced", + "--rl-sync-quorum-theta", "0.75", "--rl-sync-max-base-lag", "2", "--cluster-prefix", "c3-synced", "--rl-algorithm-spec", "examples/rl_algorithms/decoupled-ppo-c2.json") @@ -334,14 +334,13 @@ def test_synced_trainer_plan_forwards_the_sample_store_and_lanes(): def test_synced_trainer_launch_writes_island_round_cuts(): - """S27 R-2: a cut interval gives every synced trainer its own write-only cuts.""" + """S27 R-2: a cut interval gives every synced trainer, in a fixed roster or an + island pool, its own write-only cuts.""" from types import SimpleNamespace as NS from yeto.rl.adapters.miles import island_entry - with pytest.raises(ValueError, match="need --rl-island-scheduling legacy"): - _synced_plan(("--rl-cut-every", "2")) # ElasticAvgSync takes no island cuts - _, plan = _synced_plan(("--rl-cut-every", "2", "--rl-island-scheduling", "legacy")) + _, plan = _synced_plan(("--rl-cut-every", "2")) for m, island in enumerate(plan["island_requests"]): tokens = _island_tokens(island["learner_command"], str(m)) root = tokens[tokens.index("--rl-island-cut-store") + 1] @@ -351,8 +350,7 @@ def test_synced_trainer_launch_writes_island_round_cuts(): assert miles_args.yeto_rl_resume["store"] == f"{root}/island{m}" assert miles_args.yeto_rl_resume["write_only"] is True with pytest.raises(ValueError, match="does not resume from round cuts"): - _synced_plan(("--rl-cut-every", "1", "--rl-stop-after-rounds", "2", - "--rl-island-scheduling", "legacy")) + _synced_plan(("--rl-cut-every", "1", "--rl-stop-after-rounds", "2")) def test_island_refuses_a_sample_store_outside_a_synced_trainer(): diff --git a/tests/test_rl_mixed_island_placement.py b/tests/test_rl_mixed_island_placement.py index 81a9938e9..f4c2c32f1 100644 --- a/tests/test_rl_mixed_island_placement.py +++ b/tests/test_rl_mixed_island_placement.py @@ -17,7 +17,7 @@ GPUS = ("--gpu", "aws:1xa100@us-east-1,aws:2xa100@us-west-2") MIXED = (*GPUS, "--rl-island-placement", "1:placement=fixed-partition", "--rl-rollout-gpus", "1") -ELASTIC = ("--rl-island-scheduling", "elastic") +ISLAND_POOL = ("--rl-sync-quorum-theta", "0.75", "--rl-sync-max-base-lag", "2") @pytest.fixture(autouse=True) @@ -89,7 +89,7 @@ def test_without_overrides_every_island_uses_the_run_wide_flags(): def test_a_colocated_island_drops_the_fixed_partition_only_values(): - args = _cli((*MIXED, "--rl-standby-gpus", "0", *ELASTIC)) + args = _cli((*MIXED, "--rl-standby-gpus", "0", *ISLAND_POOL)) colocated, partitioned = launcher.rl_island_args(args, 0), launcher.rl_island_args(args, 1) assert (colocated.rl_placement, colocated.rollout_num_gpus) == ("colocated", None) assert (partitioned.rl_placement, partitioned.rollout_num_gpus) == ("fixed-partition", 1) @@ -98,15 +98,15 @@ def test_a_colocated_island_drops_the_fixed_partition_only_values(): # --- launch checks (before any cloud work) ------------------------------------------ -def test_mixed_placement_under_legacy_scheduling_is_refused(): +def test_mixed_placement_with_a_fixed_roster_is_refused(): args = _cli(MIXED) - with pytest.raises(ValueError, match="needs --rl-island-scheduling elastic") as error: + with pytest.raises(ValueError, match="needs island pool options") as error: launcher._prepare_rl_args(args) assert "island 0 colocated, island 1 fixed-partition" in str(error.value) -def test_mixed_placement_under_elastic_scheduling_passes(): - args = _cli((*MIXED, *ELASTIC)) +def test_mixed_placement_in_an_island_pool_passes(): + args = _cli((*MIXED, *ISLAND_POOL)) launcher._prepare_rl_args(args) # includes the per-island shape checks assert args.rl_expected_algorithm_sha256 @@ -114,14 +114,14 @@ def test_mixed_placement_under_elastic_scheduling_passes(): @pytest.mark.parametrize("flag, dest", [("--rl-elastic", "rl_elastic"), ("--rl-overlap-eval", "rl_overlap_eval")]) def test_fixed_partition_only_switches_name_the_colocated_island(flag, dest): - args = _cli((*MIXED, *ELASTIC)) + args = _cli((*MIXED, *ISLAND_POOL)) setattr(args, dest, True) with pytest.raises(ValueError, match=rf"{flag} needs .* colocated island\(s\): \[0\]"): launcher._check_ports_infra_switches(args) def test_a_colocated_island_cannot_reserve_standby_gpus(): - args = _cli((*MIXED, *ELASTIC, "--rl-island-placement", "0:standby-gpus=1")) + args = _cli((*MIXED, *ISLAND_POOL, "--rl-island-placement", "0:standby-gpus=1")) with pytest.raises(ValueError, match="island 0 is colocated and cannot reserve standby"): launcher._check_ports_infra_switches(args) @@ -131,14 +131,14 @@ def test_mixed_launch_refusals_create_no_cloud_resources(monkeypatch): import sys monkeypatch.setitem(sys.modules, "sky", None) # any sky import would fail loudly - with pytest.raises(ValueError, match="needs --rl-island-scheduling elastic"): + with pytest.raises(ValueError, match="needs island pool options"): launcher._prepare_rl_args(_cli(MIXED)) # --- per-island shapes and commands ------------------------------------------------ def test_each_island_gets_its_own_shape_and_command(monkeypatch): - args = _cli((*MIXED, *ELASTIC)) + args = _cli((*MIXED, *ISLAND_POOL)) launcher._prepare_rl_args(args) colocated, partitioned = (_learner_flags(t) for t in _island_tasks(args, monkeypatch)) assert "--rl-placement" not in colocated and "--rollout-num-gpus" not in colocated @@ -151,7 +151,7 @@ def test_each_island_gets_its_own_shape_and_command(monkeypatch): == _value(partitioned, "--rl-expected-algorithm-sha256") == args.rl_expected_algorithm_sha256) for flags in (colocated, partitioned): - assert _value(flags, "--rl-island-scheduling") == "elastic" + assert "--rl-island-pool" in flags assert "--rl-per-island-placement" in flags @@ -190,7 +190,7 @@ def test_policy_age_overlap_differs_by_island_placement(): # --- records ----------------------------------------------------------------------- def test_run_manifest_rows_record_every_island_placement(): - args = _cli((*MIXED, *ELASTIC)) + args = _cli((*MIXED, *ISLAND_POOL)) launcher._prepare_rl_args(args) rows = launcher.rl_island_placement_records(args) assert [(r["island"], r["placement"], r["standby_gpus"]) for r in rows] == [ diff --git a/tests/test_rl_object_store.py b/tests/test_rl_object_store.py index 91cf5812c..f047faf7b 100644 --- a/tests/test_rl_object_store.py +++ b/tests/test_rl_object_store.py @@ -17,7 +17,7 @@ from yeto.rl.engine.policy_release import LORA, PolicyFollower, PolicyReleaseStore from yeto.rl.engine.resume import STORE_IMPL_ENV, CheckpointStore, store_for from yeto.rl.engine.sample_buffer import FileSampleBuffer, ObjectIndex, SampleBufferError -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine.trainer_lifecycle import SampleIndexEntry from yeto.rl.engine.trainer_remote import CommandReader, CommandWriter diff --git a/tests/test_rl_outer_merge_options.py b/tests/test_rl_outer_merge_options.py index 5b86410f6..e671ee6a2 100644 --- a/tests/test_rl_outer_merge_options.py +++ b/tests/test_rl_outer_merge_options.py @@ -92,6 +92,17 @@ def test_max_base_lag_rounds_become_versions_for_the_syncer(preset, versions): assert _syncer_flags(plan)["--max-base-lag"] == versions +def test_outer_merge_options_reach_an_island_pool_syncer(monkeypatch): + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") + _, plan = _plan(("--rl-sync-island-lease-s", "30", "--rl-sync-quorum", "1", + "--rl-sync-max-base-lag", "2", "--rl-sync-outer-lr", "0.5"), + gpu=THREE_ISLANDS) + flags = _syncer_flags(plan) + assert (flags["--quorum"], flags["--max-base-lag"], flags["--outer-lr"]) == ("1", "2", "0.5") + [actor] = plan["role_syncers"] + assert "--island-lease-s 30" in actor["command"] + + @pytest.mark.parametrize( ("option", "value", "flag"), [ @@ -138,7 +149,6 @@ def test_rda_matrix_merge_reaches_the_decoupled_islands(): (("--rl-sync-quorum", "1", "--experimental-rl-sync", "--fragments", "1", "--pipeline", "1", "--merge-alpha", "0", "--wire-dtype", "f32"), "experimental-rl-sync"), - (("--rl-sync-quorum", "1", "--rl-island-scheduling", "elastic"), "legacy"), (("--rl-sync-quorum", "0"), "between 1 and the 2 islands"), (("--rl-sync-quorum", "3"), "between 1 and the 2 islands"), (("--rl-sync-grace-ms", "-1"), "grace-ms must be >= 0"), diff --git a/tests/test_rl_policy_age.py b/tests/test_rl_policy_age.py index bc9b5b1e0..eba86d589 100644 --- a/tests/test_rl_policy_age.py +++ b/tests/test_rl_policy_age.py @@ -44,17 +44,16 @@ def test_strict_handshake_refuses_an_island_with_another_limit(): assert session_contract_hash(fp, bind_policy_age(MILES, 1)) != syncer # limit 1: refused -def test_elastic_join_carries_the_bound_identity(monkeypatch): - """Elastic JOIN: the island identity bytes differ, so only that connection is refused.""" +def test_pool_hello_carries_the_bound_identity(monkeypatch): + """Island pool HELLO: the session contracts differ, so only that connection is refused.""" from yeto.rl.adapters.miles import island_entry - from yeto.rl.elastic_client import ElasticClientConfig + from yeto.rl.engine.backend_identity import session_contract_hash zero = island_entry._backend_identity_sha256(SimpleNamespace(rl_engine="ports")) one = island_entry._backend_identity_sha256(SimpleNamespace(rl_engine="ports", rl_max_policy_age=1)) assert zero == MILES and one == bind_policy_age(MILES, 1) - a = ElasticClientConfig(("x", 0), 1, backend_identity_sha256=zero).backend_identity_bytes() - b = ElasticClientConfig(("x", 0), 2, backend_identity_sha256=one).backend_identity_bytes() - assert a != b + fingerprint = bytes(32) + assert session_contract_hash(fingerprint, zero) != session_contract_hash(fingerprint, one) @pytest.mark.parametrize("backend,limit,stage", [("miles", 5, 3), ("verl", 5, 2)]) diff --git a/tests/test_rl_elastic_rejoin_final_round.py b/tests/test_rl_pool_rejoin_final_round.py similarity index 89% rename from tests/test_rl_elastic_rejoin_final_round.py rename to tests/test_rl_pool_rejoin_final_round.py index bbe3c25fe..dbbc20ebf 100644 --- a/tests/test_rl_elastic_rejoin_final_round.py +++ b/tests/test_rl_pool_rejoin_final_round.py @@ -10,11 +10,11 @@ import torch from test_rl_engine_driver import _driver, _engine, _strict_config -from yeto.rl.engine.bridges import ElasticAvgSync +from yeto.rl.engine.bridges import PoolAvgSync class _Client: - """Scripted elastic client: the outer version returned after the n-th delta.""" + """Scripted island-pool client: the outer version returned after the n-th delta.""" def __init__(self, versions_after_delta): self.versions = list(versions_after_delta) @@ -30,7 +30,7 @@ def join(self): return SimpleNamespace(catch_up=False, base_version=0) def _base(self, version): - self.latest_base = SimpleNamespace(params=tuple(self.params), outer_version=version) + self.latest_base = SimpleNamespace(params=self.params.clone(), outer_version=version, final=False) return self.latest_base def wait_base(self, *, newer_than, timeout_s): @@ -41,7 +41,7 @@ def wait_base(self, *, newer_than, timeout_s): return self._base(self.versions[len(self.deltas) - 1]) def elastic_init(self, params): - self.params = list(params) + self.params = params.clone() def raise_if_finished(self): pass @@ -59,9 +59,9 @@ def close(self): def _run(tmp_path, versions, rounds=3): engine = _engine(torch.tensor([1.0, 3.0])) client = _Client(versions) - sync = ElasticAvgSync(_strict_config(tmp_path, engine, learner_id=1, rounds=rounds), client=client, - base_wait_s=1.0) - _driver(engine, sync, tmp_path, learner_id=1).run() + sync = PoolAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=rounds), client=client, + base_wait_s=1.0) + _driver(engine, sync, tmp_path, learner_id=0).run() ev = [json.loads(l) for l in (tmp_path / "events.jsonl").read_text().splitlines()] return client, sync, ev diff --git a/tests/test_rl_elastic_restart_cursor.py b/tests/test_rl_pool_restart_cursor.py similarity index 82% rename from tests/test_rl_elastic_restart_cursor.py rename to tests/test_rl_pool_restart_cursor.py index f3884aea2..56e796684 100644 --- a/tests/test_rl_elastic_restart_cursor.py +++ b/tests/test_rl_pool_restart_cursor.py @@ -1,17 +1,16 @@ -"""0.21: an elastic island restart resumes its round id and data cursor. CPU, no Ray. +"""0.21: an island-pool island restart resumes its round id and data cursor. CPU, no Ray. -Before 0.21 ``ElasticAvgSync.start`` always returned round 0, so a restarted island +Before 0.21 the pool session always started at round 0, so a restarted island re-drew the groups it had already trained (or the ledger refused them).""" from __future__ import annotations import pytest -from yeto.rl.elastic_client import ElasticClientConfig, ElasticIslandClient -from yeto.rl.engine.bridges import ElasticAvgSync +from yeto.rl.engine.bridges import PoolAvgSync +from pool_fakes import FakePoolSyncer as FakeSyncer from test_rl_engine_driver import _strict_config -from test_rl_inter_island_elastic_client import FakeSyncer from test_rl_restart_data_cursor import (DataSourceState, MilesLikePool, _close, _events, _records) from test_rl_reconfig_e1 import _setup @@ -19,10 +18,8 @@ def _island(tmp_path, fake, *, rounds, source, die_at=None, tape="a.jsonl", ledger=True): def factory(engine): - client = ElasticIslandClient(ElasticClientConfig(("x", 0), 0, lease_s=30.0, join_timeout_s=2), - fake.key, connect=fake.connect) - return ElasticAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=rounds, tape=tape), - client=client, base_wait_s=5.0) + return PoolAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=rounds, tape=tape), + client=fake.client(0), base_wait_s=5.0) driver, ctl, *_ = _setup(tmp_path, rounds=rounds, outer="strict-avg", sync_factory=factory, ledger=ledger) pool = MilesLikePool(driver.rollout, source, die_at=die_at) @@ -40,8 +37,7 @@ def test_restart_resumes_after_ledger_rounds_and_skips_trained_groups(tmp_path): assert first.version == 2 cursor_after_1 = {r["rollout_id"]: r for r in _records(tmp_path) if r["kind"] == "prepared"}[1]["data_cursor"] - restarted = FakeSyncer(base=list(first.base)) # the syncer kept base_version 2 - restarted.version = 2 + restarted = FakeSyncer(base=first.base.tolist(), version=2) # the syncer kept base_version 2 driver, ctl, pool = _island(tmp_path, restarted, rounds=rounds, source=DataSourceState(), tape="b.jsonl") driver.run() assert pool.seeks == [cursor_after_1] # ledger path (same as strict) @@ -55,8 +51,7 @@ def test_restart_resumes_after_ledger_rounds_and_skips_trained_groups(tmp_path): def test_empty_ledger_joining_at_base_version_advances_cursor(tmp_path): rounds = 4 - fake = FakeSyncer(base=[0.0, 0.0]) - fake.version = 2 + fake = FakeSyncer(base=[0.0, 0.0], version=2) driver, ctl, pool = _island(tmp_path, fake, rounds=rounds, source=DataSourceState()) driver.run() g = driver.sync.groups_per_round @@ -90,8 +85,7 @@ def test_fresh_run_still_starts_at_zero(tmp_path): def test_1b_regression_no_ledger_island_joining_at_base_3_advances_cursor(tmp_path): rounds = 5 - fake = FakeSyncer(base=[0.0, 0.0], catch_up=True) - fake.version = 3 + fake = FakeSyncer(base=[0.0, 0.0], version=3, catch_up=True) driver, ctl, pool = _island(tmp_path, fake, rounds=rounds, source=DataSourceState(), ledger=False) assert driver.ledger is None driver.run() diff --git a/tests/test_rl_sample_buffer.py b/tests/test_rl_sample_buffer.py index 0af088342..fe685f8fa 100644 --- a/tests/test_rl_sample_buffer.py +++ b/tests/test_rl_sample_buffer.py @@ -13,7 +13,7 @@ import torch from test_rl_engine_driver import NAME, PROFILE, _driver -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine import recovery from yeto.rl.engine.bridges import LocalOnlySync from yeto.rl.engine.fake_lifecycle import ( diff --git a/tests/test_rl_sync_membership_options.py b/tests/test_rl_sync_membership_options.py new file mode 100644 index 000000000..81067719e --- /dev/null +++ b/tests/test_rl_sync_membership_options.py @@ -0,0 +1,177 @@ +"""Membership options (--rl-sync-quorum-theta, ...) for every synced RL preset. No Ray. + +strict-avg, decoupled and full-parameter with asynchronous rounds take the +same --rl-sync-* membership options. Each set option reaches the role +syncer as the syncer flag of the same name; the islands get the pool switch, +the syncer epoch and the lease. Unset, the syncer and island commands stay +the previous ones. The full-parameter sweep, SAO, --experimental-rl-sync, +several syncer shards and critics are refused. +""" + +from __future__ import annotations + +import shlex + +import pytest + +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 + +TWO_ISLANDS = "modal:2xh100,modal:2xh100" +THREE_ISLANDS = "modal:2xh100,modal:2xh100,modal:2xh100" +DECOUPLED = ("--rl-sync-preset", "decoupled", "--fragments", "4", "--pipeline", "2", + "--local-rl-rounds-per-sync", "3") +FULL_PARAMETER = ("--rl-sync-preset", "full-parameter", "--rl-placement", "fixed-partition", + "--rl-rollout-gpus", "1", "--rl-parameter-layout-sha256", "e" * 64, + "--fragments", "4") +FULL_ASYNC = FULL_PARAMETER + ("--rl-sync-full-parameter-rounds", "async") +PRESETS = pytest.mark.parametrize("preset", [(), DECOUPLED, FULL_ASYNC], + ids=["strict-avg", "decoupled", "full-parameter-async"]) +# A count quorum below the islands lets rounds merge without an island. +QUORUM = ("--rl-sync-quorum", "2") +MEMBERSHIP = ("--rl-sync-q-min", "2", "--rl-sync-soft-deadline-s", "300", + "--rl-sync-carry-gamma", "0.5", "--rl-sync-island-lease-s", "60", + "--rl-sync-final-grace-s", "45") +SYNCER_MEMBERSHIP = (" --q-min 2 --carry-gamma 0.5 --soft-deadline-s 300 --island-lease-s 60.0" + " --syncer-epoch 0 --final-grace-s 45") +ISLAND_POOL = " --rl-island-pool --rl-syncer-epoch 0 --rl-sync-island-lease-s 60.0" + + +@pytest.fixture(autouse=True) +def _hmac_key(monkeypatch): + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") + + +def _plan(extra=(), gpu=THREE_ISLANDS): + from yeto.launcher import _prepare_rl_args, dry_run_plan + + args = _launcher_args("ports", tuple(extra), gpu=gpu) + args.controller = "local" + _prepare_rl_args(args) + return args, dry_run_plan(args) + + +def _syncer(plan) -> str: + [actor] = plan["role_syncers"] + return actor["command"] + + +def _islands(plan) -> list[str]: + return [island["learner_command"] for island in plan["island_requests"]] + + +@PRESETS +def test_unset_options_keep_every_command(preset): + _, plan = _plan(preset + QUORUM) + for flag in ("--quorum-theta", "--q-min", "--carry-gamma", "--soft-deadline-s", + "--island-lease-s", "--max-learners", "--final-grace-s", "--syncer-epoch"): + assert flag not in _syncer(plan) + assert all("--rl-island-pool" not in command and "--rl-sync-island-lease-s" not in command + for command in _islands(plan)) + + +@PRESETS +def test_membership_options_add_only_their_syncer_and_island_flags(preset): + _, baseline = _plan(preset + QUORUM) + args, plan = _plan(preset + QUORUM + MEMBERSHIP) + command = _syncer(plan) + assert SYNCER_MEMBERSHIP in command + assert command.replace(SYNCER_MEMBERSHIP, "", 1) == _syncer(baseline) + for pooled, fixed in zip(_islands(plan), _islands(baseline), strict=True): + assert ISLAND_POOL in pooled + assert pooled.replace(ISLAND_POOL, "", 1) == fixed + assert args.quorum == 2 + + +@PRESETS +def test_quorum_theta_replaces_the_count_quorum(preset): + args, plan = _plan(preset + ("--rl-sync-quorum-theta", "0.75")) + tokens = shlex.split(_syncer(plan)) + assert tokens[tokens.index("--quorum") + 1] == "1" and args.quorum == 1 + assert " --quorum-theta 0.75 --island-lease-s 150.0 --syncer-epoch 0" in _syncer(plan) + with pytest.raises(ValueError, match="replaces --rl-sync-quorum"): + _plan(preset + ("--rl-sync-quorum-theta", "0.75") + QUORUM) + + +@pytest.mark.parametrize("preset", [(), DECOUPLED], ids=["strict-avg", "decoupled"]) +def test_max_learners_reaches_the_syncer(preset): + _, plan = _plan(preset + QUORUM + ("--rl-sync-max-learners", "5")) + assert " --island-lease-s 150.0 --max-learners 5 --syncer-epoch 0" in _syncer(plan) + with pytest.raises(ValueError, match="at least the 3 islands"): + _plan(preset + QUORUM + ("--rl-sync-max-learners", "2")) + + +@PRESETS +def test_options_need_rounds_that_may_merge_without_an_island(preset): + with pytest.raises(ValueError, match="merge without some island"): + _plan(preset + ("--rl-sync-island-lease-s", "30")) + + +def test_pool_island_configs_carry_the_lease_and_epoch(): + from types import SimpleNamespace + + from yeto.rl.adapters.miles.island_entry import _island_pool_membership + + assert _island_pool_membership(SimpleNamespace(rl_island_pool=False)) == {} + assert _island_pool_membership(SimpleNamespace( + rl_island_pool=True, rl_sync_island_lease_s=60.0, rl_syncer_epoch=3)) == { + "island_lease_s": 60.0, "syncer_epoch": 3} + assert _island_pool_membership(SimpleNamespace( + rl_island_pool=True, rl_sync_island_lease_s=None, rl_syncer_epoch=None)) == { + "island_lease_s": 150.0, "syncer_epoch": 0} + + +@pytest.mark.parametrize( + ("extra", "message"), + [ + (FULL_PARAMETER + ("--rl-sync-island-lease-s", "30"), + "membership options need --rl-sync-full-parameter-rounds async"), + (FULL_ASYNC + QUORUM + ("--rl-sync-max-learners", "4"), + "--rl-sync-max-learners does not apply to full-parameter"), + (("--rl-sync-quorum-theta", "1.5"), r"in \(0, 1\]"), + (QUORUM + ("--rl-sync-q-min", "0"), "q-min must be >= 1"), + ], + ids=["full-parameter-sweep", "full-parameter-max-learners", "theta-range", "q-min"], +) +def test_refusals(extra, message): + from yeto.launcher import _prepare_rl_args, syncer_command + + args = _launcher_args("ports", tuple(extra), gpu=THREE_ISLANDS) + with pytest.raises(ValueError, match=message): + _prepare_rl_args(args) + syncer_command(args, 3) + + +@pytest.mark.parametrize( + ("fields", "critic", "message"), + [ + (dict(syncer_shards=2), False, "no --syncer-shards > 1"), + (dict(experimental_rl_sync=True), False, "apply to --rl-sync-preset strict-avg, decoupled"), + (dict(rl_sync_preset="sao-streaming-full"), False, + "apply to --rl-sync-preset strict-avg, decoupled"), + ({}, True, "do not support a critic"), + ], + ids=["syncer-shards", "experimental", "sao", "critic"], +) +def test_presets_without_a_pool_session_are_refused(monkeypatch, fields, critic, message): + from types import SimpleNamespace + + from yeto.launch import rl_args + + monkeypatch.setattr(rl_args, "rl_needs_critic", lambda args: critic) + args = SimpleNamespace(**{"training_mode": "rl", "rl_sync_island_lease_s": 30.0, + "rl_sync_preset": "strict-avg", "experimental_rl_sync": False, + "syncer_shards": None, **fields}) + with pytest.raises(ValueError, match=message): + rl_args._check_island_pool(args, num_islands=2) + + +def test_role_stream_lease_needs_asynchronous_rounds(): + from dataclasses import replace + + from test_rl_full_parameter_sync import _role_stream_config + + strict = _role_stream_config("actor", 29500) + with pytest.raises(ValueError, match="island lease needs asynchronous rounds"): + replace(strict, island_lease_s=30.0) + pooled = replace(strict, partial_quorum=True, island_lease_s=30.0, syncer_epoch=2) + assert pooled.island_lease_s == 30.0 diff --git a/tests/test_rl_sync_mode_table.py b/tests/test_rl_sync_mode_table.py index f22df9eb5..93f33ef24 100644 --- a/tests/test_rl_sync_mode_table.py +++ b/tests/test_rl_sync_mode_table.py @@ -25,7 +25,7 @@ def _rows(text: str) -> list[str]: def test_bridges_docstring_has_the_sync_mode_table(): rows = _rows(bridges.__doc__) assert rows[0] == HEADER - assert [r.split("|")[1].split()[0] for r in rows[2:]] == ["strict", "decoupled", "elastic"] + assert [r.split("|")[1].split()[0] for r in rows[2:]] == ["strict", "decoupled", "island"] def test_spec_table_matches_the_docstring(): diff --git a/tests/test_rl_sync_timing.py b/tests/test_rl_sync_timing.py index a54220c6b..995b38c27 100644 --- a/tests/test_rl_sync_timing.py +++ b/tests/test_rl_sync_timing.py @@ -1,68 +1,24 @@ -"""rl-trainer-start-stop L5.6: island-side sync timing. A fake elastic syncer -delays its ELASTIC_BASE reply; the island must report that delay as +"""rl-trainer-start-stop L5.6: island-side sync timing. A fake island-pool syncer +delays its next base; the island must report that delay as ``sync/sync_wait_seconds``. CPU, no Ray.""" from __future__ import annotations import json -import socket -import threading import time import pytest import torch -from yeto.protocol import read_frame, write_frame -from yeto.rl.elastic_client import ( - DeltaTensor, ElasticBase, ElasticClientConfig, ElasticInit, ElasticIslandClient, Join, JoinAck, - Leave, decode, encode, -) -from yeto.rl.engine.bridges import DecoupledSync, ElasticAvgSync +from pool_fakes import FakePoolSyncer +from yeto.rl.engine.bridges import DecoupledSync, PoolAvgSync from yeto.rl.engine.sync_timing import SyncTiming from test_rl_engine_driver import _driver, _engine, _strict_config -KEY = b"k" DELAY_S = 0.4 -class DelayingSyncer: - """Answers each DELTA_TENSOR with base v+1 after ``delay_s``.""" - - def __init__(self, delay_s: float): - self.a, self.b = socket.socketpair() - self.delay_s = delay_s - self.version, self.base = 0, None - threading.Thread(target=self._serve, daemon=True).start() - - def _base(self): - write_frame(self.b, *encode(ElasticBase(0, self.version, tuple(self.base)), KEY)) - - def _serve(self): - while True: - try: - t, p = read_frame(self.b, 1 << 20) - except Exception: - return - msg = decode(KEY, t, p) - if isinstance(msg, Join): - write_frame(self.b, *encode(JoinAck(0, msg.island_id, 1, 0, bytes(32), False), KEY)) - elif isinstance(msg, ElasticInit): - self.base = list(msg.params) - self._base() - elif isinstance(msg, DeltaTensor): - time.sleep(self.delay_s) # the other islands / the merge - self.version += 1 - self.base = [b + u for b, u in zip(self.base, msg.update)] - self._base() - elif isinstance(msg, Leave): - self.b.close() - return - - def connect(self, addr): - return self.a - - def test_sync_timing_split(): t = SyncTiming(submit_ts=100.0, global_ready_ts=103.0, recv_seconds=0.5) assert t.sync_wait_seconds == 2.5 @@ -73,14 +29,12 @@ def test_sync_timing_split(): "sync/sync_wait_seconds": None} -def test_elastic_island_measures_the_delayed_reply(tmp_path): - fake = DelayingSyncer(DELAY_S) +def test_pool_island_measures_the_delayed_reply(tmp_path): + fake = FakePoolSyncer(base_delay_s=DELAY_S) engine = _engine(torch.tensor([1.0, 3.0])) - client = ElasticIslandClient( - ElasticClientConfig(("x", 0), 0, lease_s=30.0, join_timeout_s=2, heartbeat_connection=False), - KEY, connect=fake.connect) - sync = ElasticAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=2), client=client, - base_wait_s=5.0) + client = fake.client() + sync = PoolAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=2), client=client, + base_wait_s=5.0) _driver(engine, sync, tmp_path).run() client.close() tape = [json.loads(l) for l in (tmp_path / "events.jsonl").read_text().splitlines()] diff --git a/tests/test_rl_synced_trainers_sampler.py b/tests/test_rl_synced_trainers_sampler.py index 61f4d0966..21a2161bf 100644 --- a/tests/test_rl_synced_trainers_sampler.py +++ b/tests/test_rl_synced_trainers_sampler.py @@ -14,11 +14,13 @@ 3. the two lanes are disjoint and no buffer position is trained twice; 4. the syncer ran R OuterSteps. -Runs when YETO_TEST_ELASTIC_SYNCER points to a syncer binary (as -test_rl_inter_island_ports.py).""" +Runs when YETO_TEST_SYNCER points to a syncer binary (as +test_rl_island_pool_e2e.py).""" from __future__ import annotations +from dataclasses import replace + import asyncio import json import os @@ -32,11 +34,10 @@ import torch from test_rl_engine_driver import NAME, _driver, _engine, _strict_config -from test_rl_inter_island_ports import _free_port +from test_rl_island_pool_e2e import _free_port from yeto.rl.adapters.miles import buffer_rollout as br from yeto.rl.adapters.miles.sampler_island import wire_policy_releases -from yeto.rl.elastic_client import ElasticClientConfig, ElasticIslandClient -from yeto.rl.engine.bridges import ElasticAvgSync +from yeto.rl.engine.bridges import PoolAvgSync from yeto.rl.engine.merged_release import BufferLane from yeto.rl.engine.policy_release import PolicyFollower, PolicyReleaseStore from yeto.rl.engine.resume import CheckpointStore @@ -109,10 +110,11 @@ def _events(path: Path, kind: str) -> list[dict]: if r.get("event") == kind] -@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), - reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") -def test_two_synced_trainers_feed_one_sampler_through_merged_releases(tmp_path): +@pytest.mark.skipif(not os.environ.get("YETO_TEST_SYNCER"), + reason="needs YETO_TEST_SYNCER (syncer binary)") +def test_two_synced_trainers_feed_one_sampler_through_merged_releases(tmp_path, monkeypatch): port, key = _free_port(), "c3-e2e" + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", key) # the islands sign their HELLO out = Path(os.environ.get("YETO_TEST_ELASTIC_OUT", tmp_path)) out.mkdir(parents=True, exist_ok=True) for old in out.glob("*.jsonl"): @@ -120,11 +122,11 @@ def test_two_synced_trainers_feed_one_sampler_through_merged_releases(tmp_path): store_root = out / "store" store_root.mkdir(exist_ok=True) syncer = subprocess.Popen( - [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", + [os.environ["YETO_TEST_SYNCER"], "--port", str(port), "--learners", "2", "--total-steps", str(ROUNDS), "--event-tape", str(out / "syncer-tape.jsonl"), - "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", - "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", - "--island-lease-s", "6", "--syncer-epoch", "5"], + "--quorum", "1", "--quorum-theta", "0.75", "--carry-gamma", "0.5", + "--soft-deadline-s", "30", "--q-min", "1", "--max-base-lag", "2", "--pipeline", "1", + "--learner-weight", "tokens2-over-steps", "--island-lease-s", "6", "--syncer-epoch", "5"], env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), stdout=open(out / "syncer.log", "w"), stderr=subprocess.STDOUT) stop, loads = threading.Event(), [] @@ -136,10 +138,10 @@ def test_two_synced_trainers_feed_one_sampler_through_merged_releases(tmp_path): sampler.start() drivers, rollouts = [], [] for i, engine in enumerate(engines): - client = ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), i, lease_s=6.0, - syncer_epoch=5), key.encode()) - cfg = _strict_config(out, engine, learner_id=i, rounds=ROUNDS, tape=f"island-{i}.jsonl") - driver = _driver(engine, ElasticAvgSync(cfg, client=client, base_wait_s=30.0), out, + cfg = replace(_strict_config(out, engine, learner_id=i, rounds=ROUNDS, tape=f"island-{i}.jsonl"), + syncer_addr=("127.0.0.1", port), wan_streams=1) + sync = PoolAvgSync(cfg, base_wait_s=30.0, syncer_epoch=5, lease_s=6.0) + driver = _driver(engine, sync, out, name=f"events-{i}.jsonl", learner_id=i) lane = BufferLane(i, 2) buffer = FileSampleBuffer(CheckpointStore(store_root), batch_size=engine.groups, lane=lane) diff --git a/tests/test_rl_trainer_lifecycle.py b/tests/test_rl_trainer_lifecycle.py index 50a566e6e..541cdcf97 100644 --- a/tests/test_rl_trainer_lifecycle.py +++ b/tests/test_rl_trainer_lifecycle.py @@ -9,7 +9,7 @@ import torch from test_rl_engine_driver import NAME, PROFILE, _driver -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine import recovery from yeto.rl.engine.bridges import LocalOnlySync from yeto.rl.engine.driver import RoundFailedError diff --git a/tests/test_rl_trainer_loop.py b/tests/test_rl_trainer_loop.py index 92a79b45b..cde3b0bfc 100644 --- a/tests/test_rl_trainer_loop.py +++ b/tests/test_rl_trainer_loop.py @@ -15,7 +15,7 @@ from yeto import runs from yeto.launch import trainer_loop as launch_loop -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine import recovery from yeto.rl.engine.bridges import LocalOnlySync from yeto.rl.engine.fake_lifecycle import ( diff --git a/tests/test_rl_trainer_loop_failures.py b/tests/test_rl_trainer_loop_failures.py index d6ffa0980..e52ee2dc1 100644 --- a/tests/test_rl_trainer_loop_failures.py +++ b/tests/test_rl_trainer_loop_failures.py @@ -17,7 +17,7 @@ from yeto import runs from yeto.launch import trainer_loop as launch_loop -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine.journal import Journal from yeto.rl.engine.trainer_lifecycle import ( BacklogStats, diff --git a/tests/test_rl_trainer_remote.py b/tests/test_rl_trainer_remote.py index c5a709e19..271c39469 100644 --- a/tests/test_rl_trainer_remote.py +++ b/tests/test_rl_trainer_remote.py @@ -29,7 +29,7 @@ modal_binding, ) from yeto.rl import event_echo -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine import recovery from yeto.rl.engine.bridges import LocalOnlySync from yeto.rl.engine.fake_lifecycle import BufferFedFakeRollout, FakeRoundCuts, FakeSampler, cuttable_fake_engine diff --git a/tests/test_rl_trainer_scheduler.py b/tests/test_rl_trainer_scheduler.py index f85a89310..dac669268 100644 --- a/tests/test_rl_trainer_scheduler.py +++ b/tests/test_rl_trainer_scheduler.py @@ -4,7 +4,7 @@ import pytest -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine.journal import Journal from yeto.rl.engine.trainer_lifecycle import ( AlwaysOnScheduler, diff --git a/tests/test_rl_trainer_sky_island.py b/tests/test_rl_trainer_sky_island.py index 3a76f8c2d..242442716 100644 --- a/tests/test_rl_trainer_sky_island.py +++ b/tests/test_rl_trainer_sky_island.py @@ -31,7 +31,7 @@ sky_binding, ) from yeto.rl import event_echo -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine import recovery from yeto.rl.engine.fake_lifecycle import FakeSampler from yeto.rl.engine.journal import Journal diff --git a/tests/test_rl_verl_round_cuts.py b/tests/test_rl_verl_round_cuts.py index 6653eaa6a..475f4713f 100644 --- a/tests/test_rl_verl_round_cuts.py +++ b/tests/test_rl_verl_round_cuts.py @@ -17,7 +17,7 @@ from yeto.rl.adapters.verl import cut as vcut from yeto.rl.adapters.verl import round_cuts -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.rl.engine import recovery, trainer_remote from yeto.rl.engine.algorithm import AlgorithmSpec from yeto.rl.engine.bridges import LocalOnlySync diff --git a/tests/test_rl_verl_synced_trainers.py b/tests/test_rl_verl_synced_trainers.py index 85547c6db..542a81830 100644 --- a/tests/test_rl_verl_synced_trainers.py +++ b/tests/test_rl_verl_synced_trainers.py @@ -4,7 +4,7 @@ The trainers are the verl buffer trainer of test_rl_verl_buffer_trainer.py (fake trainer actor; real driver, fully_async rollout port, feed, publisher, -merged checkpoint) over one lane each of one buffer, with ``ElasticAvgSync``. +merged checkpoint) over one lane each of one buffer, with ``PoolAvgSync``. The sampler is the real engine sampler island (fake engine) in a thread. Checks (the same as the GPU run's, SESSION25/gpu-review/T-C3V.md): @@ -15,7 +15,7 @@ 3. the two lanes are disjoint and no buffer position is trained twice; 4. the syncer ran R OuterSteps. -Runs when YETO_TEST_ELASTIC_SYNCER points to a syncer binary (as +Runs when YETO_TEST_SYNCER points to a syncer binary (as test_rl_synced_trainers_sampler.py).""" from __future__ import annotations @@ -30,7 +30,7 @@ import pytest import torch -from test_rl_inter_island_ports import _free_port +from test_rl_island_pool_e2e import _free_port from test_rl_verl_buffer_trainer import NAMES, FakeBufferActor, _base from yeto.rl.adapters.sampler_engines.fake import FakeEngine from yeto.rl.adapters.verl import buffer_feed as bf @@ -62,8 +62,7 @@ def _island(out: Path, store_root: Path, base: Path, learner_id: int, port: int, from yeto.rl.adapters.verl.ports_impl import VerlPolicyState from yeto.rl.bridge import BridgeConfig from yeto.rl.core import CanonicalTensorSpec, canonical_layout_hash - from yeto.rl.elastic_client import ElasticClientConfig, ElasticIslandClient - from yeto.rl.engine.bridges import ElasticAvgSync + from yeto.rl.engine.bridges import PoolAvgSync from yeto.rl.engine.driver import EventTape, IslandDriver specs = (CanonicalTensorSpec(NAMES[0], (2, 4), "float32", 8), @@ -86,35 +85,34 @@ def _island(out: Path, store_root: Path, base: Path, learner_id: int, port: int, groups_per_round=GROUPS, samples_per_group=2, local_optimizer_steps=1, expected_specs=specs, base_model_revision="a" * 40, lora_config_hash="b" * 64, layout_hash=canonical_layout_hash(specs), event_tape=str(out / f"bridge-{learner_id}.jsonl"), - wan_streams=0) - client = ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), learner_id, lease_s=6.0, - syncer_epoch=5), key.encode()) + wan_streams=1) profile = far.execution_profile_for(TOLERANT, MAX_POLICY_AGE, groups_per_round=GROUPS, samples_per_group=2, sync="elastic") driver = IslandDriver( learner_id=learner_id, rollout=rollout, trainer=FullyAsyncTrainerGroup(island), policy_state=VerlPolicyState(island), publisher=publisher, placement=placement, capabilities=verl_capabilities("sha256:" + "0" * 64, fully_async_limit=MAX_POLICY_AGE), - algorithm=TOLERANT, sync=ElasticAvgSync(bridge, client=client, base_wait_s=30.0, syncer_epoch=5), + algorithm=TOLERANT, sync=PoolAvgSync(bridge, base_wait_s=30.0, syncer_epoch=5, lease_s=6.0), events=tape, profile=profile) publisher.bind(driver) return driver, island -@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), - reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") -def test_two_synced_verl_trainers_feed_one_engine_sampler(tmp_path): +@pytest.mark.skipif(not os.environ.get("YETO_TEST_SYNCER"), + reason="needs YETO_TEST_SYNCER (syncer binary)") +def test_two_synced_verl_trainers_feed_one_engine_sampler(tmp_path, monkeypatch): port, key = _free_port(), "c3v-e2e" + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", key) # the islands sign their HELLO out = tmp_path store_root = out / "store" store = CheckpointStore(store_root) base, _ = _base(tmp_path, dtype=torch.float32) syncer = subprocess.Popen( - [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", + [os.environ["YETO_TEST_SYNCER"], "--port", str(port), "--learners", "2", "--total-steps", str(ROUNDS), "--event-tape", str(out / "syncer-tape.jsonl"), - "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", - "--soft-deadline-s", "30", "--q-min", "1", "--max-carry-lag", "2", - "--island-lease-s", "6", "--syncer-epoch", "5"], + "--quorum", "1", "--quorum-theta", "0.75", "--carry-gamma", "0.5", + "--soft-deadline-s", "30", "--q-min", "1", "--max-base-lag", "2", "--pipeline", "1", + "--learner-weight", "tokens2-over-steps", "--island-lease-s", "6", "--syncer-epoch", "5"], env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), stdout=open(out / "syncer.log", "w"), stderr=subprocess.STDOUT) done, sampler_events = threading.Event(), [] diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index b4b5e9872..c7152be8e 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -242,7 +242,7 @@ def test_forward_tables_name_flags_of_both_layers(): with pytest.raises(KeyError): rl_forward_pairs("rl_island_bundle_map") # island-only: the CLI cannot forward it with pytest.raises(KeyError): - rl_forward_pairs("rl_quorum_theta") # CLI-only: goes to the syncer, not the island + rl_forward_pairs("rl_sync_quorum_theta") # CLI-only: goes to the syncer, not the island # Flag literals in the launcher's command builders that are NOT island or launch @@ -252,13 +252,13 @@ def test_forward_tables_name_flags_of_both_layers(): "--abbrev-ref", "--adapter-dir", "--add", "--address", "--allow-unattested-legacy-adapter", "--allow-unauthenticated-islands", "--async-checkpoint", "--bcast-delta-full-every", "--bind", "--cap-add", "--carry-gamma", "--checkpoint-every", "--checkpoint-path", "--codex-harness-binary", "--config-json", - "--create", "--dashboard-host", "--dearmor", "--depth", "--detach", "--dtype", + "--create", "--dashboard-host", "--max-learners", "--dearmor", "--depth", "--detach", "--dtype", "--entrypoint", "--expected-island-contract", "--expert-token-weights", "--final-ack-timeout-s", "--final-grace-s", "--follow", "--force", "--format", "--get", "--global", "--gpus", "--head", "--ignore-installed", "--include-dashboard", "--index-url", "--ipc", "--island-lease-s", - "--island-scheduling-mode", "--json", "--learner-weight", "--learners", "--level", + "--json", "--learner-weight", "--learners", "--level", "--manifest-path", "--mark-final-checkpoint", "--master", "--matrix-merge", - "--max-base-lag", "--max-carry-lag", "--shard-count", "--shard-index", "--weighted-fold", "--name", "--net", "--nnodes", + "--max-base-lag", "--shard-count", "--shard-index", "--weighted-fold", "--name", "--net", "--nnodes", "--no-build-isolation", "--no-checkout", "--no-deps", "--no-follow", "--node", "--node-ip-address", "--nproc", "--num-gpus", "--num-inference-steps", "--output-dir", "--password-stdin", "--policy-sweep-fragments", "--porcelain", "--port", "--profile", diff --git a/tests/test_secret_handling_hardening.py b/tests/test_secret_handling_hardening.py index a9def9dcb..d371978f4 100644 --- a/tests/test_secret_handling_hardening.py +++ b/tests/test_secret_handling_hardening.py @@ -82,11 +82,11 @@ def _stop(process): return process.communicate(timeout=10)[0] -def test_legacy_syncer_without_key_refuses_to_start(syncer_binary, tmp_path): +def test_syncer_without_key_refuses_to_start(syncer_binary, tmp_path): process = _syncer(syncer_binary, _port(), tmp_path) out = process.communicate(timeout=30)[0] assert process.returncode != 0 - assert "legacy mode requires an island HMAC key" in out + assert "the syncer requires an island HMAC key" in out def test_legacy_hello_with_right_key_is_admitted_wrong_key_refused(syncer_binary, tmp_path): @@ -147,12 +147,12 @@ def test_bad_pinned_contract_value_is_refused(syncer_binary, tmp_path): def _rl_args(**kw): import argparse - base = dict(training_mode="rl", rl_island_scheduling="legacy", rl_island_contract_sha256=None) + base = dict(training_mode="rl", rl_island_contract_sha256=None) base.update(kw) return argparse.Namespace(**base) -def test_launcher_generates_one_legacy_key_per_run(monkeypatch): +def test_launcher_generates_one_key_per_run(monkeypatch): from yeto import launcher monkeypatch.delenv("YETO_ISLAND_HMAC_KEY", raising=False) @@ -163,14 +163,6 @@ def test_launcher_generates_one_legacy_key_per_run(monkeypatch): assert launcher.island_hmac_secret(_rl_args(training_mode="sft")) == {} -def test_launcher_elastic_still_needs_a_user_key(monkeypatch): - from yeto import launcher - - monkeypatch.delenv("YETO_ISLAND_HMAC_KEY", raising=False) - with pytest.raises(ValueError, match="YETO_ISLAND_HMAC_KEY"): - launcher.island_hmac_secret(_rl_args(rl_island_scheduling="elastic")) - - def test_launcher_passes_the_pinned_contract_to_the_syncer(): from yeto import launcher diff --git a/tests/test_spot_droppable.py b/tests/test_spot_droppable.py index 3f3662c1f..1fe276cac 100644 --- a/tests/test_spot_droppable.py +++ b/tests/test_spot_droppable.py @@ -14,7 +14,7 @@ def _args(**kw): - base = dict(training_mode="rl", rl_island_scheduling="elastic", spot=False, rl_q_min=None, + base = dict(training_mode="rl", rl_sync_island_lease_s=150.0, spot=False, rl_sync_q_min=None, rl_checkpoint_store=None, rl_island_role=["1:droppable", "2:droppable"]) base.update(kw) return SimpleNamespace(**base) @@ -36,11 +36,14 @@ def test_no_roles_is_a_noop(): assert dr.manifest_fields({}) == {} -def test_legacy_refused_with_spec_reason(): - with pytest.raises(ValueError, match="可丢弃岛只在 elastic 下支持"): - dr.admit(_args(rl_island_scheduling=None), ["nebius", "aws", "aws"]) - with pytest.raises(ValueError, match="可丢弃岛只在 elastic 下支持"): - dr.admit(_args(rl_island_scheduling="legacy"), ["nebius", "aws", "aws"]) +def test_fixed_roster_refused_with_spec_reason(): + with pytest.raises(ValueError, match="可丢弃岛只在岛池下支持"): + dr.admit(_args(rl_sync_island_lease_s=None), ["nebius", "aws", "aws"]) + + +def test_full_parameter_islands_cannot_be_droppable(): + with pytest.raises(ValueError, match="full-parameter island cannot join a running pool"): + dr.admit(_args(rl_sync_preset="full-parameter"), ["nebius", "aws", "aws"]) def test_needs_an_anchor(): @@ -59,9 +62,9 @@ def test_global_spot_refused(): def test_q_min_must_fit_the_anchors(): - with pytest.raises(ValueError, match="--rl-q-min 2 > 1 anchor"): - dr.admit(_args(rl_q_min=2), ["nebius", "aws", "aws"]) - assert dr.admit(_args(rl_q_min=2, rl_island_role=["2:droppable"]), ["nebius", "nebius", "aws"]) + with pytest.raises(ValueError, match="--rl-sync-q-min 2 > 1 anchor"): + dr.admit(_args(rl_sync_q_min=2), ["nebius", "aws", "aws"]) + assert dr.admit(_args(rl_sync_q_min=2, rl_island_role=["2:droppable"]), ["nebius", "nebius", "aws"]) def test_verda_anchor_cut_not_on_local_volume(): @@ -229,30 +232,6 @@ def test_droppable_rejoin_takes_the_current_base_with_zero_weight(): # --- 5.2 CPU part: after a reclaim LEAVE the still-running island must not rejoin ----------- -def test_client_after_leave_never_rejoins_on_not_member_errors(): - """The syncer answers a left island's delta/heartbeat with "not a member - (rejoin required)". The droppable island LEAVEs on reclaim and is killed by - the cloud later; until then it must not JOIN again by itself (the rejoin - comes from the relaunched island, catch_up from the syncer base).""" - from yeto.rl.elastic_client import JoinAck, ElasticClientConfig, ElasticIslandClient - - sent = [] - c = ElasticIslandClient(ElasticClientConfig(syncer_addr=("x", 1), island_id=2), b"k") - c.ack = JoinAck(1, 2, 3, 1, bytes(32), False) - c.sock = object() - c._send = sent.append - c.leave() - assert type(sent[-1]).__name__ == "Leave" - c._on_error("2 is not a member (rejoin required)") - assert not c._rejoin.is_set() - assert c.rejoin() is True and c.rejoins == 0 - c.check() - assert [type(m).__name__ for m in sent] == ["Leave"] - - -# --- 5.2: the Modal reclaim signal reaches the learner (grandchild) --------------------- - - def test_island_script_forwards_reclaim_signal_to_group_on_droppable_only(): import json as _json import signal as _signal diff --git a/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py index be6ecc85d..1429e8b92 100644 --- a/tests/test_tools_eval_efficiency.py +++ b/tests/test_tools_eval_efficiency.py @@ -462,7 +462,7 @@ def test_e11_arms_read_one_row_set(): assert f"--rl-allow-local-data --data {spec.YETO_TASK_FILE}" in single and "--gpu modal:4xh100" in single multi = spec.yeto_miles_islands_launcher_args(e11, 2, "e11-y", data="t.jsonl", island_tasks="t.json") assert "--gpu modal:2xh100,modal:2xh100" in multi and "--rollout-batch-size 16" in multi - assert "--rl-rollout-gpus 1" in multi and "--controller head" in multi and "--rl-island-scheduling elastic" in multi + assert "--rl-rollout-gpus 1" in multi and "--controller head" in multi and "--rl-sync-quorum-theta 0.75" in multi and "--rl-sync-island-lease-s 150" in multi split = spec.island_task_split(e11.groups_per_round, e11.rounds, 2) ids = sorted(int(t[3:]) for v in split.values() for t in v) assert ids == list(range(1, e11.train_rows + 1)) diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py index 0efc0a138..c82db98ce 100644 --- a/tools/eval/efficiency/spec.py +++ b/tools/eval/efficiency/spec.py @@ -545,9 +545,9 @@ def yeto_miles_islands_launcher_args(spec: EffSpec, islands: int, run_id: str, * "--no-sglang-deterministic-inference --rl-lr-schedule constant --rl-activation-recompute " f"--rl-max-tokens-per-gpu {max_tokens_per_gpu(spec)} --rl-gradient-accumulation-fusion " f"--trust-remote-code --rl-algorithm-spec {MILES_ALGORITHM_SPEC} " - # elastic: the multi-island LoRA path with real-card runs (S24 d4, S25 c3); legacy - # scheduling hung in the first apply on both islands (e11-yeto-2i-1, no round in 13 min) - "--rl-island-scheduling elastic " + # island pool: the multi-island LoRA path with real-card runs (S24 d4, S25 c3); a + # fixed roster hung in the first apply on both islands (e11-yeto-2i-1, no round in 13 min) + "--rl-sync-quorum-theta 0.75 --rl-sync-carry-gamma 0.5 --rl-sync-soft-deadline-s 900 --rl-sync-max-base-lag 2 --rl-sync-learner-weight tokens2-over-steps --rl-sync-island-lease-s 150 " f"--rl-placement fixed-partition --rl-rollout-gpus {per // 2} " f"--rl-max-policy-age {spec.max_policy_age}") if kwargs: diff --git a/tools/eval/miles_head.sh b/tools/eval/miles_head.sh index dac653735..61769544d 100755 --- a/tools/eval/miles_head.sh +++ b/tools/eval/miles_head.sh @@ -30,7 +30,7 @@ LORA="--tuning lora --lora-r 16" MODALX="--modal-gpu-exact --modal-retries 0 --modal-cpu 8 --modal-memory 64 --modal-tape-volume $TAPEVOL" STORE="--rl-checkpoint-store modal-volume://$STOREVOL/$P" COMMON="--training-mode rl --rl-engine ports $MODALX $MODEL $DATA $LORA $STORE --rl-algorithm-spec examples/rl_algorithms/decoupled-ppo-c2.json --rl-max-policy-age 2 --rl-observe-timeline --reward-function gsm8k_reward:score --fragments 1 --pipeline 1 --rollout-batch-size 32 --n-samples-per-prompt 4 --rollout-max-response-len 1024 --seq-len 2048 --inner-lr 1e-5 --seed 17 --apply-chat-template-kwargs '{\"enable_thinking\": false}' --trust-remote-code" -TARGS="launch --controller head --syncer-region $HEAD_REGION --preflight-threads off --rl-stall-timeout 2400 --on-demand --no-island-relaunch --modal-timeout-s $HARD --cluster-prefix $P --gpu modal:1x$GPU,modal:1x$GPU --rl-island-scheduling elastic --rl-sample-role trainer --rl-sample-wait-s 900 $COMMON --total-steps $STEPS" +TARGS="launch --controller head --syncer-region $HEAD_REGION --preflight-threads off --rl-stall-timeout 2400 --on-demand --no-island-relaunch --modal-timeout-s $HARD --cluster-prefix $P --gpu modal:1x$GPU,modal:1x$GPU --rl-sync-quorum-theta 0.75 --rl-sync-carry-gamma 0.5 --rl-sync-soft-deadline-s 900 --rl-sync-max-base-lag 2 --rl-sync-learner-weight tokens2-over-steps --rl-sync-island-lease-s 150 --rl-sample-role trainer --rl-sample-wait-s 900 $COMMON --total-steps $STEPS" SARGS="launch --preflight-threads off --controller local --cluster-prefix $SP --gpu modal:1x$GPU --rl-single-island-no-sync --rl-sample-role sampler --rl-sampler-max-backlog-rounds 4 $COMMON --total-steps $SROUNDS" git -C $REPO rev-parse $SHA > $R/yeto_sha.txt; echo "$TARGS" > $R/args.txt; echo "$SARGS" > $R/sargs.txt echo "$APP $SAPP $HEAD_CL" > $R/resources.txt diff --git a/tools/eval/modal_verl.py b/tools/eval/modal_verl.py index 351bb1103..e5402c30b 100644 --- a/tools/eval/modal_verl.py +++ b/tools/eval/modal_verl.py @@ -55,7 +55,7 @@ # then stop STOPPED), printed island environment as JSON (S25 C4 modal_c4.py) COMMAND_WRITER = ( "import json, sys\n" - "from yeto.rl.elastic_client import LeaveReason\n" + "from yeto.rl.pool_client import LeaveReason\n" "from yeto.rl.engine import trainer_remote as t\n" "from yeto.rl.engine.resume import store_for\n" "w = t.CommandWriter(t.LocalStoreFiles(store_for(sys.argv[1])), 0)\n" @@ -131,9 +131,10 @@ def syncer(run: str, cfg: dict) -> dict: began = time.time() argv = ["/opt/yeto-syncer", "--port", str(SYNCER_PORT), "--learners", str(cfg["learners"]), "--total-steps", str(cfg["rounds"]), "--event-tape", "/tmp/syncer-tape.jsonl", - "--island-scheduling-mode", "elastic", "--quorum-theta", "0.75", "--carry-gamma", "0.5", - "--soft-deadline-s", "600", "--q-min", "1", "--max-carry-lag", "2", "--syncer-epoch", "0", - "--island-lease-s", "150.0", "--final-grace-s", "300"] + "--quorum", "1", "--quorum-theta", "0.75", "--carry-gamma", "0.5", + "--soft-deadline-s", "600", "--q-min", "1", "--max-base-lag", "2", "--syncer-epoch", "0", + "--island-lease-s", "150.0", "--final-grace-s", "300", "--pipeline", "1", + "--learner-weight", "tokens2-over-steps"] meta: dict[str, Any] = {"argv": argv, "container": os.environ.get("MODAL_TASK_ID")} try: with modal.forward(SYNCER_PORT, unencrypted=True) as tunnel, (out / "syncer.log").open("w") as log: @@ -306,8 +307,8 @@ def trainer(run: str, cfg: dict, learner_id: int) -> dict: time.sleep(2.0) if not address: raise RuntimeError("no syncer address within 1200 s") - argv += ["--syncer", address, "--rl-island-scheduling", "elastic", "--rl-syncer-epoch", "0", - "--rl-island-lease-s", "150.0", "--rl-sample-store", f"{STORE}/{run}", + argv += ["--syncer", address, "--rl-island-pool", "--rl-syncer-epoch", "0", + "--rl-sync-island-lease-s", "150.0", "--rl-sample-store", f"{STORE}/{run}", "--rl-sample-lanes", str(cfg["learners"])] else: argv += ["--rl-single-island-no-sync", "--rl-resume-store", f"{STORE}/{run}"] diff --git a/tools/probes/c4_verl/modal_c4.py b/tools/probes/c4_verl/modal_c4.py index 5b0bbfd21..ccb519a2b 100644 --- a/tools/probes/c4_verl/modal_c4.py +++ b/tools/probes/c4_verl/modal_c4.py @@ -94,7 +94,7 @@ def island(run: str, phase: str, gpu: str) -> dict: if phase == "a": # the head's commands for incarnation 0, written through yeto's own writer rc = _sh(f"{PY} -c " + shlex.quote( "import sys; from yeto.rl.engine import trainer_remote as t; " - "from yeto.rl.engine.resume import store_for; from yeto.rl.elastic_client import LeaveReason; " + "from yeto.rl.engine.resume import store_for; from yeto.rl.pool_client import LeaveReason; " "w = t.CommandWriter(t.LocalStoreFiles(store_for(sys.argv[1])), 0); " "w.send(t.RUN, rounds=2); w.send(t.STOP, reason=LeaveReason.STOPPED); " "print(t.island_env(0, None, poll_s=2.0))") + f" {store}", log) diff --git a/yeto/cloud/droppable.py b/yeto/cloud/droppable.py index abd1e3759..4415334cd 100644 --- a/yeto/cloud/droppable.py +++ b/yeto/cloud/droppable.py @@ -7,11 +7,11 @@ move the syncer base) and the cut point (its checkpoint store, when one is given). * A droppable island runs on spot. It never resumes from a checkpoint store. When it comes back after a reclaim it starts fresh, JOINs with catch-up and takes the - current syncer base (``ElasticAvgSync.start``: optimizer reset); the data cursor + current syncer base (``PoolAvgSync.start``: optimizer reset); the data cursor comes from the syncer version (0.21 fresh-machine rule). So it syncs from the line of weights the anchor keeps current, never from an archive. * If every droppable island is gone at once, the anchor alone still closes rounds: - admission requires ``--rl-q-min`` <= number of anchors. + admission requires ``--rl-sync-q-min`` <= number of anchors. Admission (task 5.1, spec "推理型岛的定义"): elastic only (legacy refused), at least one anchor, no anchor on a cloud whose GPUs are always preemptible (Modal), no global @@ -34,7 +34,7 @@ ROLE_ENV = "YETO_ISLAND_ROLE" # JSON {"role", "cloud", "region"} on every island of a role run ANCHOR, DROPPABLE = "anchor", "droppable" ROLES = (ANCHOR, DROPPABLE) -LEGACY_REFUSED = "可丢弃岛只在 elastic 下支持 (droppable islands need --rl-island-scheduling elastic)" +LEGACY_REFUSED = "可丢弃岛只在岛池下支持 (droppable islands need island pool options, e.g. --rl-sync-island-lease-s)" def parse_roles(items, num_islands: int) -> dict[int, str]: @@ -69,9 +69,14 @@ def admit(args, clouds: list[str]) -> dict[int, str]: return {} if getattr(args, "training_mode", "sft") != "rl": raise ValueError("--rl-island-role needs --training-mode rl") - mode = getattr(args, "rl_island_scheduling", None) or "legacy" - if mode != "elastic": + from ..launch.syncer import island_pool + + if not island_pool(args): raise ValueError(LEGACY_REFUSED) + if getattr(args, "rl_sync_preset", None) == "full-parameter": + # A droppable island comes back as a fresh process after a reclaim. + raise ValueError("--rl-island-role: a full-parameter island cannot join a running pool " + "after the initial cut yet, so it cannot be droppable") roles = parse_roles(items, len(clouds)) anchors = [i for i, r in roles.items() if r == ANCHOR] if not anchors: @@ -79,10 +84,10 @@ def admit(args, clouds: list[str]) -> dict[int, str]: if getattr(args, "spot", False): raise ValueError("--rl-island-role with --spot: billing comes from the role " "(anchor on-demand, droppable spot); drop --spot") - q_min = getattr(args, "rl_q_min", None) + q_min = getattr(args, "rl_sync_q_min", None) q_min = 1 if q_min is None else int(q_min) if q_min > len(anchors): - raise ValueError(f"--rl-q-min {q_min} > {len(anchors)} anchor island(s): rounds would stop " + raise ValueError(f"--rl-sync-q-min {q_min} > {len(anchors)} anchor island(s): rounds would stop " "when every droppable island is reclaimed at once") store = str(getattr(args, "rl_checkpoint_store", None) or "") for i, role in roles.items(): diff --git a/yeto/island_overrides.py b/yeto/island_overrides.py index d3360376e..6e67729a5 100644 --- a/yeto/island_overrides.py +++ b/yeto/island_overrides.py @@ -22,7 +22,7 @@ (10-09 ruling of the main agent for the user: the original ``--rl-resume`` / ``yeto export`` do not exist on main; resume is ``--rl-checkpoint-store``, -export is ``yeto merge``. ``--rl-max-carry-lag`` is a syncer parameter and +export is ``yeto merge``. ``--rl-sync-max-base-lag`` is a syncer parameter and cannot be overridden per island.) This module is imported by the launcher AND by the island (env helpers), so it @@ -79,7 +79,7 @@ def _salt(value: str) -> str: "rl_max_policy_age": (_age, "落后上限"), "identity_test_salt": (_salt, "岛身份测试扰动"), } -SYNCER_ONLY = {"rl_max_carry_lag": "--rl-max-carry-lag"} +SYNCER_ONLY = {"rl_sync_max_base_lag": "--rl-sync-max-base-lag"} def _norm_key(key: str) -> str: @@ -134,8 +134,9 @@ def _old_value(args, key: str): if key == "rl_lr_schedule": value = value or "auto" if value == "auto": # resolved as run_config.resolve_lr_schedule does - mode = getattr(args, "rl_island_scheduling", None) or "legacy" - run_until_stop = getattr(args, "sync_preset", None) == "decoupled" or mode == "elastic" + from .launch.syncer import island_pool + + run_until_stop = getattr(args, "sync_preset", None) == "decoupled" or island_pool(args) return "constant" if run_until_stop else "linear" return value if key == "rl_max_policy_age": @@ -154,11 +155,11 @@ def check_island_args(args) -> None: sched = getattr(args, "rl_lr_schedule", None) or "auto" _lr(sched) if sched == "linear": - from .launcher import island_scheduling_mode + from .launcher import island_pool - if getattr(args, "sync_preset", None) == "decoupled" or island_scheduling_mode(args) == "elastic": + if getattr(args, "sync_preset", None) == "decoupled" or island_pool(args): raise ValueError("--rl-lr-schedule linear needs a known local step count; decoupled and " - "elastic islands run until the syncer stops them (use auto or constant)") + "island pool islands run until the syncer stops them (use auto or constant)") from .launcher import _rl_backend_module from .rl.engine.policy_age import validate_limit diff --git a/yeto/launch/fleet.py b/yeto/launch/fleet.py index 4bf7ce154..7bd57c72d 100644 --- a/yeto/launch/fleet.py +++ b/yeto/launch/fleet.py @@ -7,16 +7,15 @@ import threading from dataclasses import dataclass -from .syncer import island_scheduling_mode +from .syncer import island_pool from .sky_ops import SYNCER_REFUSAL_MARKERS def controller_fixed_roster(args) -> bool: - """FleetController ``fixed_roster``: RL runs whose island scheduling is NOT - elastic. Elastic membership is dynamic: an island that fails for good is - dropped and the pool goes on (10-09 ruling of the main agent for the user; - before, every RL run, elastic included, was a fixed roster).""" - return getattr(args, "training_mode", "sft") == "rl" and island_scheduling_mode(args) != "elastic" + """FleetController ``fixed_roster``: RL runs without island pool options. + A pool's membership is dynamic: an island that fails for good is dropped + and the pool goes on (10-09 ruling of the main agent for the user).""" + return getattr(args, "training_mode", "sft") == "rl" and not island_pool(args) # FleetController states. diff --git a/yeto/launch/peer_recovery.py b/yeto/launch/peer_recovery.py index 86f373303..19f97e1e4 100644 --- a/yeto/launch/peer_recovery.py +++ b/yeto/launch/peer_recovery.py @@ -42,7 +42,7 @@ def peer_recovery_store(args) -> str | None: def check_peer_recovery(args) -> None: """Refuse ``--rl-peer-recovery`` options that do not apply. Every recovery option without ``--rl-peer-recovery`` is refused too.""" - from .syncer import island_scheduling_mode, rl_full_parameter_async, syncer_shard_count + from .syncer import island_pool, rl_full_parameter_async, syncer_shard_count options = [flag for dest, flag in ( ("rl_recovery_bound_steps", "--rl-recovery-bound-steps"), @@ -58,8 +58,9 @@ def check_peer_recovery(args) -> None: if not rl_full_parameter_async(args): raise ValueError("--rl-peer-recovery needs --rl-sync-preset full-parameter with " "--rl-sync-full-parameter-rounds async") - if island_scheduling_mode(args) != "legacy": - raise ValueError("--rl-peer-recovery needs --rl-island-scheduling legacy") + if island_pool(args): + # Both set the syncer epoch; recovery is only wired for a fixed roster. + raise ValueError("--rl-peer-recovery does not take island pool options yet") if syncer_shard_count(args) > 1: # Recovery is one exchange with one syncer; shards would each need it. raise ValueError("--rl-peer-recovery needs one syncer process (--syncer-shards 1)") diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index a681ec1eb..7fa3d9a7c 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -20,7 +20,7 @@ rl_island_count, rl_island_placement_args, rl_island_spec, rl_min_nodes, rl_trainer_shape, ) from .syncer import ( - ROLE_SYNCERS, island_scheduling_mode, rl_full_parameter_async, rl_needs_critic, + ROLE_SYNCERS, island_pool, rl_full_parameter_async, rl_needs_critic, rl_syncer_profile, rl_trainable_roles, syncer_shard_count, ) from .storage import _rl_checkpoint_mount @@ -99,8 +99,8 @@ def _prepare_sao_streaming(args, *, num_islands: int) -> None: problems.append("--rl-sao-contracts ") if args.experimental_rl_sync: problems.append("no --experimental-rl-sync") - if island_scheduling_mode(args) != "legacy": - problems.append("--rl-island-scheduling legacy (no fragment frames in elastic mode)") + if island_pool(args): + problems.append("a fixed island roster (no island pool options)") if rl_trainable_roles(args) != ("actor", "critic"): problems.append("an algorithm with a critic (the SAO spec trains actor and critic)") if getattr(args, "rl_single_island_no_sync", False): @@ -185,9 +185,9 @@ def _check_bcast_delta(args) -> None: if args.rl_sync_preset not in ("strict-avg", FULL_PARAMETER_SYNC_PRESET): raise ValueError("--rl-sync-bcast-delta-full-every only applies to --rl-sync-preset " f"strict-avg or {FULL_PARAMETER_SYNC_PRESET}") - if island_scheduling_mode(args) != "legacy" or args.experimental_rl_sync: - raise ValueError("--rl-sync-bcast-delta-full-every needs --rl-island-scheduling legacy " - "and no --experimental-rl-sync") + if island_pool(args) or args.experimental_rl_sync: + raise ValueError("--rl-sync-bcast-delta-full-every needs a fixed island roster (no island " + "pool options) and no --experimental-rl-sync") # S27 P2: asynchronous rounds keep at most two fragment rounds in flight. @@ -219,8 +219,9 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: else "--local-rl-rounds-per-sync >= 1") if args.fragments < 1: problems.append("--fragments >= 1") - if island_scheduling_mode(args) != "legacy": - problems.append("--rl-island-scheduling legacy (no fragment frames in elastic mode)") + if island_pool(args) and not async_rounds: + problems.append("a fixed island roster (membership options need " + "--rl-sync-full-parameter-rounds async)") if rl_needs_critic(args): problems.append("a critic-free algorithm (the full-parameter preset trains the actor only)") fragment_pipeline = int(getattr(args, "rl_sync_fragment_pipeline", 1)) @@ -316,8 +317,7 @@ def _apply_rl_outer_merge_options(args, *, num_islands: int) -> None: Every option left at its default keeps the preset's value, so the syncer command and semantic profile are unchanged. The full-parameter preset takes them only with asynchronous rounds (S27 P2; the sweep keeps its strict - roster); SAO and --experimental-rl-sync take no overrides, and the elastic - scheduler has its own quorum and weights.""" + roster); SAO and --experimental-rl-sync take no overrides.""" changed = [option for dest, option, default in _RL_OUTER_MERGE_OPTIONS if getattr(args, dest, default) != default] if not changed: @@ -333,8 +333,6 @@ def _apply_rl_outer_merge_options(args, *, num_islands: int) -> None: if args.experimental_rl_sync: raise ValueError(f"{where} do not apply with --experimental-rl-sync " "(it takes the generic syncer flags)") - if island_scheduling_mode(args) != "legacy": - raise ValueError(f"{where} need --rl-island-scheduling legacy") quorum = args.rl_sync_quorum if quorum is not None: if not 1 <= quorum <= num_islands: @@ -363,6 +361,50 @@ def _apply_rl_outer_merge_options(args, *, num_islands: int) -> None: f"or {FULL_PARAMETER_SYNC_PRESET} with asynchronous rounds") +def _check_island_pool(args, *, num_islands: int) -> None: + """Membership options (--rl-sync-quorum-theta, ...; dynamic membership) + apply to the strict-avg, decoupled and asynchronous full-parameter + presets without a critic. The rounds must be able to merge without some + island: a capacity quorum (--rl-sync-quorum-theta), a count quorum below + the islands, a grace window or a carried base lag (the syncer refuses + membership options with a strict fixed roster). + --rl-sync-quorum-theta replaces the count quorum, so the syncer's + --quorum is 1.""" + if not island_pool(args): + if getattr(args, "rl_syncer_epoch", None) is not None and not getattr(args, "rl_peer_recovery", False): + raise ValueError("--rl-syncer-epoch needs island pool options (e.g. --rl-sync-island-lease-s) " + "or --rl-peer-recovery") + return + preset = args.rl_sync_preset + if args.experimental_rl_sync or not ( + preset in ("strict-avg", "decoupled") or rl_full_parameter_async(args)): + raise ValueError("membership options (--rl-sync-quorum-theta, --rl-sync-island-lease-s, " + "...) apply to --rl-sync-preset strict-avg, decoupled, or " + f"{FULL_PARAMETER_SYNC_PRESET} with --rl-sync-full-parameter-rounds async") + if rl_needs_critic(args): + raise ValueError("membership options do not support a critic yet") + if int(getattr(args, "syncer_shards", None) or 1) > 1: + raise ValueError("membership options need one syncer (no --syncer-shards > 1)") + max_learners = getattr(args, "rl_sync_max_learners", None) + if max_learners is not None: + if preset == FULL_PARAMETER_SYNC_PRESET: + raise ValueError(f"--rl-sync-max-learners does not apply to {FULL_PARAMETER_SYNC_PRESET}: " + "its island session binds the launched roster") + if int(max_learners) < num_islands: + raise ValueError(f"--rl-sync-max-learners must be at least the {num_islands} islands") + if getattr(args, "rl_sync_quorum_theta", None) is not None: + if getattr(args, "rl_sync_quorum", None) is not None: + raise ValueError("--rl-sync-quorum-theta replaces --rl-sync-quorum; give one of them") + args.quorum = 1 + fixed = (args.rl_sync_max_base_lag == 0 and args.quorum == num_islands + and int(args.grace_ms) == 0) + if fixed: + raise ValueError( + "membership options need rounds that may merge without some island: set " + "--rl-sync-quorum-theta, --rl-sync-quorum below the island count, --rl-sync-grace-ms " + "or --rl-sync-max-base-lag >= 1") + + def _prepare_node_inventory(args) -> None: """Validate ``--rl-node-inventory`` before any cloud or SSH work and keep the parsed nodes on ``args`` (rl-ssh-node-pool: the launcher reads the nodes from @@ -440,7 +482,7 @@ def _prepare_rl_args( expert_full_count=getattr(args, "expert_full_count", 0) or 0, rollout_num_gpus=getattr(iargs, "rollout_num_gpus", None), placement=getattr(iargs, "rl_placement", "colocated"), - island_scheduling=getattr(args, "rl_island_scheduling", "legacy"), + island_scheduling="elastic" if island_pool(args) else "legacy", ) if resolve_model_kind(args.model, args.model_kind) != "causal-lm": raise ValueError("RL v0 supports only causal language models") @@ -953,6 +995,7 @@ def _prepare_rl_args( ) _apply_rl_outer_merge_options(args, num_islands=len(specs)) syncer_shard_count(args) # refuse --syncer-shards outside asynchronous rounds now + _check_island_pool(args, num_islands=len(specs)) check_peer_recovery(args) if args.spot: _rl_checkpoint_mount(args.rl_completed_groups_path) diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index c4a84c71b..37e208ee2 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -15,7 +15,7 @@ from .rl_backend import _rl_backend_module from .remote_setup import read_island_pre_run from .topology import rl_island_count, rl_island_placement_args -from .syncer import island_lease_s, island_scheduling_mode, island_syncer_epoch +from .syncer import island_lease_s, island_pool, island_syncer_epoch from .storage import ELASTIC_ISLAND_STATE_DIR, rl_checkpoint_store_plan from .island_tasks import _check_island_tasks @@ -306,9 +306,6 @@ def _check_resume_launch_flags(args) -> bool: if island_cuts_requested(args) and (getattr(args, "rl_backend", None) or "miles") != "miles": raise ValueError("island round cuts (--rl-cut-every / --rl-cut-auto-mtbf-s on a synced " "launch) are wired for --rl-backend miles only") - if island_cuts_requested(args) and island_scheduling_mode(args) == "elastic": - raise ValueError("island round cuts (--rl-cut-every / --rl-cut-auto-mtbf-s on a synced " - "launch) need --rl-island-scheduling legacy (strict or decoupled)") rl_checkpoint_store_plan(args) # validates the form return False if elastic: @@ -393,7 +390,7 @@ def _check_ports_infra_switches(args) -> None: from ..rl.engine.execution_profile import UNKNOWN, check_elastic_placement, check_overlap_eval placements = island_placement.check(args, rl_island_count(args), - island_scheduling=island_scheduling_mode(args)) + island_scheduling="elastic" if island_pool(args) else "legacy") # The checks below need a fixed partition; any colocated island fails them. placement = (island_placement.FIXED_PARTITION if all(p.placement == island_placement.FIXED_PARTITION for p in placements) @@ -601,10 +598,10 @@ def _ports_infra_flags(args) -> tuple[str, str]: flags += sample_role_flags(args) # L4.6: role both (default) adds nothing if (getattr(args, "rl_lr_schedule", None) or "auto") != "auto": # auto adds nothing flags += f" --rl-lr-schedule {args.rl_lr_schedule}" - if island_scheduling_mode(args) == "elastic": # legacy adds nothing - flags += (" --rl-island-scheduling elastic" + if island_pool(args): # a fixed roster adds nothing + flags += (" --rl-island-pool" f" --rl-syncer-epoch {island_syncer_epoch(args)}" - f" --rl-island-lease-s {island_lease_s(args)!r}") + f" --rl-sync-island-lease-s {island_lease_s(args)!r}") resume_store = (rl_checkpoint_store_plan(args) if getattr(args, "rl_checkpoint_store", None) and not getattr(args, "rl_elastic", False) else None) diff --git a/yeto/launch/storage.py b/yeto/launch/storage.py index 706a17947..4b31a766c 100644 --- a/yeto/launch/storage.py +++ b/yeto/launch/storage.py @@ -10,7 +10,7 @@ import shlex from pathlib import PurePosixPath -from .syncer import island_scheduling_mode +from .syncer import island_pool ELASTIC_ISLAND_STATE_DIR = "~/yeto-rl/elastic-state" @@ -52,10 +52,10 @@ def resumes_in_new_container(args) -> bool: def droppable_rejoins_in_new_container(args, modal_cfg) -> bool: - """rl-spot-cost-saving 5.2: on an elastic run a Modal re-run of a droppable island + """rl-spot-cost-saving 5.2: on an island-pool run a Modal re-run of a droppable island (reclaim / ``modal container stop``) is expected: it rejoins from the syncer base. Every other island keeps the fail-closed container guard.""" - if island_scheduling_mode(args) != "elastic": + if not island_pool(args): return False from ..cloud import droppable diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index 0c0845a43..e3b178066 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -303,9 +303,13 @@ def rl_syncer_profile(args, num_learners: int): }) -_ISLAND_SCHEDULING_PARAMS = rl_option_pairs( - "rl_quorum_theta", "rl_carry_gamma", "rl_soft_deadline_s", "rl_q_min", "rl_max_carry_lag", - "rl_syncer_epoch", "rl_island_lease_s", "rl_final_grace_s", +# Membership options (--rl-sync-*): they put the RL islands in a pool with +# dynamic fragment-round syncer membership (leases and re-JOIN, capacity +# quorum, soft deadline, carry discount, admitted island ids, final grace +# window). None set = fixed roster. +_ISLAND_POOL_PARAMS = rl_option_pairs( + "rl_sync_quorum_theta", "rl_sync_carry_gamma", "rl_sync_soft_deadline_s", "rl_sync_q_min", + "rl_sync_island_lease_s", "rl_sync_max_learners", "rl_sync_final_grace_s", layer=CLI, ) @@ -313,14 +317,14 @@ def rl_syncer_profile(args, num_learners: int): def elastic_debug_delay_env(args) -> dict[str, str]: - """0.25 test switch: islands delay each DELTA_TENSOR (soft deadline / carried_over - on real hardware). Elastic only; {} by default.""" + """0.25 test switch: islands delay each delta push (soft deadline / carried_over + on real hardware). Island pool only; {} by default.""" spec = getattr(args, "rl_elastic_debug_delay_s", None) if not spec: return {} - if island_scheduling_mode(args) != "elastic": - raise ValueError("--rl-elastic-debug-delay-s needs --rl-island-scheduling elastic") - from ..rl.elastic_client import parse_elastic_debug_delay + if not island_pool(args): + raise ValueError("--rl-elastic-debug-delay-s needs island pool options (e.g. --rl-sync-island-lease-s)") + from ..rl.pool_client import parse_elastic_debug_delay parse_elastic_debug_delay(spec) return {ELASTIC_DEBUG_DELAY_ENV: str(spec)} @@ -333,13 +337,13 @@ def elastic_debug_pause_env(args) -> dict[str, str]: """Test switch: the given island's link to the syncer goes silent (no heartbeat, sends/receives held) for S seconds after base AFTER_V; the process keeps running. Used to exercise lease expiry -> "rejoin required" -> automatic re-JOIN on real - hardware. Elastic only; {} by default.""" + hardware. Island pool only; {} by default.""" spec = getattr(args, "rl_elastic_debug_pause", None) if not spec: return {} - if island_scheduling_mode(args) != "elastic": - raise ValueError("--rl-elastic-debug-pause needs --rl-island-scheduling elastic") - from ..rl.elastic_client import parse_elastic_debug_pause + if not island_pool(args): + raise ValueError("--rl-elastic-debug-pause needs island pool options (e.g. --rl-sync-island-lease-s)") + from ..rl.pool_client import parse_elastic_debug_pause parse_elastic_debug_pause(spec) return {ELASTIC_DEBUG_PAUSE_ENV: str(spec)} @@ -349,10 +353,10 @@ def elastic_debug_pause_env(args) -> dict[str, str]: def island_lease_s(args) -> float: - value = getattr(args, "rl_island_lease_s", None) + value = getattr(args, "rl_sync_island_lease_s", None) value = DEFAULT_ISLAND_LEASE_S if value is None else float(value) if not value > 0: - raise ValueError("--rl-island-lease-s must be > 0") + raise ValueError("--rl-sync-island-lease-s must be > 0") return value @@ -365,9 +369,9 @@ def island_syncer_epoch(args) -> int: def syncer_status_reader(*, local_tape: str | None = None, ssh_host: str | None = None): - """0.17: a zero-argument reader of the elastic syncer's status.json (next to its - --event-tape): a local file on a head node, else ``ssh cat``. - Returns the parsed snapshot, or None when absent / not the elastic schema.""" + """0.17: a zero-argument reader of the syncer's island-pool status.json (next to + its --event-tape): a local file on a head node, else ``ssh cat``. + Returns the parsed snapshot, or None when absent / not the status schema.""" from ..rl.engine.island_status import read_syncer_status, STATUS_SCHEMA if local_tape is not None: @@ -388,11 +392,11 @@ def read() -> dict | None: def resolve_island_syncer_epoch(args, read_status, *, timeout_s: float = 120.0, poll_s: float = 2.0, sleep=time.sleep, clock=time.monotonic) -> int: - """0.17: elastic without an explicit --rl-syncer-epoch takes the term the running - syncer reports in status.json (and records it on ``args`` so every island - command carries it). Unreadable within ``timeout_s``: keep 0 with a warning. - Legacy, or an explicit value: unchanged, nothing is read.""" - if island_scheduling_mode(args) != "elastic" or getattr(args, "rl_syncer_epoch", None) is not None: + """0.17: an island pool without an explicit --rl-syncer-epoch takes the term the + running syncer reports in status.json (and records it on ``args`` so every + island command carries it). Unreadable within ``timeout_s``: keep 0 with a + warning. A fixed roster, or an explicit value: unchanged, nothing is read.""" + if not island_pool(args) or getattr(args, "rl_syncer_epoch", None) is not None: return island_syncer_epoch(args) deadline = clock() + timeout_s while True: @@ -409,19 +413,11 @@ def resolve_island_syncer_epoch(args, read_status, *, timeout_s: float = 120.0, sleep(poll_s) -def island_scheduling_mode(args) -> str: - """``--rl-island-scheduling`` (rl-inter-island-scheduling 0.13); legacy unless - elastic is given explicitly. Elastic-only parameters with legacy are refused.""" - mode = getattr(args, "rl_island_scheduling", None) or "legacy" - if mode not in ("legacy", "elastic"): - raise ValueError(f"--rl-island-scheduling must be legacy|elastic, got {mode!r}") - if mode == "legacy": - given = [f for n, f in _ISLAND_SCHEDULING_PARAMS if getattr(args, n, None) is not None - # S27 E-3: peer recovery fences fragment-round HELLOs with the syncer epoch - and not (n == "rl_syncer_epoch" and getattr(args, "rl_peer_recovery", False))] - if given: - raise ValueError(", ".join(given) + " need --rl-island-scheduling elastic") - return mode +def island_pool(args) -> bool: + """True when any island pool option is set: the RL islands then join the + syncer's dynamic membership (leases and re-JOIN, capacity quorum, soft + deadline, carry discount, final grace); otherwise the roster is fixed.""" + return any(getattr(args, name, None) is not None for name, _ in _ISLAND_POOL_PARAMS) ISLAND_HMAC_KEY_ENV = "YETO_ISLAND_HMAC_KEY" @@ -432,19 +428,15 @@ def island_hmac_secret(args) -> dict[str, str]: ``secrets=``, Modal ``cfg.envs`` which becomes a ``modal.Secret``) -- never spliced into a command line, log or process list. - elastic: the key must come from the launching environment (unchanged). - legacy (secret-handling-hardening): the legacy syncer also requires the key - (HELLO HMAC). Without one in the environment the launcher generates a - random per-run key once and keeps it in ``os.environ`` so the syncer task, - the head job and every island receive the same value. SFT: ``{}`` (its + secret-handling-hardening: the syncer requires the key (HELLO HMAC). + Without one in the environment the launcher generates a random per-run key + once and keeps it in ``os.environ`` so the syncer task, the head job and + every island receive the same value. SFT: ``{}`` (its syncer runs with --allow-unauthenticated-islands, see design.md).""" if getattr(args, "training_mode", "sft") != "rl": return {} key = os.environ.get(ISLAND_HMAC_KEY_ENV, "") if not key: - if island_scheduling_mode(args) != "legacy": - raise ValueError(f"--rl-island-scheduling elastic needs {ISLAND_HMAC_KEY_ENV} " - "in the launching environment (island message HMAC key)") import secrets as _secrets key = os.environ[ISLAND_HMAC_KEY_ENV] = _secrets.token_hex(32) @@ -467,40 +459,46 @@ def _syncer_expected_contract(args) -> str: return f" --expected-island-contract {pinned}" if pinned else "" -def _syncer_island_scheduling(args) -> str: - """Syncer flags for the island-scheduling mode: "" in legacy, so the legacy - command line is byte-for-byte the previous one. The syncer folds these into - its session contract hash (the Python side does not compute that hash).""" - if island_scheduling_mode(args) == "legacy": +def _syncer_island_pool(args) -> str: + """Syncer membership flags of an island pool; "" for a fixed roster, so + that command line is byte-for-byte the previous one. Each option maps to + the syncer flag of the same name; unset options keep the syncer default.""" + if not island_pool(args): + if getattr(args, "rl_syncer_epoch", None) is not None and not getattr(args, "rl_peer_recovery", False): + raise ValueError("--rl-syncer-epoch needs island pool options (e.g. --rl-sync-island-lease-s)") return "" - from ..rl.engine.island_ledger import (DEFAULT_GAMMA, DEFAULT_MAX_CARRY_LAG, - DEFAULT_QUORUM_MIN, DEFAULT_THETA) - - theta = getattr(args, "rl_quorum_theta", None) - gamma = getattr(args, "rl_carry_gamma", None) - soft = getattr(args, "rl_soft_deadline_s", None) - if soft is None: - soft = getattr(args, "rl_elastic_quorum_timeout_s", None) or 900 - q_min = getattr(args, "rl_q_min", None) - lag = getattr(args, "rl_max_carry_lag", None) - theta = DEFAULT_THETA if theta is None else float(theta) - gamma = DEFAULT_GAMMA if gamma is None else float(gamma) - q_min = DEFAULT_QUORUM_MIN if q_min is None else int(q_min) - lag = DEFAULT_MAX_CARRY_LAG if lag is None else int(lag) - if not 0.0 < theta <= 1.0 or not 0.0 <= gamma <= 1.0 or q_min < 1 or lag < 0: - raise ValueError("--rl-quorum-theta in (0,1], --rl-carry-gamma in [0,1], " - "--rl-q-min >= 1, --rl-max-carry-lag >= 0") - if float(soft) != int(soft) or int(soft) <= 0: - raise ValueError("--rl-soft-deadline-s must be positive whole seconds") - return (" --island-scheduling-mode elastic" - f" --quorum-theta {theta!r}" - f" --carry-gamma {gamma!r}" - f" --soft-deadline-s {int(soft)}" - f" --q-min {q_min}" - f" --max-carry-lag {lag}" - f" --syncer-epoch {island_syncer_epoch(args)}" - f" --island-lease-s {island_lease_s(args)!r}" - f" --final-grace-s {_final_grace_s(args, int(soft))}") + theta = getattr(args, "rl_sync_quorum_theta", None) + gamma = getattr(args, "rl_sync_carry_gamma", None) + soft = getattr(args, "rl_sync_soft_deadline_s", None) + q_min = getattr(args, "rl_sync_q_min", None) + max_learners = getattr(args, "rl_sync_max_learners", None) + if theta is not None and not 0.0 < float(theta) <= 1.0: + raise ValueError("--rl-sync-quorum-theta must be in (0, 1]") + if gamma is not None and not 0.0 < float(gamma) <= 1.0: + raise ValueError("--rl-sync-carry-gamma must be in (0, 1]") + if q_min is not None and int(q_min) < 1: + raise ValueError("--rl-sync-q-min must be >= 1") + if soft is not None and (float(soft) != int(soft) or int(soft) <= 0): + raise ValueError("--rl-sync-soft-deadline-s must be positive whole seconds") + if max_learners is not None and int(max_learners) < 1: + raise ValueError("--rl-sync-max-learners must be >= 1") + flags = "" + if theta is not None: + flags += f" --quorum-theta {float(theta)!r}" + if q_min is not None: + flags += f" --q-min {int(q_min)}" + if gamma is not None: + flags += f" --carry-gamma {float(gamma)!r}" + if soft is not None: + flags += f" --soft-deadline-s {int(soft)}" + flags += f" --island-lease-s {island_lease_s(args)!r}" + if max_learners is not None: + flags += f" --max-learners {int(max_learners)}" + flags += f" --syncer-epoch {island_syncer_epoch(args)}" + grace = _final_grace_s(args) + if grace is not None: + flags += f" --final-grace-s {grace}" + return flags def peer_recovery_syncer_flags(args) -> str: @@ -518,13 +516,17 @@ def peer_recovery_syncer_flags(args) -> str: return flags -def _final_grace_s(args, soft_deadline_s: int) -> int: - """0.26: the syncer's FINISHED window (whole seconds); default the soft deadline.""" - value = getattr(args, "rl_final_grace_s", None) - value = soft_deadline_s if value is None else int(value) - if value < 0: - raise ValueError("--rl-final-grace-s must be >= 0") - return value + +def _final_grace_s(args) -> int | None: + """0.26: the syncer's final grace window (whole seconds); default the soft deadline.""" + value = getattr(args, "rl_sync_final_grace_s", None) + if value is None: + value = getattr(args, "rl_sync_soft_deadline_s", None) + if value is None: + return None + if int(value) < 0: + raise ValueError("--rl-sync-final-grace-s must be >= 0") + return int(value) def _rl_max_base_lag(args) -> int: @@ -568,7 +570,7 @@ def _rl_role_syncer_command(args, num_learners: int, binary: str, f" --delta-correction {args.delta_correction}" f" --total-steps {total_steps}" f"{_syncer_quorum_timeout(args)}" - f"{_syncer_island_scheduling(args)}" + f"{_syncer_island_pool(args)}" f"{_syncer_expected_contract(args)}" f"{_syncer_profile_binding(args)}" f"{_syncer_bcast_delta(args)}" @@ -593,8 +595,8 @@ def syncer_shard_count(args) -> int: Sharding needs fragment-round scheduling and a learner that routes per fragment: the torch causal-LM learner (yeto.learner) in SFT, and in RL the full-parameter preset with asynchronous rounds (the actor role - stream). Strict policy sweeps, the elastic scheduler, critics and the - other RL presets are refused.""" + stream). Strict policy sweeps, island pools, critics and the other RL + presets are refused.""" count = getattr(args, "syncer_shards", None) count = 1 if count is None else int(count) if count < 1: @@ -608,8 +610,8 @@ def syncer_shard_count(args) -> int: if not rl_full_parameter_async(args): raise ValueError("--syncer-shards > 1 with --training-mode rl needs --rl-sync-preset " f"{FULL_PARAMETER_SYNC_PRESET} --rl-sync-full-parameter-rounds async") - if island_scheduling_mode(args) != "legacy": - raise ValueError("--syncer-shards > 1 needs --rl-island-scheduling legacy") + if island_pool(args): + raise ValueError("--syncer-shards > 1 does not take island pool options yet") if rl_needs_critic(args): raise ValueError("--syncer-shards > 1 supports actor-only algorithms") return count @@ -685,8 +687,8 @@ def syncer_command(args, num_learners: int, binary: str = "~/yeto-syncer") -> st + "".join(f"{{ {command} & }} && " for command in commands.values()) + actor ) - if island_scheduling_mode(args) != "legacy": - raise ValueError("--rl-island-scheduling elastic only applies to --training-mode rl") + if island_pool(args): + raise ValueError("island pool options only apply to --training-mode rl") ports = syncer_shard_ports(args) if len(ports) == 1: return _sft_syncer_command(args, num_learners, binary) diff --git a/yeto/launcher.py b/yeto/launcher.py index d8702b642..50d35a9b4 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -40,7 +40,7 @@ from .launch.rl_backend import _rl_backend_module from .launch.topology import rl_island_args, rl_island_count, rl_island_placement_records, rl_island_spec from .launch.syncer import ( - SYNCER_PORT, island_scheduling_mode, island_syncer_addrs, make_syncer_task, probe_syncer_ports, + SYNCER_PORT, island_pool, island_syncer_addrs, make_syncer_task, probe_syncer_ports, resolve_island_syncer_epoch, rl_role_syncers, role_syncer_commands, role_syncer_shards, syncer_event_tape, syncer_shard_addrs, syncer_status_reader, ) @@ -351,7 +351,7 @@ def run(args, on_clusters=None, local_syncer=None, on_instance_ids=None) -> int: syncer_addr = syncer_shard_addrs(str(syncer_handle.head_ip), args) print(f"[launcher] syncer up at {syncer_addr}") probe_syncer_ports(args, str(syncer_handle.head_ip)) - if not no_sync and island_scheduling_mode(args) == "elastic" \ + if not no_sync and island_pool(args) \ and getattr(args, "rl_syncer_epoch", None) is None: # 0.17: the islands JOIN with the term the syncer actually runs. resolve_island_syncer_epoch(args, syncer_status_reader( @@ -660,7 +660,7 @@ def container_changed() -> bool: (lambda name: name in event_collectors and event_collectors[name].left_pool) if echo_names else None ), - elastic=island_scheduling_mode(args) == "elastic", + elastic=island_pool(args), no_island_relaunch=bool(getattr(args, "no_island_relaunch", False)), wan_reject_probe=( (lambda name: event_collectors[name].wan_rejects if name in event_collectors else 0) @@ -736,7 +736,7 @@ def drain_tapes(limit: float = NO_SYNC_EVENT_DRAIN_S) -> None: drain_tapes() if container_changed(): return CONTAINER_CHANGED_EXIT - if (no_sync_incomplete and head_mode and island_scheduling_mode(args) == "elastic"): + if (no_sync_incomplete and head_mode and island_pool(args)): # S16 §6.6: an island the syncer removed from the pool for an excused # reason (finished, stopped, preempted, lease expiry) owes no # rl_learner_finalized; report it, do not exit 3. Failure codes still fail. diff --git a/yeto/protocol.py b/yeto/protocol.py index e131432a9..06c068003 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -67,6 +67,8 @@ MSG_BUDGET_DONE = 14 # S26 E3b: link-bandwidth probe connection (yeto/wan_probe.py; syncer serve_bw_probe). MSG_BW_PROBE = 26 +# An island leaves the pool: u32 learner_id | u8 reason (yeto.rl.pool_client.LeaveReason). +MSG_LEAVE = 27 FINALIZATION_REVISION = 1 FINALIZATION_TIMEOUT = 900.0 @@ -1569,6 +1571,13 @@ def _expert_tag_block( def heartbeat(self, local_step: int) -> None: self._enqueue(0, self._frame(MSG_HEARTBEAT, struct.pack(" bool: + """Tell the syncer why this island leaves; True once the frame was + written. The caller closes the client afterwards.""" + sent = threading.Event() + frame = self._frame(MSG_LEAVE, struct.pack(" int: diff --git a/yeto/rl/adapters/miles/elastic_wiring.py b/yeto/rl/adapters/miles/elastic_wiring.py index 41d147d48..30057bfa0 100644 --- a/yeto/rl/adapters/miles/elastic_wiring.py +++ b/yeto/rl/adapters/miles/elastic_wiring.py @@ -91,7 +91,7 @@ def build_elastic( **({} if max_recovery_attempts is None else {"max_recovery_attempts": int(max_recovery_attempts)}), # rl-multinode-island Q4 (C5): --rl-elastic-checkpoint-store (None = node0-local) **({} if not checkpoint_store else {"checkpoint_store": checkpoint_store}), - # rl-inter-island-scheduling 0.13 (--rl-island-scheduling, default legacy) + # "elastic" = island pool options given, "legacy" = fixed roster (journal field) island_scheduling=island_scheduling, # 0.18: scheduling fields from the syncer's status.json (elastic only; # the controller never calls the probe under legacy) diff --git a/yeto/rl/adapters/miles/entry.py b/yeto/rl/adapters/miles/entry.py index e81cf2a64..a0a62a7d6 100644 --- a/yeto/rl/adapters/miles/entry.py +++ b/yeto/rl/adapters/miles/entry.py @@ -644,13 +644,13 @@ def build_sync( if getattr(miles_args, "yeto_rl_sync_preset", "strict-avg") == "decoupled": return DecoupledSync(miles_args), DecoupledIslandProgress(miles_args) progress = StrictIslandProgress(miles_args) - if getattr(miles_args, "yeto_rl_island_scheduling", "legacy") == "elastic": - # rl-inter-island-scheduling 0.15: absent attribute = legacy (unchanged below). - from yeto.rl.engine.bridges import ElasticAvgSync + if getattr(miles_args, "yeto_rl_island_pool", False): + # Dynamic island pool; absent attribute = fixed roster (unchanged below). + from yeto.rl.engine.bridges import PoolAvgSync if getattr(miles_args, "use_critic", False): - raise ValueError("--rl-island-scheduling elastic does not support a critic yet") - return ElasticAvgSync(miles_args.yeto_rl_bridge_config, progress=progress, + raise ValueError("island pool options do not support a critic yet") + return PoolAvgSync(miles_args.yeto_rl_bridge_config, progress=progress, syncer_epoch=int(getattr(miles_args, "yeto_rl_syncer_epoch", 0)), lease_s=getattr(miles_args, "yeto_rl_island_lease_s", None), groups_per_round=int(getattr(miles_args, "rollout_batch_size", 0) or 0) or None, @@ -914,7 +914,7 @@ def _wire_droppable_reclaim(sync: Any, driver: Any, environ: Any = None) -> Any: return None def leave() -> bool: - from yeto.rl.elastic_client import LeaveReason + from yeto.rl.pool_client import LeaveReason return sync.leave(LeaveReason.PREEMPTED) # once per process; the driver's close won't resend diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 99641e990..2340513d4 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -54,6 +54,7 @@ ) from yeto.rl.local_learner import MATRIX_MERGES, ComponentIdentity, layout_stream_role from yeto.rl.moe_layout import RoutedTokens, check_expert_partition, fragment_expert_tags +from yeto.rl.pool_client import LeaseHeartbeat from yeto.rl.trajectory_evidence import TrajectoryBatchEvidence # Seconds between progress lines while an island waits for its initial cut. @@ -153,6 +154,12 @@ class RoleStreamConfig: # during the push's outer interval; the syncer (--expert-token-weights) # weights that expert's tensors by them. Needs a layout with experts. expert_token_weights: bool = False + # Island pool (--rl-sync-island-lease-s; asynchronous rounds only): the + # syncer's lease on this island. The stream heartbeats every lease/3; a + # lease expiry or a lost link redials (re-JOIN) and the syncer's catch-up + # may skip fragment versions. None = fixed roster. Not part of the + # session contract (membership is a syncer option). + island_lease_s: float | None = None @property def tensor_blocks_active(self) -> bool: @@ -246,6 +253,9 @@ def __post_init__(self) -> None: raise ValueError("role-stream matrix merge must be avg or rda") if self.overlap and self.partial_quorum: raise ValueError("sync/train overlap needs the strict sweep roster") + if self.island_lease_s is not None and ( + not self.partial_quorum or not self.island_lease_s > 0): + raise ValueError("role-stream island lease needs asynchronous rounds and a positive lease") if type(self.push_int4_block) is not int or ( not I4_BLOCK_MIN <= self.push_int4_block <= I4_BLOCK_MAX or self.push_int4_block % 2 ): @@ -267,6 +277,8 @@ def __post_init__(self) -> None: raise ValueError("peer recovery needs asynchronous rounds without broadcast deltas") if self.peer_recovery is not None and self.syncer_shards: raise ValueError("peer recovery needs one syncer process, not a sharded head") + if self.peer_recovery is not None and self.island_lease_s is not None: + raise ValueError("peer recovery is wired for a fixed roster, not an island pool") if self.wait_timeout <= 0 or not 0 < self.poll_seconds <= 1: raise ValueError("role-stream wait bounds are invalid") if ( @@ -381,6 +393,7 @@ def __init__(self, config: RoleStreamConfig, *, ray_module=None) -> None: self.group = None self.adapter: MilesChunkedFullParameterAdapter | None = None self.client = None + self.heartbeat: LeaseHeartbeat | None = None self.sink = None self.current: ReferencedPolicyCut | None = None # Keep the latest post-train/pre-broadcast cut until the next local @@ -520,6 +533,12 @@ async def initialize(self, group) -> None: recover=self._recovering, expert_tokens=self.config.expert_token_weights, ) + heartbeat = ( + LeaseHeartbeat(client, self.config.island_lease_s, + progress=lambda: self.local_rounds, + name=f"{self.config.role}-heartbeat") + if self.config.island_lease_s is not None else None + ) sink = adapter.fragment_sink(initial, ray_module=self.ray_module) client.install_inbound_chunk_sink(MilesRayInboundChunkSink(sink)) if self._recovering: @@ -527,6 +546,8 @@ async def initialize(self, group) -> None: return try: client.start() + if heartbeat is not None: + heartbeat.start() # the initial cut may take longer than the lease init_parts = None if self.config.learner_id == 0: def init_parts(fragment_id: int): @@ -547,6 +568,8 @@ def init_parts(fragment_id: int): ) adapter.verify_initial_fragments(initial, authoritative) except BaseException: + if heartbeat is not None: + heartbeat.stop() client.close() sink.release_all() adapter.release(initial) @@ -555,6 +578,7 @@ def init_parts(fragment_id: int): self.group = group self.adapter = adapter self.client = client + self.heartbeat = heartbeat self.sink = sink self.current = initial self._anchors = list(authoritative) @@ -698,6 +722,12 @@ def _wait_for_initial_fragments( for update in client.drain_updates(): stored = self._stored_update(update) if update.version != 0 or stored.fragment_id in observed: + if self.config.island_lease_s is not None and update.version != 0: + raise RuntimeError( + f"{self.config.algorithm} {self.config.role}: the island pool is past " + "its initial cut; a full-parameter island cannot join a running pool " + "yet (it can re-JOIN after a lease expiry while its process runs)" + ) raise RuntimeError( f"{self.config.algorithm} {self.config.role} initial broadcast changed" ) @@ -927,6 +957,8 @@ async def apply_staged_broadcast( finally: rebased.release() # the applied target holds its own references self._staged.pop(key, None) + for older in [k for k in self._staged if k[0] == fragment_id and k[1] < version]: + self._staged.pop(older).release() if self.config.partial_quorum: # The round closed: an unanswered PULL of it can never commit. permit = self._pending.get(fragment_id) @@ -1363,6 +1395,8 @@ def discard_detached_submission_cut(self) -> None: self.submission_fragment_versions = None def close(self) -> None: + if self.heartbeat is not None: + self.heartbeat.stop() if self.client is not None: self.client.close() if self.sink is not None: diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 575eebdf0..f178541e3 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -413,7 +413,7 @@ def apply_ports_infra_switches(args, miles_args, environ=None) -> None: miles_args.yeto_rl_elastic["accept_rebind"] = True if getattr(args, "rl_elastic_checkpoint_store", None): miles_args.yeto_rl_elastic["checkpoint_store"] = str(args.rl_elastic_checkpoint_store) - if getattr(args, "rl_island_scheduling", "legacy") == "elastic": + if getattr(args, "rl_island_pool", False): miles_args.yeto_rl_elastic["island_scheduling"] = "elastic" for name in _ELASTIC_PAUSE: if getattr(args, name, None) is not None: @@ -611,7 +611,7 @@ def _require_ports_supported(args, extra_argv: Sequence[str] = ()) -> None: expert_full_count=getattr(args, "expert_full_count", 0) or 0, rollout_num_gpus=getattr(args, "rollout_num_gpus", None), placement=getattr(args, "rl_placement", "colocated"), - island_scheduling=getattr(args, "rl_island_scheduling", "legacy"), + island_scheduling="elastic" if getattr(args, "rl_island_pool", False) else "legacy", extra_argv=tuple(extra_argv), ) @@ -932,8 +932,9 @@ def _dense_full_sync_config( overlap=bool(getattr(args, "sweep_overlap", False)), partial_quorum=async_rounds, matrix_merge=getattr(args, "rl_sync_matrix_merge", "avg"), - syncer_epoch=_island_syncer_epoch(args), peer_recovery=peer_recovery_options(args), + # an island pool sets the epoch (and lease); else peer recovery may + **{"syncer_epoch": _island_syncer_epoch(args), **_island_pool_membership(args)}, ) return FullParameterSyncConfig((actor,), sweep=not async_rounds) @@ -969,6 +970,19 @@ def peer_recovery_options(args, environ=None): ) +def _island_pool_membership(args) -> dict: + """Island pool (``--rl-island-pool``, set by the launcher with any + --rl-sync-* membership option): the syncer lease the session heartbeats + for and the syncer epoch of its HELLO. {} for a fixed roster.""" + if not getattr(args, "rl_island_pool", False): + return {} + from yeto.rl.pool_client import DEFAULT_ISLAND_LEASE_S + + lease = getattr(args, "rl_sync_island_lease_s", None) + return {"island_lease_s": DEFAULT_ISLAND_LEASE_S if lease is None else float(lease), + "syncer_epoch": int(getattr(args, "rl_syncer_epoch", 0) or 0)} + + def _load_sao_streaming_runtime(args): """Load the hash-bound SAO streaming contract before any GPU work. @@ -2251,6 +2265,7 @@ def run_miles( code_pack_digest=code_pack_digest, compat_group=_island_compat_group(args), matrix_merge=matrix_merge, + **_island_pool_membership(args), ) else: miles_args.yeto_rl_bridge_config = BridgeConfig( @@ -2275,11 +2290,11 @@ def run_miles( compat_group=_island_compat_group(args), bcast_delta=bool(getattr(args, "bcast_delta", False)), ) - if getattr(args, "rl_island_scheduling", "legacy") == "elastic": - # rl-inter-island-scheduling 0.15 (legacy sets nothing) - miles_args.yeto_rl_island_scheduling = "elastic" + if getattr(args, "rl_island_pool", False): + # a fixed roster sets nothing + miles_args.yeto_rl_island_pool = True miles_args.yeto_rl_syncer_epoch = int(getattr(args, "rl_syncer_epoch", 0)) - miles_args.yeto_rl_island_lease_s = getattr(args, "rl_island_lease_s", None) + miles_args.yeto_rl_island_lease_s = getattr(args, "rl_sync_island_lease_s", None) # 4.2.1: the critic syncer (second channel); its BridgeConfig needs # the critic specs/layout, built by build_sync once Miles created it. miles_args.yeto_rl_critic_syncer_addr = ( @@ -2657,7 +2672,7 @@ def _rejoin_exhausted(error: BaseException | None) -> bool: except BaseException as error: # noqa: BLE001 if _rejoin_exhausted(error): # 0.22: the island gave up re-joining the elastic pool (not 4/6). - from yeto.rl.elastic_client import ELASTIC_REJOIN_FAILED_EXIT + from yeto.rl.pool_client import ELASTIC_REJOIN_FAILED_EXIT print(f"[yeto-island] elastic re-JOIN exhausted: {error}", file=sys.stderr) raise SystemExit(ELASTIC_REJOIN_FAILED_EXIT) from error diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 043781d57..febb0d818 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -270,6 +270,10 @@ ) _ISLAND = ( + # Set by the launcher when any island pool option is given (dynamic membership). + island_flag("--rl-island-pool", action="store_true", + help="this island joins a dynamic island pool on the syncer: it heartbeats " + "its lease, re-JOINs after an expiry and may see the outer version jump"), island_flag("--model", required=True), island_flag("--model-revision", required=True), island_flag("--data", required=True), diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 22da4c7d9..504f168e3 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -265,7 +265,7 @@ def run_sampler_island(miles_args: Any, *, environ: Any = None, from miles.utils.async_utils import Disposer from miles.utils.orchestration_utils import init_orchestration_script - from yeto.rl.elastic_client import LeaveReason + from yeto.rl.pool_client import LeaveReason from yeto.rl.engine.events import append_tape_event from yeto.rl.engine.resume import store_for from yeto.rl.engine.sampler_loop import SamplerLoop diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index 5fec3a02a..b499f46c6 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -64,9 +64,9 @@ def parse_args(argv): p.add_argument("--bcast-delta", action="store_true") p.add_argument("--apply-chat-template-kwargs", default=None) p.add_argument("--rl-algorithm-spec", default=None) - p.add_argument("--rl-island-scheduling", default="legacy") + p.add_argument("--rl-island-pool", action="store_true") p.add_argument("--rl-syncer-epoch", type=int, default=0) - p.add_argument("--rl-island-lease-s", type=float, default=None) + p.add_argument("--rl-sync-island-lease-s", type=float, default=None) p.add_argument("--rl-echo-events", action="store_true") p.add_argument("--rl-max-policy-age", type=int, default=0) # S25 C2: a trainer island over the sample buffer (sample_roles; the sampler is @@ -329,8 +329,8 @@ def main(argv=None) -> int: ] if args.rl_single_island_no_sync: sync = "none" - elif args.rl_island_scheduling == "elastic": - sync = "elastic" + elif args.rl_island_pool: + sync = "pool" else: sync = "strict" if sync != "none" and not args.syncer: @@ -362,7 +362,7 @@ def main(argv=None) -> int: "bcast_delta": bool(args.bcast_delta), "thresholds_key": ["verl", "fsdp2", "vllm-" + str(manifest["versions"].get("vllm")), _threshold_card(manifest)], - "algorithm_spec": spec_dict, "syncer_epoch": args.rl_syncer_epoch, "island_lease_s": args.rl_island_lease_s, + "algorithm_spec": spec_dict, "syncer_epoch": args.rl_syncer_epoch, "island_lease_s": args.rl_sync_island_lease_s, "test_exit_after_version": exit_after, "publish_selftest": os.environ.get("YETO_VERL_PUBLISH_SELFTEST") == "1", "run_config": run.to_dict(), "overrides": overrides, "asserted": {k: _override_value(overrides, k) for k in asserted_keys}, diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index 08c72e83c..d7f2969e1 100644 --- a/yeto/rl/adapters/verl/trainer.py +++ b/yeto/rl/adapters/verl/trainer.py @@ -167,8 +167,8 @@ def publish_selftest(island, publisher, policy_state, specs) -> dict: def build_sync(plan: dict, specs, *, lora_hash: str, layout_hash: str, identity_sha256: str, compat_group: str | None = None): - """Outer sync of the island plan: none / strict / elastic (both trainer paths).""" - from yeto.rl.engine.bridges import ElasticAvgSync, LocalOnlySync, StrictAvgSync + """Outer sync of the island plan: none / strict / pool (both trainer paths).""" + from yeto.rl.engine.bridges import LocalOnlySync, PoolAvgSync, StrictAvgSync learner_id = int(plan["learner_id"]) mode = plan["sync"] @@ -190,8 +190,8 @@ def build_sync(plan: dict, specs, *, lora_hash: str, layout_hash: str, identity_ bcast_delta=bool(plan.get("bcast_delta", False))) if mode == "strict": sync = StrictAvgSync(bridge, progress=None) - elif mode == "elastic": - sync = ElasticAvgSync(bridge, progress=None, syncer_epoch=int(plan.get("syncer_epoch", 0)), + elif mode == "pool": + sync = PoolAvgSync(bridge, progress=None, syncer_epoch=int(plan.get("syncer_epoch", 0)), lease_s=plan.get("island_lease_s")) else: raise ValueError(f"unknown sync mode {mode!r}") diff --git a/yeto/rl/bridge.py b/yeto/rl/bridge.py index 7cae92564..259591146 100644 --- a/yeto/rl/bridge.py +++ b/yeto/rl/bridge.py @@ -808,26 +808,6 @@ def _append_event(self, event: dict) -> None: # ----------------------------------------------------------------------------- elastic -def flatten_state(state: CanonicalLoraState, specs) -> list[float]: - """Flat f32 vector in canonical spec order (elastic wire format).""" - parts = [state.tensors[spec.name].detach().to(torch.float32).reshape(-1) for spec in specs] - return torch.cat(parts).tolist() if parts else [] - - -def unflatten_state(values, version: int, template: CanonicalLoraState, specs) -> CanonicalLoraState: - flat = torch.tensor(values, dtype=torch.float32) - tensors, off = {}, 0 - for spec in specs: - tensors[spec.name] = flat[off:off + spec.numel].reshape(spec.shape).clone() - off += spec.numel - if off != flat.numel(): - raise StrictRlInvariantError("layout_hash_mismatch", - f"elastic base has {flat.numel()} values, layout needs {off}") - return canonical_state(version, tensors, base_model_revision=template.base_model_revision, - lora_config_hash=template.lora_config_hash, - layout_hash=template.layout_hash, expected_specs=specs) - - def _session_contract(layout, identity_sha256: str | None) -> bytes | None: """HELLO session contract (decoupling 6.2): layout fingerprint bound to the backend identity.""" if identity_sha256 is None: diff --git a/yeto/rl/decoupled.py b/yeto/rl/decoupled.py index 36e25181c..f5f92905b 100644 --- a/yeto/rl/decoupled.py +++ b/yeto/rl/decoupled.py @@ -9,6 +9,7 @@ from .engine.backend_identity import config_island_contract_sha256 from ..protocol import DTYPE_F32, FinalManifest, PullRequest, SyncerClient +from .pool_client import DEFAULT_MAX_REJOIN_FAILURES, LeaseHeartbeat from ..tensor_io import ( apply_fragment, fragment_flat, @@ -54,6 +55,12 @@ class DecoupledBridgeConfig: compat_group: str | None = None # S27 P1: --rl-sync-matrix-merge (yeto.rl.core.build_rl_fragment_layout). matrix_merge: str = "avg" + # Island pool (--rl-sync-island-lease-s): the syncer's lease on this + # island; the bridge heartbeats every lease/3 and redials (re-JOIN) when + # the lease expired or the link was lost. None = fixed roster. + island_lease_s: float | None = None + # Island pool: the syncer incarnation this island was launched for. + syncer_epoch: int | None = None def __post_init__(self) -> None: if self.num_fragments < 2: @@ -69,6 +76,12 @@ def __post_init__(self) -> None: raise ValueError("decoupled RL local horizon must be at least 2") if self.learner_budget_steps is not None and self.learner_budget_steps < 1: raise ValueError("decoupled RL learner budget must be positive") + if self.island_lease_s is not None and not self.island_lease_s > 0: + raise ValueError("decoupled RL island lease must be positive") + + @property + def island_pool(self) -> bool: + return self.island_lease_s is not None @dataclass(frozen=True) @@ -152,8 +165,18 @@ def __init__( session_contract_hash=_session_contract( self.layout, config_island_contract_sha256(config)), compat_group=getattr(config, "compat_group", None), - max_reconnects=(None if config.learner_budget_steps is not None else 0), + max_reconnects=self._max_reconnects(config), + syncer_epoch=config.syncer_epoch, + ) + # Island pool: keep the lease while the island trains between pushes. + self._optimizer_steps = 0 + self._heartbeat = ( + LeaseHeartbeat(self.client, config.island_lease_s, + progress=lambda: self._optimizer_steps, + name="decoupled-heartbeat") + if config.island_pool else None ) + self._connection_generation: int | None = None count = self.layout.num_fragments self._anchors: list[torch.Tensor | None] = [None] * count self._fragment_versions: list[int | None] = [None] * count @@ -164,6 +187,15 @@ def __init__( self.startup_final_manifest: FinalManifest | None = None self.final_payload_bytes_received = 0 + @staticmethod + def _max_reconnects(config: DecoupledBridgeConfig) -> int | None: + """A fixed roster fails on a lost link (no redial) unless it runs a + learner budget; an island pool redials a bounded number of times + (its rounds merge without it meanwhile).""" + if config.learner_budget_steps is not None: + return None + return DEFAULT_MAX_REJOIN_FAILURES if config.island_pool else 0 + @property def fragment_versions(self) -> tuple[int, ...]: if any(version is None for version in self._fragment_versions): @@ -195,6 +227,8 @@ def finalizing(self) -> bool: def start(self) -> None: self.client.start() + if self._heartbeat is not None: + self._heartbeat.start() if self.config.learner_id == 0: params = dict(self.initial.tensors) for fragment_id, fragment in enumerate(self.layout.fragments): @@ -271,6 +305,7 @@ def drain_broadcasts( ) -> BroadcastBatch: if optimizer_steps < 0 or action_tokens < 0: raise ValueError("decoupled RL progress must be non-negative") + self._optimizer_steps = optimizer_steps return self._drain_broadcast_updates(local, self._fragment_versions) def _drain_broadcast_updates( @@ -329,6 +364,10 @@ def commit_broadcasts( self._fragment_versions[update.fragment_id] = update.version self._steps_at_anchor[update.fragment_id] = optimizer_steps self._tokens_at_anchor[update.fragment_id] = action_tokens + permit = self._pending.get(update.fragment_id) + if self.config.island_pool and permit is not None and permit.global_step <= update.version: + # The round closed without this island; its PULL can never commit. + self._pending.pop(update.fragment_id) def submit_ready( self, @@ -338,6 +377,7 @@ def submit_ready( action_tokens: int, ) -> tuple[FragmentSubmission, ...]: self.client.check_health() + self._optimizer_steps = optimizer_steps if self.finalizing: return () state = self._canonical(local) @@ -560,15 +600,33 @@ def acknowledge_finalization(self, manifest: FinalManifest) -> None: self.client.acknowledge_finalization(manifest) def close(self) -> None: + if self._heartbeat is not None: + self._heartbeat.stop() self.client.close() def _queue_pulls(self) -> None: + if self.config.island_pool: + # A new connection generation (re-JOIN after a lease expiry or a + # lost link): the syncer replays the PULLs this island still owes. + generation = int(getattr(self.client, "connection_generation", 0)) + if self._connection_generation != generation: + self._pending.clear() + self._answered.clear() + self._connection_generation = generation for permit in self.client.drain_pulls(): self._validate_permit(permit) key = (permit.global_step, permit.round_attempt) if key in self._answered: continue previous = self._pending.get(permit.fragment_id) + if (self.config.island_pool and previous is not None + and previous.global_step != permit.global_step): + # Island pool: a round may close without this island; the + # syncer permits a fragment's next round only after the + # previous one committed, so the newer permit supersedes. + if permit.global_step > previous.global_step: + self._pending[permit.fragment_id] = permit + continue if previous is not None and previous != permit: raise RuntimeError("decoupled RL received conflicting fragment permits") if previous is None: diff --git a/yeto/rl/elastic_client.py b/yeto/rl/elastic_client.py deleted file mode 100644 index 5e660be06..000000000 --- a/yeto/rl/elastic_client.py +++ /dev/null @@ -1,927 +0,0 @@ -"""Island-side client for the elastic (inter-island scheduling) syncer. - -Only used with ``--rl-island-scheduling elastic``; the legacy path never imports -this module. Frame layout mirrors ``syncer/src/elastic.rs`` (ElasticMsg::encode -/ decode) byte for byte: - -* wire frame: the usual `` int: - return MAX_ELASTIC_BASE_FRAME if msg_type == MSG_ELASTIC_BASE else MAX_ELASTIC_FRAME -HMAC_KEY_ENV = "YETO_ISLAND_HMAC_KEY" - - -class ElasticProtocolError(RuntimeError): - pass - - -def hmac_sha256(key: bytes, msg: bytes) -> bytes: - return _hmac.new(key, msg, hashlib.sha256).digest() - - -def seal(key: bytes, msg_type: int, body: bytes) -> bytes: - return body + hmac_sha256(key, bytes([msg_type]) + body) - - -def open_frame(key: bytes, msg_type: int, payload: bytes) -> bytes: - if len(payload) < HMAC_LEN: - raise ElasticProtocolError("elastic frame shorter than its HMAC") - body, mac = payload[:-HMAC_LEN], payload[-HMAC_LEN:] - if not _hmac.compare_digest(hmac_sha256(key, bytes([msg_type]) + body), mac): - raise ElasticProtocolError("elastic frame HMAC mismatch") - return body - - -# ----------------------------------------------------------------------------- messages -@dataclass(frozen=True) -class Join: - syncer_epoch: int - island_id: int - incarnation: int - capacity: float - backend_identity: bytes = bytes(32) # BackendIdentity.sha256() bytes; zero = not declared - compat_group: str = "" # 7.7c: "-"; empty = not declared - TYPE = MSG_JOIN - - -@dataclass(frozen=True) -class JoinAck: - syncer_epoch: int - learner_slot: int - membership_epoch: int - base_version: int - policy_hash: bytes - catch_up: bool - TYPE = MSG_JOIN_ACK - - -@dataclass(frozen=True) -class Leave: - syncer_epoch: int - island_id: int - reason: int = LeaveReason.FINISHED - TYPE = MSG_LEAVE - - -@dataclass(frozen=True) -class LeaseHeartbeat: - syncer_epoch: int - island_id: int - membership_epoch: int - inner_step: int - round_wall_s: float - TYPE = MSG_LEASE_HEARTBEAT - - -@dataclass(frozen=True) -class SampleIndex: - """Wire form of ``sample_pool.SampleIndexEntry`` (elastic.rs SampleIndexEntry): - island_id is the numeric island id, policy_hash the raw 32-byte sha256. - Convert with :func:`sample_index_from_entry` / :func:`sample_index_to_entry`.""" - syncer_epoch: int - island_id: int - outer_version: int - inner_step: int - policy_hash: bytes - group_id: str - prompt_id: str - n: int - uri: str - size_bytes: int - sha256: str - has_behavior_logprob: bool - advantage_included: bool - created_at: float - schema: str - TYPE = MSG_SAMPLE_INDEX - - -@dataclass(frozen=True) -class Finished: - syncer_epoch: int - final_outer_version: int - policy_hash: bytes - TYPE = MSG_FINISHED - - -class ElasticFinished(Exception): - """The syncer finished the run (FINISHED, or its connection went away after the - final version was known): the island ends normally.""" - - def __init__(self, final_outer_version: int | None, reason: str) -> None: - super().__init__(f"elastic syncer finished at outer version {final_outer_version} ({reason})") - self.final_outer_version = final_outer_version - self.reason = reason - - -@dataclass(frozen=True) -class SampleVerdict: - syncer_epoch: int - verdict: int # VERDICT_* - reason: str - outer_lag: int - TYPE = MSG_SAMPLE_VERDICT - - @property - def verdict_name(self) -> str: - return VERDICT_NAMES.get(self.verdict, f"UNKNOWN({self.verdict})") - - -def island_number(island_id: str | int, island_ids: Mapping[str, int] | None = None) -> int: - """String island name -> u32 wire id: explicit map first, else a decimal string.""" - if isinstance(island_id, int) and not isinstance(island_id, bool): - n = island_id - elif island_ids is not None and island_id in island_ids: - n = int(island_ids[island_id]) - elif isinstance(island_id, str) and island_id.isdigit(): - n = int(island_id) - else: - raise ElasticProtocolError(f"island {island_id!r} has no numeric id (pass island_ids)") - if not 0 <= n < 1 << 32: - raise ElasticProtocolError(f"island id {n} does not fit u32") - return n - - -def policy_hash_bytes(value: str | bytes) -> bytes: - """Hex sha256 string (sample pool form) -> raw 32 bytes (wire form).""" - if isinstance(value, (bytes, bytearray)): - return _hash32(bytes(value)) - try: - raw = bytes.fromhex(value) - except ValueError as exc: - raise ElasticProtocolError(f"policy_hash {value!r} is not hex") from exc - return _hash32(raw) - - -def sample_index_from_entry(entry: Any, *, syncer_epoch: int, - island_ids: Mapping[str, int] | None = None) -> SampleIndex: - """``sample_pool.SampleIndexEntry`` -> SAMPLE_INDEX frame.""" - return SampleIndex(int(syncer_epoch), island_number(entry.island_id, island_ids), - int(entry.outer_version), int(entry.inner_step), - policy_hash_bytes(entry.policy_hash), entry.group_id, entry.prompt_id, - int(entry.n), entry.uri, int(entry.size_bytes), entry.sha256, - bool(entry.has_behavior_logprob), bool(entry.advantage_included), - float(entry.created_at), entry.schema) - - -def sample_index_to_entry(msg: SampleIndex, *, island_names: Mapping[int, str] | None = None): - """SAMPLE_INDEX frame -> ``sample_pool.SampleIndexEntry`` (validated).""" - from .engine.sample_pool import SampleIndexEntry - - name = (island_names or {}).get(msg.island_id, str(msg.island_id)) - return SampleIndexEntry(name, msg.outer_version, msg.inner_step, msg.policy_hash.hex(), - msg.group_id, msg.prompt_id, msg.n, msg.uri, msg.size_bytes, - msg.sha256, msg.has_behavior_logprob, msg.advantage_included, - msg.created_at, msg.schema) - - -@dataclass(frozen=True) -class DeltaReady: - syncer_epoch: int - island_id: int - base_version: int - c_tokens: int - c_steps: int - TYPE = MSG_DELTA_READY - - -@dataclass(frozen=True) -class ElasticInit: - """Initial flat f32 parameter vector (the first one received wins).""" - syncer_epoch: int - island_id: int - params: tuple[float, ...] - TYPE = MSG_ELASTIC_INIT - - -@dataclass(frozen=True) -class DeltaTensor: - """theta - base as a flat f32 vector (same sign convention as PUSH_FRAGMENT).""" - syncer_epoch: int - island_id: int - base_version: int - c_tokens: int - c_steps: int - update: tuple[float, ...] - TYPE = MSG_DELTA_TENSOR - - -@dataclass(frozen=True) -class ElasticBase: - syncer_epoch: int - outer_version: int - params: tuple[float, ...] - TYPE = MSG_ELASTIC_BASE - - -def _hash32(value: bytes) -> bytes: - if len(value) != 32: - raise ElasticProtocolError("policy_hash must be 32 bytes") - return bytes(value) - - -def _f32s(values) -> bytes: - """u64 count + little-endian f32 values (elastic.rs put_f32s).""" - import array - - arr = array.array("f", values) - if sys.byteorder != "little": - arr.byteswap() - return struct.pack(" bytes: - """u32 length + utf-8 (elastic.rs put_str).""" - raw = value.encode("utf-8") - return struct.pack(" tuple[int, bytes]: - b = struct.pack(" None: - self.data, self.pos = data, 0 - - def take(self, n: int) -> bytes: - if self.pos + n > len(self.data): - raise ElasticProtocolError("truncated elastic frame") - out = self.data[self.pos:self.pos + n] - self.pos += n - return out - - def unpack(self, fmt: str): - s = struct.Struct("<" + fmt) - return s.unpack(self.take(s.size)) - - def f32s(self) -> tuple[float, ...]: - import array - - (n,) = self.unpack("Q") - arr = array.array("f") - arr.frombytes(self.take(4 * n)) - if sys.byteorder != "little": - arr.byteswap() - return tuple(arr) - - def text(self) -> str: - (n,) = self.unpack("I") - try: - return self.take(n).decode("utf-8") - except UnicodeDecodeError as exc: - raise ElasticProtocolError("string is not utf-8") from exc - - def flag(self) -> bool: - v = self.take(1)[0] - if v not in (0, 1): - raise ElasticProtocolError(f"invalid bool byte {v}") - return bool(v) - - -def decode(key: bytes, msg_type: int, payload: bytes) -> Any: - r = _Reader(open_frame(key, msg_type, payload)) - (epoch,) = r.unpack("Q") - if msg_type == MSG_JOIN: - island, incarnation, capacity = r.unpack("IQd") - try: - identity = r.take(32) - except ElasticProtocolError as exc: - raise ElasticProtocolError( - "JOIN without backend identity (island older than the syncer?)") from exc - try: - group = r.text() - except ElasticProtocolError as exc: - raise ElasticProtocolError( - "JOIN without compat_group (island older than the syncer?)") from exc - msg = Join(epoch, island, incarnation, capacity, identity, group) - if not (msg.capacity > 0 and msg.capacity != float("inf")): - raise ElasticProtocolError("JOIN capacity must be > 0") - elif msg_type == MSG_JOIN_ACK: - slot, mepoch, base = r.unpack("IQQ") - msg = JoinAck(epoch, slot, mepoch, base, r.take(32), r.flag()) - elif msg_type == MSG_LEAVE: - msg = Leave(epoch, *r.unpack("IB")) - elif msg_type == MSG_LEASE_HEARTBEAT: - msg = LeaseHeartbeat(epoch, *r.unpack("IQQd")) - elif msg_type == MSG_SAMPLE_INDEX: - island, version, inner = r.unpack("IQQ") - h = r.take(32) - group, prompt = r.text(), r.text() - (n,) = r.unpack("Q") - uri = r.text() - (size,) = r.unpack("Q") - sha = r.text() - has_lp, adv = r.flag(), r.flag() - (created,) = r.unpack("d") - msg = SampleIndex(epoch, island, version, inner, h, group, prompt, n, uri, size, sha, - has_lp, adv, created, r.text()) - elif msg_type == MSG_FINISHED: - (final,) = r.unpack("Q") - msg = Finished(epoch, final, r.take(32)) - elif msg_type == MSG_SAMPLE_VERDICT: - verdict = r.take(1)[0] - reason = r.text() - (lag,) = r.unpack("Q") - msg = SampleVerdict(epoch, verdict, reason, lag) - elif msg_type == MSG_DELTA_READY: - msg = DeltaReady(epoch, *r.unpack("IQQQ")) - elif msg_type == MSG_ELASTIC_INIT: - (island,) = r.unpack("I") - msg = ElasticInit(epoch, island, r.f32s()) - elif msg_type == MSG_DELTA_TENSOR: - island, base, tokens, steps = r.unpack("IQQQ") - msg = DeltaTensor(epoch, island, base, tokens, steps, r.f32s()) - elif msg_type == MSG_ELASTIC_BASE: - (version,) = r.unpack("Q") - msg = ElasticBase(epoch, version, r.f32s()) - else: - raise ElasticProtocolError(f"message type {msg_type} is not decoded by the island") - if r.pos != len(r.data): - raise ElasticProtocolError(f"trailing bytes in elastic frame type {msg_type}") - return msg - - -DEBUG_DELAY_ENV = "YETO_RL_ELASTIC_DEBUG_DELAY" - - -def parse_elastic_debug_delay(spec: str | None) -> dict[int, float]: - """``"ISLAND:SECONDS[,ISLAND:SECONDS]"`` (bare ``SECONDS`` = every island).""" - out: dict[int, float] = {} - for item in (spec or "").split(","): - item = item.strip() - if not item: - continue - island, sep, seconds = item.rpartition(":") - key = int(island) if sep else -1 - value = float(seconds) - if value < 0 or (sep and key < 0): - raise ValueError(f"--rl-elastic-debug-delay-s: bad item {item!r}") - out[key] = value - return out - - -DEBUG_PAUSE_ENV = "YETO_RL_ELASTIC_DEBUG_PAUSE" - - -def parse_elastic_debug_pause(spec: str | None) -> dict[int, tuple[int, float]]: - """``"ISLAND:AFTER_V:SECONDS[,...]"``: once the island applied an ELASTIC_BASE with - outer version >= AFTER_V, its link to the syncer goes silent for SECONDS (no - heartbeats, sends held, received frames held) while the process keeps training.""" - out: dict[int, tuple[int, float]] = {} - for item in (spec or "").split(","): - item = item.strip() - if not item: - continue - parts = item.split(":") - try: - island, after_v, seconds = int(parts[0]), int(parts[1]), float(parts[2]) - if len(parts) != 3: - raise ValueError - except (ValueError, IndexError): - raise ValueError(f"--rl-elastic-debug-pause: bad item {item!r} (want ISLAND:AFTER_V:SECONDS)") from None - if island < 0 or after_v < 0 or seconds <= 0: - raise ValueError(f"--rl-elastic-debug-pause: bad item {item!r}") - out[island] = (after_v, seconds) - return out - - -def hmac_key_from_env(environ: dict | None = None) -> bytes: - key = (os.environ if environ is None else environ).get(HMAC_KEY_ENV, "") - if not key: - raise ElasticProtocolError(f"elastic island needs {HMAC_KEY_ENV}") - return key.encode("utf-8") - - -# ----------------------------------------------------------------------------- client -@dataclass -class ElasticClientConfig: - syncer_addr: tuple[str, int] - island_id: int - capacity: float = 1.0 - incarnation: int = 0 - lease_s: float = DEFAULT_ISLAND_LEASE_S - syncer_epoch: int = 0 - join_timeout_s: float = 60.0 - # 0.22: consecutive failed automatic re-JOINs before the island gives up. - max_rejoin_failures: int = 5 - # 0.22: heartbeats on their own TCP connection (a multi-MB DELTA_TENSOR / - # ELASTIC_BASE on the main connection must not delay lease renewal). - heartbeat_connection: bool = True - # decoupling 6.2a: BackendIdentity.sha256() hex sent in every JOIN; None = not - # declared (32 zero bytes). The syncer refuses a JOIN whose identity differs - # from the one the session's first JOIN pinned. - backend_identity_sha256: str | None = None - # 7.7c: compat_group sent in every JOIN ("" = not declared); the syncer pins - # the first one and refuses another group, naming both. - compat_group: str | None = None - - def backend_identity_bytes(self) -> bytes: - value = self.backend_identity_sha256 - if value is None: - return bytes(32) - if not isinstance(value, str) or len(value) != 64: - raise ElasticProtocolError("backend_identity_sha256 must be a 64-char sha256 hex digest") - return bytes.fromhex(value) - - -REJOIN_MARKER = "rejoin required" -# learner exit code when the island gave up re-joining (not the launcher's 4/6). -ELASTIC_REJOIN_FAILED_EXIT = 7 - - -class ElasticJoinRefused(ElasticProtocolError): - """The syncer answered the first JOIN with an error (e.g. backend identity or - compat_group differ): the island never became a member, so no LEAVE goes - out; the departure is CONFIG_MISMATCH on the island tape only (design D7 Q6).""" - - -class ElasticRejoinError(ElasticProtocolError): - """``max_rejoin_failures`` consecutive automatic re-JOINs failed.""" - - -class ElasticIslandClient: - """JOIN → JOIN_ACK, LEASE_HEARTBEAT every lease_s/3, DELTA_READY, LEAVE.""" - - def __init__(self, config: ElasticClientConfig, key: bytes, *, - connect: Callable[[tuple[str, int]], socket.socket] | None = None, - on_event: Callable[[dict], None] | None = None) -> None: - self.config = config - self.key = key - self._connect = connect or (lambda addr: socket.create_connection(addr, timeout=20.0)) - self._on_event = on_event or (lambda ev: None) - self.sock: socket.socket | None = None - self._send_lock = threading.Lock() - self._stop = threading.Event() - self._ack = threading.Event() - self._left = threading.Event() - self._base_cv = threading.Condition() - self.ack: JoinAck | None = None - self.errors: list[str] = [] - self.latest_base: ElasticBase | None = None - # L5.6: (Unix time it became visible to the island, seconds from header to last byte) - # of ``latest_base``; set together with it under ``_base_cv``. - self.latest_base_timing: tuple[float, float] | None = None - self._verdict_cv = threading.Condition() - self._verdicts: list[SampleVerdict] = [] - self._index_lock = threading.Lock() # one SAMPLE_INDEX in flight: replies are in order - self.inner_step = 0 - self.round_wall_s = 0.0 - self._threads: list[threading.Thread] = [] - self.hb_sock: socket.socket | None = None - self._hb_lock = threading.Lock() - self._rejoin = threading.Event() # set by a "rejoin required" MSG_ERROR - self._rejoin_lock = threading.Lock() - self.rejoins = 0 # successful automatic re-JOINs - self.rejoin_failures = 0 # consecutive failed ones - self.fatal: ElasticRejoinError | None = None - self.join_refusal: str | None = None # first MSG_ERROR before any JOIN_ACK - self.finished: Finished | None = None # 0.26 - self.disconnected = False - # 0.26: outer version at which the run ends (the syncer's total_steps); a - # lost connection once a base at/after it arrived counts as finished. - self.final_outer_version: int | None = None - self._base_seq = 0 # ELASTIC_BASE frames received - self._rejoin_base_seq = 0 # _base_seq when the last re-JOIN was sent - self._pause_until = 0.0 # debug link pause (monotonic); 0 = never paused - - @property - def heartbeat_period_s(self) -> float: - return self.config.lease_s / 3.0 - - def pause_link(self, seconds: float) -> None: - """Test switch (--rl-elastic-debug-pause): simulate a network outage of - ``seconds``. Heartbeats are skipped, sends wait until the outage ends and - received frames are held back; nothing else in the process stops.""" - t0 = time.time() - self._pause_until = time.monotonic() + float(seconds) - self._on_event({"event": "elastic_debug_pause", "pause_s": float(seconds), "start_unix": t0}) - - def _end() -> None: - if self._stop.wait(float(seconds)): - return - self._on_event({"event": "elastic_debug_pause_end", "pause_s": float(seconds), - "start_unix": t0, "end_unix": time.time()}) - th = threading.Thread(target=_end, name="elastic-debug-pause", daemon=True) - th.start() - - def link_paused(self) -> bool: - return self._pause_until > 0.0 and time.monotonic() < self._pause_until - - def _link_wait(self) -> None: - while self._pause_until > 0.0 and not self._stop.is_set(): - left = self._pause_until - time.monotonic() - if left <= 0: - return - self._stop.wait(min(left, 0.5)) - - def _send(self, msg: Any) -> None: - self._link_wait() - t, payload = encode(msg, self.key) - with self._send_lock: - if self.sock is None: - raise ElasticProtocolError("elastic client is not connected") - try: - write_frame(self.sock, t, payload) - except (BrokenPipeError, ConnectionResetError, ConnectionRefusedError): - if self.finished is not None or self.known_final(): - raise ElasticFinished( - self.finished.final_outer_version if self.finished else self.latest_base.outer_version, - "connection closed after the final version") from None - raise - self._on_event({"event": "rl_elastic_sent", "type": t, **_fields(msg)}) - - def _reader(self) -> None: - while not self._stop.is_set(): - header_at = 0.0 - - def limit(msg_type: int) -> int: - nonlocal header_at - header_at = time.monotonic() # header read; payload follows - return _frame_limit(msg_type) - - try: - t, payload = read_frame(self.sock, limit) - recv_seconds = time.monotonic() - header_at - except (OSError, ConnectionError, ValueError) as exc: - if not self._stop.is_set(): - self.errors.append(f"connection: {exc}") - self.disconnected = True - self._ack.set() - with self._base_cv: - self._base_cv.notify_all() - return - self._link_wait() - if t == MSG_ERROR: - text = payload.decode("utf-8", "replace") - if self.ack is None and self.join_refusal is None: - self.join_refusal = text - self._on_error(text) - self._ack.set() # a refused JOIN must not wait for the timeout - continue - try: - msg = decode(self.key, t, payload) - except ElasticProtocolError as exc: - self.errors.append(str(exc)) - continue - if isinstance(msg, JoinAck): - self.ack = msg - self._ack.set() - elif isinstance(msg, Finished): - with self._base_cv: - self.finished = msg - self._base_cv.notify_all() - elif isinstance(msg, SampleVerdict): - with self._verdict_cv: - self._verdicts.append(msg) - self._verdict_cv.notify_all() - elif isinstance(msg, ElasticBase): - with self._base_cv: - self._base_seq += 1 - if self.latest_base is None or msg.outer_version >= self.latest_base.outer_version: - self.latest_base = msg - self.latest_base_timing = (time.time(), recv_seconds) - self._base_cv.notify_all() - self._on_event({"event": "rl_elastic_received", "type": t, **_fields(msg)}) - - def _on_error(self, text: str) -> None: - self.errors.append(text) - self._on_event({"event": "rl_elastic_error", "error": text}) - if REJOIN_MARKER in text and self.ack is not None and not self._left.is_set(): - self._rejoin.set() - - def _hb_reader(self, sock: socket.socket) -> None: - """The heartbeat connection only ever gets MSG_ERROR replies (not a member).""" - while not self._stop.is_set(): - try: - t, payload = read_frame(sock, _frame_limit) - except (OSError, ConnectionError, ValueError): - return - self._link_wait() - if t == MSG_ERROR: - self._on_error(payload.decode("utf-8", "replace")) - - def _open_heartbeat_connection(self) -> None: - if not self.config.heartbeat_connection: - return - try: - sock = self._connect(self.config.syncer_addr) - sock.settimeout(None) - except OSError as exc: # fall back to the main connection - self.errors.append(f"heartbeat connection: {exc}") - return - if sock is self.sock: # test doubles with a single socket - return - self.hb_sock = sock - t = threading.Thread(target=self._hb_reader, args=(sock,), name="elastic-hb-reader", daemon=True) - t.start() - self._threads.append(t) - - def rejoin(self) -> bool: - """0.22: re-JOIN with a new incarnation after a lease expiry; the syncer then - answers JOIN_ACK (catch_up) + the current ELASTIC_BASE. False on failure; - after ``max_rejoin_failures`` consecutive failures ``fatal`` is set.""" - with self._rejoin_lock: - if not self._rejoin.is_set() or self._left.is_set(): - return True - c = self.config - c.incarnation += 1 - self._ack.clear() - old = self.ack - self.ack = None - n_err = len(self.errors) - with self._base_cv: - self._rejoin_base_seq = self._base_seq - try: - self._send(Join(c.syncer_epoch, c.island_id, c.incarnation, float(c.capacity), - c.backend_identity_bytes(), c.compat_group or "")) - ok = self._ack.wait(c.join_timeout_s) and self.ack is not None - except Exception as exc: # noqa: BLE001 - self.errors.append(f"rejoin: {exc}") - ok = False - if ok: - self._rejoin.clear() - self.rejoins += 1 - self.rejoin_failures = 0 - self._on_event({"event": "elastic_rejoin", "incarnation": c.incarnation, - "base_version": self.ack.base_version, "catch_up": self.ack.catch_up, - "membership_epoch": self.ack.membership_epoch}) - return True - self.ack = self.ack or old - self.rejoin_failures += 1 - self._on_event({"event": "elastic_rejoin_failed", "incarnation": c.incarnation, - "failures": self.rejoin_failures, "errors": self.errors[n_err:][-3:]}) - if self.rejoin_failures >= c.max_rejoin_failures: - self.fatal = ElasticRejoinError( - f"{self.rejoin_failures} consecutive re-JOINs failed: " + "; ".join(self.errors[-3:])) - return False - - def check(self) -> None: - """Raise the client's fatal error (rejoin exhausted), else re-JOIN if needed.""" - if self.fatal is not None: - raise self.fatal - if self._rejoin.is_set(): - self.rejoin() - if self.fatal is not None: - raise self.fatal - - def _heartbeat(self) -> None: - while not self._stop.wait(self.heartbeat_period_s): - if self._left.is_set(): - return - if self.link_paused(): - continue # debug link pause: no lease renewal - try: - if self._rejoin.is_set() and self.fatal is None: - self.rejoin() # also while the main thread trains - elif self.fatal is None: - self.heartbeat() - except Exception as exc: # reported, the reader decides liveness - self.errors.append(f"heartbeat: {exc}") - - def join(self) -> JoinAck: - self.sock = self._connect(self.config.syncer_addr) - self.sock.settimeout(None) - reader = threading.Thread(target=self._reader, name="elastic-reader", daemon=True) - reader.start() - self._threads.append(reader) - c = self.config - self._send(Join(c.syncer_epoch, c.island_id, c.incarnation, float(c.capacity), - c.backend_identity_bytes(), c.compat_group or "")) - if not self._ack.wait(c.join_timeout_s) or self.ack is None: - if self.join_refusal is not None: - raise ElasticJoinRefused("JOIN refused: " + self.join_refusal) - raise ElasticProtocolError("JOIN refused or timed out: " + "; ".join(self.errors)) - self._open_heartbeat_connection() - hb = threading.Thread(target=self._heartbeat, name="elastic-heartbeat", daemon=True) - hb.start() - self._threads.append(hb) - return self.ack - - def heartbeat(self) -> None: - if self.ack is None: - raise ElasticProtocolError("heartbeat before JOIN_ACK") - msg = LeaseHeartbeat(self.config.syncer_epoch, self.config.island_id, - self.ack.membership_epoch, int(self.inner_step), float(self.round_wall_s)) - if self.hb_sock is None: - self._send(msg) - return - t, payload = encode(msg, self.key) - with self._hb_lock: - write_frame(self.hb_sock, t, payload) - self._on_event({"event": "rl_elastic_sent", "type": t, **_fields(msg)}) - - def delta_ready(self, *, base_version: int, c_tokens: int, c_steps: int) -> None: - self._send(DeltaReady(self.config.syncer_epoch, self.config.island_id, - int(base_version), int(c_tokens), int(c_steps))) - - def elastic_init(self, params) -> None: - self._send(ElasticInit(self.config.syncer_epoch, self.config.island_id, tuple(params))) - - def delta_tensor(self, *, base_version: int, c_tokens: int, c_steps: int, update) -> None: - self._send(DeltaTensor(self.config.syncer_epoch, self.config.island_id, int(base_version), - int(c_tokens), int(c_steps), tuple(update))) - - def sample_index(self, entry: Any, *, island_ids: Mapping[str, int] | None = None, - timeout_s: float = 30.0) -> SampleVerdict: - """Send one ``SampleIndexEntry`` (or a ready SampleIndex) and wait for its SAMPLE_VERDICT.""" - msg = entry if isinstance(entry, SampleIndex) else sample_index_from_entry( - entry, syncer_epoch=self.config.syncer_epoch, island_ids=island_ids) - with self._index_lock: - with self._verdict_cv: - self._verdicts.clear() - n_errors = len(self.errors) - self._send(msg) - deadline = time.monotonic() + timeout_s - with self._verdict_cv: - while not self._verdicts: - left = deadline - time.monotonic() - if left <= 0 or self._stop.is_set() or len(self.errors) > n_errors: - raise ElasticProtocolError("no SAMPLE_VERDICT: " + "; ".join(self.errors)) - self._verdict_cv.wait(min(left, 0.1)) - return self._verdicts.pop(0) - - def wait_base(self, *, newer_than: int | None, timeout_s: float) -> ElasticBase | None: - """Block until an ELASTIC_BASE newer than ``newer_than`` (any if None) arrived. - - 0.22: after an automatic re-JOIN during the wait (our delta may have been - refused while we were not a member) the first base received after the - re-JOIN is returned even when it is not newer (catch-up).""" - deadline = time.monotonic() + timeout_s - rejoins0 = self.rejoins - with self._base_cv: - while True: - if self.fatal is not None: - raise self.fatal - b = self.latest_base - if b is not None and (newer_than is None or b.outer_version > newer_than): - return b - self.raise_if_finished() - if (b is not None and self.rejoins > rejoins0 - and self._base_seq > self._rejoin_base_seq): - return b - left = deadline - time.monotonic() - if left <= 0: - return None - self._base_cv.wait(min(left, 0.1)) - - def known_final(self) -> bool: - b = self.latest_base - return (self.final_outer_version is not None and b is not None - and b.outer_version >= self.final_outer_version) - - def raise_if_finished(self) -> None: - """0.26: FINISHED received, or the syncer is gone after the final version.""" - if self.finished is not None: - raise ElasticFinished(self.finished.final_outer_version, "FINISHED") - if self.disconnected and self.known_final(): - raise ElasticFinished(self.latest_base.outer_version, "connection closed after the final version") - - def leave(self, reason: int = LeaveReason.FINISHED) -> None: - self._left.set() # no heartbeat after LEAVE - if self.ack is not None and self.sock is not None: - self._send(Leave(self.config.syncer_epoch, self.config.island_id, reason)) - - def close(self) -> None: - self._stop.set() - if self.hb_sock is not None: - try: - self.hb_sock.shutdown(socket.SHUT_RDWR) - except OSError: - pass - self.hb_sock.close() - if self.sock is not None: - try: - self.sock.shutdown(socket.SHUT_RDWR) - except OSError: - pass - self.sock.close() - for t in self._threads: - t.join(timeout=2.0) - - -def _fields(msg: Any) -> dict: - out = {} - for k, v in vars(msg).items(): - if isinstance(v, bytes): - v = v.hex() - elif isinstance(v, tuple) and len(v) > 8: - v = f"<{len(v)} floats>" - out[k] = v - return out diff --git a/yeto/rl/engine/bridges.py b/yeto/rl/engine/bridges.py index 2849a6760..db954b6e9 100644 --- a/yeto/rl/engine/bridges.py +++ b/yeto/rl/engine/bridges.py @@ -22,7 +22,7 @@ |---|---|---|---| | strict (``StrictAvgSync``, ``DualStrictAvgSync`` for the critic) | syncer ``--learner-weight equal`` (launcher RL default); ``tokens2-over-steps`` = ``merge::learner_weight`` when chosen | ``--max-base-lag 0``: every island trains from the exact global version v and pushes against v; quorum = all islands | stop when the committed version reaches ``global_rounds``; ``finish`` checks the final policy equals the committed one | | decoupled (``DecoupledSync``) | same syncer flags as strict (``equal``, ``--max-base-lag 0``), per fragment | fragments rotate with ``--pipeline`` fragments in flight; each fragment is merged against its own fragment version | stop when the syncer is finalizing or a submission reaches ``total_fragment_steps``; ``learner_budget_steps`` consolidates into one final cut | -| elastic (``ElasticAvgSync``) | ``merge_weight`` = ``merge::learner_weight`` (tokens^2 / steps); a catch-up island has weight 0 on its join version; a late delta is carried with ``carry_gamma^lag`` | a delta with lag > ``DEFAULT_MAX_CARRY_LAG`` (2) is rejected (``carry_lag_exceeded``); samples at most ``MAX_SAMPLE_OUTER_LAG`` (2) outer steps stale; step at quorum ``theta`` (0.75) of capacity, at least ``q_min`` (1) | the syncer answers FINISHED after ``total_steps``, and keeps answering for ``final_grace_s`` (default: the soft deadline) | +| island pool (``PoolAvgSync``) | syncer ``--learner-weight`` (``tokens2-over-steps`` by default here); a late delta is carried with ``carry_gamma^lag`` | a delta whose base lags the round by more than ``--rl-sync-max-base-lag`` is rejected (``carry_lag_exceeded``); a round merges at quorum ``theta`` of capacity (at least ``q_min``) or at the soft deadline; islands hold leases and re-JOIN after an expiry | the syncer sends the final cut after ``total_steps`` and waits ``final_grace_s`` for the islands still connected | """ from __future__ import annotations @@ -45,12 +45,12 @@ policy_hash, policy_tensor_hash, ) -from yeto.rl import elastic_client -from yeto.rl.elastic_client import ( +from yeto.rl import pool_client +from yeto.rl.pool_client import ( DEFAULT_ISLAND_LEASE_S, - ElasticClientConfig, - ElasticFinished, LeaveReason, + PoolClientConfig, + PoolFinished, ) from . import progress as legacy_progress @@ -447,9 +447,9 @@ def close(self, reason: int = LeaveReason.FINISHED) -> None: self.bridge.client.close() -def _elastic_local_round_fields(stats: LocalRoundStats, base_version: int, payload_bytes: int) -> dict: +def _pool_local_round_fields(stats: LocalRoundStats, base_version: int, payload_bytes: int) -> dict: """0.24: the strict bridge's ``rl_local_round`` fields (yeto/rl/bridge.py - ``_record_submission``) for an elastic round; base = the syncer's outer version.""" + ``_record_submission``) for an island-pool round; base = the syncer's outer version.""" return { **asdict(stats), "rl/active_groups": stats.active_groups, @@ -476,16 +476,39 @@ def _elastic_local_round_fields(stats: LocalRoundStats, base_version: int, paylo } -class ElasticAvgSync: - """``--rl-island-scheduling elastic`` sync session (rl-inter-island-scheduling 0.15). +def _flat(state: CanonicalLoraState, specs) -> torch.Tensor: + """The LoRA vector as one flat f32 tensor in ``specs`` order.""" + parts = [state.tensors[spec.name].detach().to(torch.float32).reshape(-1) for spec in specs] + return torch.cat(parts) if parts else torch.zeros(0, dtype=torch.float32) - Same session interface as :class:`StrictAvgSync`, backed by the elastic - syncer: JOIN at start (ELASTIC_INIT when no base exists yet), one - DELTA_TENSOR per boundary (theta - base as a flat f32 vector; c_tokens = - ``stats.action_tokens``, c_steps = ``local_optimizer_steps``), then the next - ELASTIC_BASE is reset-applied. The driver's rollout ids stay local round - ids; the syncer's outer version is tracked separately (``base_version``) - because a late island may see it jump. LEAVE at finish/close. + +def _unflat(flat: torch.Tensor, version: int, template: CanonicalLoraState, specs) -> CanonicalLoraState: + tensors, offset = {}, 0 + for spec in specs: + tensors[spec.name] = flat[offset:offset + spec.numel].reshape(spec.shape).clone() + offset += spec.numel + if offset != flat.numel(): + from yeto.rl.core import StrictRlInvariantError + + raise StrictRlInvariantError("layout_hash_mismatch", + f"pool base has {flat.numel()} values, layout needs {offset}") + return canonical_state(version, tensors, base_model_revision=template.base_model_revision, + lora_config_hash=template.lora_config_hash, + layout_hash=template.layout_hash, expected_specs=specs) + + +class PoolAvgSync(RoundCutHooks): + """Sync session of a LoRA island in a dynamic island pool. + + Same session interface as :class:`StrictAvgSync`, on the fragment-round + syncer with membership options (lease, capacity quorum, soft deadline, + carry discount, final grace): join at start (island 0 seeds the initial + vector when the syncer has none yet), one delta per boundary (theta - base as one + flat f32 fragment; c_tokens = ``stats.action_tokens``, c_steps = + ``local_optimizer_steps``), then the next base is reset-applied. The + driver's rollout ids stay local round ids; the syncer's outer version is + tracked separately (``base_version``) because a late island may see it + jump. Leaves the pool at finish/close. """ OUTER_SYNC_KIND = "elastic" @@ -507,7 +530,7 @@ def __init__(self, config: BridgeConfig, *, progress: StrictIslandProgress | Non self.client = client self.template: CanonicalLoraState | None = None self.specs = None - self.base: list[float] | None = None + self.base: torch.Tensor | None = None self.base_version: int | None = None self.in_boundary = False self.left = False @@ -515,32 +538,28 @@ def __init__(self, config: BridgeConfig, *, progress: StrictIslandProgress | Non def _client(self): if self.client is None: + from yeto.rl.bridge import _session_contract + from yeto.rl.core import build_avg_layout + + c = self.config + layout = build_avg_layout(self.specs) + config = PoolClientConfig( + c.syncer_addr, c.learner_id, lease_s=self.lease_s, syncer_epoch=self.syncer_epoch, + session_contract_hash=_session_contract(layout, config_island_contract_sha256(c)), + compat_group=getattr(c, "compat_group", None), wan_streams=c.wan_streams) # looked up on the module at call time: tests substitute the client class - self.client = elastic_client.ElasticIslandClient( - self._client_config(), elastic_client.hmac_key_from_env(), - on_event=self._client_event) + self.client = pool_client.PoolIslandClient(config, layout, on_event=self._client_event) return self.client - def _client_config(self) -> ElasticClientConfig: - c = self.config - return ElasticClientConfig(c.syncer_addr, c.learner_id, syncer_epoch=self.syncer_epoch, - lease_s=self.lease_s, - backend_identity_sha256=config_island_contract_sha256(c), - compat_group=getattr(c, "compat_group", None)) - _TAPE_CLIENT_EVENTS = ("elastic_rejoin", "elastic_rejoin_failed", "elastic_debug_pause", "elastic_debug_pause_end") def _client_event(self, ev: dict) -> None: - """0.22: re-JOIN outcomes go onto the island tape (sent/received frames do not).""" + """0.22: re-JOIN outcomes go onto the island tape.""" driver = getattr(self, "_driver", None) if driver is not None and ev.get("event") in self._TAPE_CLIENT_EVENTS: fields = {k: v for k, v in ev.items() if k != "event"} driver.emit(ev["event"], **fields) - limit = getattr(getattr(self.client, "config", None), "max_rejoin_failures", None) - if ev["event"] == "elastic_rejoin_failed" and limit and ev.get("failures", 0) >= limit: - # 0.27: the launcher reads this as "left the elastic pool" (exit 7) - driver.emit("elastic_left_pool", reason="rejoin_exhausted", failures=ev["failures"]) def _base_timing(self, submit_ts: float, base) -> SyncTiming: """L5.6 timing of ``base``; the client records receive time only @@ -554,24 +573,21 @@ def _base_timing(self, submit_ts: float, base) -> SyncTiming: def _wait(self, newer_than): base = self.client.wait_base(newer_than=newer_than, timeout_s=self.base_wait_s) if base is None: - raise RuntimeError(f"no ELASTIC_BASE newer than {newer_than} within " + raise RuntimeError(f"no syncer base newer than {newer_than} within " f"{self.base_wait_s}s; errors: {self.client.errors}") return base def _apply(self, driver: IslandDriver, base, rollout_id: int) -> TrainableState: - from yeto.rl.bridge import unflatten_state - - self.base, self.base_version = list(base.params), base.outer_version - state = unflatten_state(self.base, rollout_id, self.template, self.specs) + self.base, self.base_version = base.params, int(base.outer_version) + state = _unflat(self.base, rollout_id, self.template, self.specs) return driver.apply_policy(TrainableState.from_lora(state), optimizer="reset", local_step=rollout_id * self.config.local_optimizer_steps, outer_version=int(base.outer_version)) def start(self, driver: IslandDriver) -> SyncStart: - from yeto.rl.bridge import flatten_state - self.template = _lora(driver.export_local()) - self.specs = tuple(self.config.expected_specs or ()) or tuple(self.template.specs) + # The syncer layout orders tensors by spec (core.build_avg_layout). + self.specs = tuple(sorted(self.config.expected_specs or ())) or tuple(sorted(self.template.specs)) self._driver = driver injected = self.client is not None # a client built by _client() already reports to the tape client = self._client() @@ -588,13 +604,14 @@ def chained(ev, _inner=inner): driver.phase("elastic_join", catch_up=ack.catch_up, base_version=ack.base_version) base = client.wait_base(newer_than=None, timeout_s=1.0) if base is None: - client.elastic_init(flatten_state(self.template, self.specs)) + if int(self.config.learner_id) == 0: # the syncer takes INIT from island 0 only + client.elastic_init(_flat(self.template, self.specs)) base = self._wait(None) # 0.21: a restarted island resumes after its own trained rounds and never # behind the syncer's version (the driver then seeks the data cursor). ledger = getattr(driver, "ledger", None) ledger_next = ledger.next_rollout_id() if ledger is not None else 0 - start = max(ledger_next, int(ack.base_version), 0) + start = max(ledger_next, int(base.outer_version), 0) if start > 0: driver.phase("elastic_resume", rollout_id=start, ledger_next=ledger_next, base_version=ack.base_version, outer_version=base.outer_version) @@ -606,7 +623,7 @@ def debug_delay_s(self) -> float: """0.25: ``YETO_RL_ELASTIC_DEBUG_DELAY`` = "ISLAND:S[,...]" (bare S = all).""" import os - from yeto.rl.elastic_client import DEBUG_DELAY_ENV, parse_elastic_debug_delay + from yeto.rl.pool_client import DEBUG_DELAY_ENV, parse_elastic_debug_delay spec = os.environ.get(DEBUG_DELAY_ENV, "") if not spec: @@ -619,7 +636,7 @@ def debug_pause(self) -> tuple[int, float] | None: """``YETO_RL_ELASTIC_DEBUG_PAUSE`` = "ISLAND:AFTER_V:S[,...]" -> (after_v, s) for this island.""" import os - from yeto.rl.elastic_client import DEBUG_PAUSE_ENV, parse_elastic_debug_pause + from yeto.rl.pool_client import DEBUG_PAUSE_ENV, parse_elastic_debug_pause spec = os.environ.get(DEBUG_PAUSE_ENV, "") if not spec: @@ -644,14 +661,13 @@ def restart_cursor_fallback(self, current: Mapping[str, int] | None, return whole_round_restart_cursor(current, start_rollout_id, self.groups_per_round) def _reached_final(self) -> bool: - """S17 M1: an elastic island ends once it applied the syncer's final outer + """S17 M1: a pool island ends once it applied the syncer's final outer version, not after ``global_rounds`` LOCAL rounds. A re-JOINed island lost - rounds (its delta refused while it was not a member, then the catch-up - entry weighted 0) yet its local rollout id kept counting, so on the local - count it left before the last outer step (s17-m1-20261008a: island 1 left - after v5, v6 merged island 0 alone). When base == rollout_id + 1 (no - rejoin) this is the old rule. Extra rounds past the local schedule run at - the trainer's tail learning rate (linear decay: 0) -- known limitation.""" + rounds (its delta refused while it was not a member) yet its local rollout + id kept counting, so on the local count it would leave before the last + outer step. When base == rollout_id + 1 (no rejoin) both rules agree. + Extra rounds past the local schedule run at the trainer's tail learning + rate (linear decay: 0) -- known limitation.""" return self.base_version is not None and self.base_version >= self.config.global_rounds def is_final_round(self, driver, *, rollout_id: int) -> bool: @@ -660,23 +676,20 @@ def is_final_round(self, driver, *, rollout_id: int) -> bool: return rollout_id + 1 >= self.config.global_rounds def boundary(self, driver, *, rollout_id, stats) -> SyncBoundary: - from yeto.rl.bridge import flatten_state - if self.base is None: - raise RuntimeError("elastic sync called outside an active round") + raise RuntimeError("pool sync called outside an active round") if rollout_id + 1 > self.config.global_rounds: driver.phase("elastic_extra_round", rollout_id=rollout_id, base_version=self.base_version) self.in_boundary = True driver.phase("export_push", rollout_id=rollout_id, base_version=self.base_version) - local = flatten_state(_lora(driver.export_local()), self.specs) - update = [a - b for a, b in zip(local, self.base)] + update = _flat(_lora(driver.export_local()), self.specs) - self.base sent_on = self.base_version - stats = replace(stats, delta_l2_norm=math.sqrt(sum(u * u for u in update))) + stats = replace(stats, delta_l2_norm=float(torch.linalg.vector_norm(update))) delay = self.debug_delay_s() if delay > 0: # 0.25 test switch (heartbeats continue on their own thread) driver.phase("elastic_debug_delay", rollout_id=rollout_id, delay_s=delay) self._sleep(delay) - driver.emit("rl_local_round", **_elastic_local_round_fields(stats, sent_on, 4 * len(update)), + driver.emit("rl_local_round", **_pool_local_round_fields(stats, sent_on, 4 * update.numel()), **{"sync/base_outer_version": sent_on}) self.client.inner_step += self.config.local_optimizer_steps try: @@ -693,20 +706,18 @@ def boundary(self, driver, *, rollout_id, stats) -> SyncBoundary: **self._base_timing(submit_ts, base).fields()) state = self._apply(driver, base, rollout_id + 1) self._maybe_pause_link(driver) - except ElasticFinished as done: + except PoolFinished as done: return self._finished(driver, done, rollout_id=rollout_id, sent_on=sent_on) self.in_boundary = False return SyncBoundary(state, self._reached_final()) def _finished(self, driver, done, *, rollout_id: int, sent_on: int) -> SyncBoundary: """0.26: the syncer ended the run while this (late) island was still - pushing: end normally on the newest base (exit 0), not BrokenPipe/exit 1.""" + pushing: end normally on the final cut (exit 0), not BrokenPipe/exit 1.""" base = self.client.latest_base driver.emit("elastic_finished", rollout_id=rollout_id, reason=done.reason, final_outer_version=done.final_outer_version, sent_on=sent_on, base_version=None if base is None else base.outer_version) - # FINISHED: LEAVE first so the syncer exits as soon as every island left - # instead of waiting out its final grace window; best-effort otherwise. self.leave(LeaveReason.FINISHED) self.in_boundary = False if base is not None and base.outer_version > sent_on: @@ -725,27 +736,27 @@ def finish(self, driver) -> None: self.leave(LeaveReason.FINISHED) def leave(self, reason: int) -> bool: - """Send LEAVE once per process; True when this call sent it (the first reason wins: a reclaim - listener's PREEMPTED is not overwritten by the driver's close). A failed - send is reported as ``elastic_leave_failed``, never raised: the driver + """Leave the pool once per process; True when this call left (the first + reason wins: a reclaim listener's PREEMPTED is not overwritten by the + driver's close). The reason goes on the island tape; a failure to + close is reported as ``elastic_leave_failed``, never raised: the driver may be propagating the original error (design D2).""" with self._leave_lock: if self.client is None or self.left: return False self.left = True code = sendable(reason) + driver = getattr(self, "_driver", None) try: self.client.leave(code) - return True except Exception as error: # noqa: BLE001 - the syncer may already be gone - if isinstance(error, (ElasticFinished, BrokenPipeError, ConnectionResetError)) \ - and code is LeaveReason.FINISHED: - return False # 0.26: the syncer already finished - driver = getattr(self, "_driver", None) if driver is not None: driver.emit("elastic_leave_failed", reason=int(code), reason_name=code.name, error_type=type(error).__name__, error=str(error)) return False + if driver is not None: + driver.emit("elastic_leave", reason=int(code), reason_name=code.name) + return True def close(self, reason: int = LeaveReason.FINISHED) -> None: if self.client is not None: diff --git a/yeto/rl/engine/convertible_remote.py b/yeto/rl/engine/convertible_remote.py index 1246c1988..d1768c039 100644 --- a/yeto/rl/engine/convertible_remote.py +++ b/yeto/rl/engine/convertible_remote.py @@ -615,7 +615,7 @@ def release(self, job: _NodeJob) -> bool: def release_incarnation(self, incarnation: int) -> bool: """A trainer the dead head started: tell it to stop (SCHEDULED cut); the node converts back on its own.""" - from yeto.rl.elastic_client import LeaveReason + from yeto.rl.pool_client import LeaveReason append_trainer_command(self.files, incarnation, STOP, reason=LeaveReason.STOPPED.name) return True diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index ab81668a2..2dff41301 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -72,7 +72,7 @@ StrictRlInvariantError, require_nonzero_learning_rate, ) -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from . import phase_profile from .algorithm import AlgorithmSpec diff --git a/yeto/rl/engine/engine_sampler_entry.py b/yeto/rl/engine/engine_sampler_entry.py index 45481f986..838a3fbae 100644 --- a/yeto/rl/engine/engine_sampler_entry.py +++ b/yeto/rl/engine/engine_sampler_entry.py @@ -258,7 +258,7 @@ def run(spec: Mapping[str, Any], *, environ: Mapping[str, str] | None = None, **fields)`` sees every tape record after it is written.""" import os # noqa: PLC0415 - from yeto.rl.elastic_client import LeaveReason # noqa: PLC0415 + from yeto.rl.pool_client import LeaveReason # noqa: PLC0415 from .sampler_resume import ( # noqa: PLC0415 RECLAIM_MARKER_ENV, diff --git a/yeto/rl/engine/fake_lifecycle.py b/yeto/rl/engine/fake_lifecycle.py index a73bd034c..2c7d2d4df 100644 --- a/yeto/rl/engine/fake_lifecycle.py +++ b/yeto/rl/engine/fake_lifecycle.py @@ -21,7 +21,7 @@ import torch -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from . import recovery from .fake import FakeEngine, FakeRolloutPool, FakeTrainerGroup diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index 7ca0bbd31..29a93c3b3 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -64,7 +64,7 @@ from typing import TYPE_CHECKING, Any, Protocol from yeto.protocol import PartialMessageGenerationLost -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from yeto.wan_streams import upload_window from .bridges import RoundCutHooks @@ -129,6 +129,10 @@ async def stage_available_broadcasts(self) -> tuple[tuple[int, int], ...]: ... def staged_fragment_versions(self) -> tuple[tuple[int, int], ...]: ... def has_staged_broadcast(self, fragment_id: int, version: int) -> bool: ... async def apply_staged_broadcast(self, fragment_id: int, version: int) -> tuple[int, ...]: ... + # Asynchronous rounds (profile.partial_quorum) only: monotonic_broadcasts + # (the stream never delivers an older version of a fragment after a + # newer one) lets apply_staged_broadcast(..., catch_up=True) overwrite a + # fragment with a version that skips older ones (re-JOIN catch-up). # S26 E3 sync/train overlap (profile.overlap) only: # complete_local_round(..., hold_submission=True) keeps the round base and # the previous submission; apply_staged_broadcast(..., rebase=True) applies @@ -269,6 +273,7 @@ def __init__(self, streams: Mapping[str, RoleStream], profile: FullParameterSync # (catch-ups included), and the rounds catch-up jumps skipped. self.overwritten_rounds = 0 self.caught_up_rounds = 0 + self.caught_up_fragments = 0 # catch-up jumps (one per fragment and jump) # -- fail-stop --------------------------------------------------------- @@ -500,6 +505,7 @@ async def _commit_partial_quorum(self, applied: dict[str, list[int]]) -> None: f == fragment_id and v < version for f, v in staged ): await self._overwrite_or_apply_pair(fragment_id, version, applied, catch_up=True) + self.caught_up_fragments += 1 self.caught_up_rounds += fragment_rounds(fragment_id, base, fragments) - ( fragment_rounds(fragment_id, current, fragments)) @@ -744,6 +750,7 @@ def _partial_quorum_fields(self, versions: tuple[int, ...]) -> dict[str, int]: "sync/overwritten_rounds": pairing.overwritten_rounds, "sync/overwritten_fragments": overwritten, "sync/caught_up_rounds": pairing.caught_up_rounds, + "sync/caught_up_fragments": pairing.caught_up_fragments, "sync/outer_version": merged_version(versions), } diff --git a/yeto/rl/engine/island_placement.py b/yeto/rl/engine/island_placement.py index 2490e4cc4..fff346c5f 100644 --- a/yeto/rl/engine/island_placement.py +++ b/yeto/rl/engine/island_placement.py @@ -13,9 +13,9 @@ ``use-gpus-per-node`` ``--rl-island-use-gpus-per-node`` integer >= 1 ======================== =============================== ================= -Mixed placement (islands with different placements) is allowed only under -inter-island elastic scheduling (``--rl-island-scheduling elastic``): the -legacy (strict) schedule is not part of this change and is refused. A +Mixed placement (islands with different placements) is allowed only in an +island pool (any island pool option, e.g. ``--rl-sync-island-lease-s``): a fixed +roster is not part of this change and is refused. A colocated island has no dedicated rollout or standby GPUs, so in a mixed run it drops the run-wide ``--rl-rollout-gpus`` / ``--rl-standby-gpus`` it would otherwise inherit; an explicit per-island ``standby-gpus`` above 0 on a @@ -176,10 +176,10 @@ def check(args: Any, num_islands: int, *, island_scheduling: str) -> tuple[Islan if is_mixed(placements) and island_scheduling != "elastic": kinds = ", ".join(f"island {p.island} {p.placement}" for p in placements) raise ValueError( - f"mixed island placement ({kinds}) needs --rl-island-scheduling elastic: the " - "legacy (strict) schedule waits for every island each round and is not " - "supported with different placements; use one placement for every island " - "or add --rl-island-scheduling elastic") + f"mixed island placement ({kinds}) needs island pool options: a fixed " + "roster waits for every island each round and is not supported with " + "different placements; use one placement for every island or add island " + "pool options (e.g. --rl-sync-quorum-theta, --rl-sync-island-lease-s)") # One placement on every island: the execution-profile checks give the # usual error. A mixed run names its colocated islands here. for flag, dest in (("--rl-elastic", "rl_elastic"), ("--rl-overlap-eval", "rl_overlap_eval")): diff --git a/yeto/rl/engine/island_status.py b/yeto/rl/engine/island_status.py index 77ea9d0b9..6a94c7787 100644 --- a/yeto/rl/engine/island_status.py +++ b/yeto/rl/engine/island_status.py @@ -1,10 +1,11 @@ -"""Read the elastic syncer's status snapshot (rl-inter-island-scheduling 0.17/0.18). +"""Read the syncer's island-pool status snapshot (rl-inter-island-scheduling 0.17/0.18). -The Rust syncer in elastic mode writes ``status.json`` next to ``--event-tape`` -(schema ``yeto.syncer.elastic-status/v1``, see syncer/src/elastic.rs -``status_json``). This module only reads it: a missing file, unparsable JSON or -a different schema gives ``None`` -- callers then leave scheduling fields -``None`` (never estimated). Legacy mode never calls into here. +With island pool options the Rust syncer writes ``status.json`` next to +``--event-tape`` (schema ``yeto.syncer.elastic-status/v1``, see +syncer/src/server.rs ``write_pool_status``). This module only reads it: a +missing file, unparsable JSON or a different schema gives ``None`` -- callers +then leave scheduling fields ``None`` (never estimated). A fixed roster never +calls into here. """ from __future__ import annotations diff --git a/yeto/rl/engine/leave_reason.py b/yeto/rl/engine/leave_reason.py index 83000f3de..45a33d04f 100644 --- a/yeto/rl/engine/leave_reason.py +++ b/yeto/rl/engine/leave_reason.py @@ -17,7 +17,7 @@ import sys -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason _E = "yeto.rl.engine." @@ -33,7 +33,6 @@ (_E + "controller", "TransactionFailed"), (_E + "controller", "RebuildRefused"), (_E + "trainer_transition", "TrainerTransitionFailed"), - ("yeto.rl.elastic_client", "ElasticRejoinError"), # storage / resume state (D7 Q4) (_E + "cut", "CutError"), (_E + "journal", "JournalError"), @@ -51,7 +50,6 @@ (_E + "multinode", "TopologyError"), (_E + "backends", "UnknownBackend"), (_E + "trainable_state", "UnsupportedLayoutError"), - ("yeto.rl.elastic_client", "ElasticJoinRefused"), # tape only: never a member (D7 Q6) )), (LeaveReason.TRAINING_FAILED, ( (_E + "driver", "DriverError"), # Publication/PolicyIdentity/RoundFailed and checks diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 88a4a4973..b0f8940ed 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -1034,7 +1034,7 @@ def resolve_rl_run_config( rollout_batch_size=args.groups_per_round, n_samples_per_prompt=args.samples_per_group, global_batch=global_batch, - island_scheduling=getattr(args, "rl_island_scheduling", None) or "legacy", + island_scheduling="elastic" if getattr(args, "rl_island_pool", False) else "legacy", requested=getattr(args, "rl_lr_schedule", None) or "auto", ), seed=args.seed, diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 915833f8e..6c93732f5 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -407,8 +407,8 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. help="give island ISLAND its own placement value (repeatable); KEY is " "placement (colocated | fixed-partition), standby-gpus or use-gpus-per-node; " "other islands keep --rl-placement / --rl-standby-gpus / " - "--rl-island-use-gpus-per-node. Different placements need " - "--rl-island-scheduling elastic"), + "--rl-island-use-gpus-per-node. Different placements need island pool " + "options (e.g. --rl-sync-island-lease-s)"), cli_flag("--rl-island-use-gpus-per-node", type=int, default=None, metavar="M", help="rl-multinode-island: provision the --gpu machines (N GPUs/node, billed) but " "give the island only local GPUs 0..M-1 on every node (M <= N): ray start " @@ -631,13 +631,6 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. help="--rl-elastic: declare the --rl-elastic-cells names to the fork as its " "rollout engine cells (placement map rollout_cells: started on the rollout " "GPUs, then stopped on standby GPUs, then stopped unbound; needs fork F-R1)"), - # rl-inter-island-scheduling 0.13: explicit old/new outer-scheduling choice. - # 3.8 strict pause budget inputs (defaults: syncer 900 s, margin 0.5). - # rl-inter-island-scheduling 0.13 (launcher forwards only "elastic") - shared_flag("--rl-island-scheduling", choices=("legacy", "elastic"), default="legacy", - help="legacy (default): existing syncer behaviour, fixed members, every " - "island must arrive; elastic: inter-island scheduling (capacity-weighted " - "stepping, late deltas carried over with a discount, join/leave)"), cli_flag("--rl-island-contract-sha256", default=None, type=sha256_hex_arg, help="secret-handling-hardening: island contract v3 (backend identity bound to " @@ -652,37 +645,51 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. # S17 N16: LR schedule choice (yeto.rl.engine.run_config.resolve_lr_schedule) shared_flag("--rl-lr-schedule", choices=("auto", "linear", "constant"), default="auto", help="island optimizer LR schedule: auto (default) = constant for decoupled " - "and elastic islands, linear decay over global_rounds x optimizer_steps " + "and island-pool islands, linear decay over global_rounds x optimizer_steps " "otherwise; constant = fixed LR in every mode (single island included); " - "linear is refused for decoupled/elastic"), - cli_flag("--rl-quorum-theta", type=float, default=None, metavar="F", - help="elastic: step when arrived capacity >= F x total (default 0.75)"), - cli_flag("--rl-carry-gamma", type=float, default=None, metavar="F", - help="elastic: discount per outer step for a late delta (default 0.5)"), - cli_flag("--rl-soft-deadline-s", type=int, default=None, metavar="S", - help="elastic: soft deadline of an outer round (default: the quorum " - "timeout, syncer default 900)"), + "linear is refused for decoupled and island pools"), + cli_flag("--rl-sync-quorum-theta", type=float, default=None, metavar="F", + help="membership: merge a round once the arrived islands hold F of the " + "capacity the round expects (replaces --rl-sync-quorum). Any membership " + "option (this, --rl-sync-q-min, --rl-sync-soft-deadline-s, " + "--rl-sync-carry-gamma, --rl-sync-island-lease-s, --rl-sync-max-learners, " + "--rl-sync-final-grace-s) puts the islands in a pool with dynamic membership; " + "unset ones stay off. strict-avg, decoupled and full-parameter with " + "--rl-sync-full-parameter-rounds async"), + cli_flag("--rl-sync-q-min", type=int, default=None, metavar="N", + help="membership: minimum arrived islands for a capacity-quorum or " + "soft-deadline merge (default 1)"), + cli_flag("--rl-sync-soft-deadline-s", type=int, default=None, metavar="S", + help="membership: merge what arrived after S seconds even below quorum " + "(default off)"), + cli_flag("--rl-sync-carry-gamma", type=float, default=None, metavar="F", + help="membership: weight discount F^lag of a delta whose base lags the round " + "(default 1: no discount); the largest lag is --rl-sync-max-base-lag"), + shared_flag("--rl-sync-island-lease-s", type=float, default=None, metavar="S", + help="membership: island lease on the syncer (--island-lease-s); an island " + "silent this long is removed and re-JOINs; islands heartbeat every S/3. " + "Default 150 in a pool"), + cli_flag("--rl-sync-max-learners", type=int, default=None, metavar="N", + help="membership: the syncer admits island ids 0..N-1, so islands may join " + "after the launch (default: the launched islands; not with the " + "full-parameter preset, whose session binds the roster)"), + cli_flag("--rl-sync-final-grace-s", type=int, default=None, metavar="S", + help="membership: after the last step the syncer waits this long for the " + "islands still connected to take the final cut, then finishes without the " + "missing ones (default: --rl-sync-soft-deadline-s; unset both: wait for all)"), cli_flag("--rl-elastic-debug-delay-s", default=None, metavar="ISLAND:S[,ISLAND:S]", - help="elastic TEST switch: the given island(s) wait S seconds before " + help="membership TEST switch: the given island(s) wait S seconds before " "submitting each round's delta (bare S = every island); used to trigger " "the soft deadline and carried_over on real hardware"), cli_flag("--rl-elastic-debug-pause", default=None, metavar="ISLAND:AFTER_V:S[,...]", - help="elastic TEST switch: after the island applied outer version >= AFTER_V " - "its link to the syncer goes silent for S seconds (no heartbeats, sends and " - "received frames held; training continues) to exercise lease expiry and " - "automatic re-JOIN on real hardware"), - cli_flag("--rl-final-grace-s", type=int, default=None, metavar="S", - help="elastic: after total steps the syncer keeps answering FINISHED for " - "this long so late islands end cleanly (default: the soft deadline)"), - shared_flag("--rl-island-lease-s", type=float, default=None, metavar="S", - help="elastic: island lease on the syncer (--island-lease-s); an island " - "silent this long is removed; islands heartbeat every S/3. Default 150"), - cli_flag("--rl-q-min", type=int, default=None, metavar="N", - help="elastic: minimum arrived islands for a step (default 1)"), + help="membership TEST switch: after the island applied outer version >= AFTER_V " + "it stops heartbeating and holds its pushes for S seconds (training continues) " + "to exercise lease expiry and automatic re-JOIN on real hardware"), shared_flag("--rl-syncer-epoch", type=int, default=None, metavar="N", - help="elastic, or full-parameter async with --rl-peer-recovery: syncer term " - "(incremented on each syncer restart); passed to the syncer as --syncer-epoch " - "and to every island, whose JOIN/HELLO an older term fences (default 0)", + help="island pool, or full-parameter async with --rl-peer-recovery: syncer " + "term (incremented on each syncer restart); passed to the syncer as " + "--syncer-epoch and to every island's HELLO (default: read from the running " + "syncer, else 0)", island={"default": 0}), cli_flag("--rl-peer-recovery", action="store_true", default=False, help="full-parameter async: an island relaunched after it lost its trainer " @@ -705,8 +712,6 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. shared_flag("--rl-peer-recovery-ttl-s", type=float, default=None, metavar="S", help="--rl-peer-recovery: snapshots older than this are removed (a donor " "or recovering island died); longer than a snapshot transfer (default 3600)"), - cli_flag("--rl-max-carry-lag", type=int, default=None, metavar="N", - help="elastic: oldest late delta still carried over, in outer steps (default 2)"), shared_flag("--rl-elastic-quorum-timeout-s", type=float, default=None, metavar="S", help="--rl-elastic: syncer --quorum-timeout-s (whole seconds), also the " "island's strict pause budget input (default: syncer default 900)"), diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index 3c6d98b12..7ab9ce7b8 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -342,7 +342,7 @@ def check_launch(args: Any) -> None: "buffer and the policy releases live in it)") if getattr(args, "rl_elastic", False): # --rl-elastic is placement reconfiguration (T/R cells), not the island sync: - # synced trainers use --rl-island-scheduling elastic (or strict sync) + # synced trainers use an island pool (or strict sync) raise ValueError(f"--rl-sample-role {role} does not support --rl-elastic " "(placement reconfiguration)") if role in (TRAINER, CONVERTIBLE) and not int(getattr(args, "rl_max_policy_age", 0) or 0) >= 1: diff --git a/yeto/rl/engine/selection.py b/yeto/rl/engine/selection.py index acfa2cde9..3a4d27372 100644 --- a/yeto/rl/engine/selection.py +++ b/yeto/rl/engine/selection.py @@ -76,15 +76,11 @@ def ports_rejections( # The island loop, the LoRA layouts and the rollout engine are causal-LM only. if model_kind != "causal-lm": reasons.append(f"model kind {model_kind!r} (ports supports causal LMs only)") - full_parameter = ( - tuning == "full" - or parameter_mode == "full" - or sync_preset in FULL_PARAMETER_SYNC_PRESETS - ) - if full_parameter and island_scheduling == "elastic": - # P0-D7 follow-up (tasks 1.7b.6): the elastic syncer has no fragment frames yet. - reasons.append("full-parameter training with --rl-island-scheduling elastic") - elif tuning not in ("lora", "full") or parameter_mode not in ("lora", "full"): + if island_scheduling == "elastic" and sync_preset == SAO_STREAMING_SYNC_PRESET: + # The SAO island session has no pool membership (heartbeat, re-JOIN); + # the other presets' limits are checked by the launcher (rl_args). + reasons.append("the sao-streaming-full preset with membership options (island pool)") + if tuning not in ("lora", "full") or parameter_mode not in ("lora", "full"): # Input check: no other parameter mode exists. reasons.append(f"tuning {tuning!r} / parameter mode {parameter_mode!r}") if sync_preset not in SYNC_PRESETS: diff --git a/yeto/rl/engine/trainer_lifecycle.py b/yeto/rl/engine/trainer_lifecycle.py index ec3416aee..9ef492e7b 100644 --- a/yeto/rl/engine/trainer_lifecycle.py +++ b/yeto/rl/engine/trainer_lifecycle.py @@ -40,7 +40,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from . import recovery from .leave_reason import leave_reason_for diff --git a/yeto/rl/engine/trainer_loop.py b/yeto/rl/engine/trainer_loop.py index cdac13003..12949f022 100644 --- a/yeto/rl/engine/trainer_loop.py +++ b/yeto/rl/engine/trainer_loop.py @@ -46,7 +46,7 @@ from dataclasses import replace from typing import Any -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from .trainer_lifecycle import ( BacklogStats, diff --git a/yeto/rl/engine/trainer_remote.py b/yeto/rl/engine/trainer_remote.py index bac6f980b..36d5ad4f6 100644 --- a/yeto/rl/engine/trainer_remote.py +++ b/yeto/rl/engine/trainer_remote.py @@ -47,7 +47,7 @@ from pathlib import Path from typing import Any, Protocol, runtime_checkable -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from . import recovery from .resume import FLAT_MANIFEST, LATEST, STATE_DIR, CheckpointStore diff --git a/yeto/rl/engine/trainer_scheduler.py b/yeto/rl/engine/trainer_scheduler.py index 78791aa0e..1e17193ab 100644 --- a/yeto/rl/engine/trainer_scheduler.py +++ b/yeto/rl/engine/trainer_scheduler.py @@ -29,7 +29,7 @@ from dataclasses import dataclass from typing import Any -from yeto.rl.elastic_client import LeaveReason +from yeto.rl.pool_client import LeaveReason from .trainer_lifecycle import BacklogStats, Decision, Hold, LifecycleState, Start, Stop diff --git a/yeto/rl/pool_client.py b/yeto/rl/pool_client.py new file mode 100644 index 000000000..33a9eb7d3 --- /dev/null +++ b/yeto/rl/pool_client.py @@ -0,0 +1,352 @@ +"""Island side of the island pool: one RL island on the fragment-round syncer +with dynamic membership (leases, re-JOIN, capacity quorum, soft deadline, +carry discount, final grace window). + +The wire protocol is the ordinary fragment protocol (``yeto.protocol``): the +canonical LoRA vector is one f32 fragment (``core.build_avg_layout``). This +module only adds what an island in a dynamic pool needs on top of +:class:`yeto.protocol.SyncerClient`: + +* heartbeats every ``lease_s / 3`` so the syncer keeps the island's lease; +* a pushed delta answers the newest PULL the island has not answered yet, so + a late island's delta is carried into the next round (the syncer applies + ``--max-base-lag`` and the carry discount); +* bases are read from broadcasts, and the final cut (FINAL_MANIFEST) is + acknowledged and handed over as the last base; +* a lease expiry makes the client redial with a new generation (a re-JOIN); + each one goes on the island tape as ``elastic_rejoin``. + +Tape event names (``elastic_*``) are kept: dashboards and the launcher read them. +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass +from enum import IntEnum +from typing import Any + +import torch + +from ..protocol import DTYPE_F32, FragmentLayout, PullRequest, SyncerClient +from ..tensor_io import pack_tensor, unpack_fragment + + +class LeaveReason(IntEnum): + """Why an island left the pool (one byte of MSG_LEAVE; the syncer's + ``pool_leave`` record names it, see ``protocol.rs`` LEAVE_REASON_NAMES). + ``LEASE_EXPIRED`` is assigned by the syncer only.""" + + FINISHED = 0 + LEASE_EXPIRED = 1 + STOPPED = 2 + PREEMPTED = 3 + INFRA_FAILED = 4 + TRAINING_FAILED = 5 + CONFIG_MISMATCH = 6 + + +# Island lease on the syncer; the island heartbeats every lease/3. Same +# default as the launcher's --rl-sync-island-lease-s. +DEFAULT_ISLAND_LEASE_S = 150.0 +# Consecutive failed redials (lease re-JOIN or lost link) before the island gives up. +DEFAULT_MAX_REJOIN_FAILURES = 5 + +DEBUG_DELAY_ENV = "YETO_RL_ELASTIC_DEBUG_DELAY" +DEBUG_PAUSE_ENV = "YETO_RL_ELASTIC_DEBUG_PAUSE" + + +def parse_elastic_debug_delay(spec: str | None) -> dict[int, float]: + """``"ISLAND:SECONDS[,ISLAND:SECONDS]"`` (bare ``SECONDS`` = every island).""" + out: dict[int, float] = {} + for item in (spec or "").split(","): + item = item.strip() + if not item: + continue + island, sep, seconds = item.rpartition(":") + key = int(island) if sep else -1 + value = float(seconds) + if value < 0 or (sep and key < 0): + raise ValueError(f"--rl-elastic-debug-delay-s: bad item {item!r}") + out[key] = value + return out + + +def parse_elastic_debug_pause(spec: str | None) -> dict[int, tuple[int, float]]: + """``"ISLAND:AFTER_V:SECONDS[,...]"``: once the island applied a base with + outer version >= AFTER_V, it stops heartbeating and holds its pushes for + SECONDS while the process keeps training (exercises lease expiry and the + automatic re-JOIN).""" + out: dict[int, tuple[int, float]] = {} + for item in (spec or "").split(","): + item = item.strip() + if not item: + continue + parts = item.split(":") + try: + if len(parts) != 3: + raise ValueError + island, after_v, seconds = int(parts[0]), int(parts[1]), float(parts[2]) + except (ValueError, IndexError): + raise ValueError(f"--rl-elastic-debug-pause: bad item {item!r} (want ISLAND:AFTER_V:SECONDS)") from None + if island < 0 or after_v < 0 or seconds <= 0: + raise ValueError(f"--rl-elastic-debug-pause: bad item {item!r}") + out[island] = (after_v, seconds) + return out + + +class LeaseHeartbeat: + """Keep one island's lease on the syncer: a daemon thread sends a + heartbeat every ``lease_s / 3`` on ``client`` (a :class:`SyncerClient`). + + A lease expiry is handled by the client itself (it redials and the + syncer treats the new generation as a re-JOIN); ``on_tick`` runs before + each heartbeat (e.g. to log re-JOINs) and ``paused`` lets a test hold + the heartbeats. ``progress`` is the local step reported in each one.""" + + def __init__(self, client: Any, lease_s: float, *, + progress: Callable[[], int] = lambda: 0, + paused: Callable[[], bool] = lambda: False, + on_tick: Callable[[], None] = lambda: None, + name: str = "pool-heartbeat") -> None: + if not lease_s > 0: + raise ValueError("island lease must be > 0 seconds") + self.client = client + self.period_s = float(lease_s) / 3.0 + self.errors: list[str] = [] + self._progress = progress + self._paused = paused + self._on_tick = on_tick + self._stop = threading.Event() + self._thread = threading.Thread(target=self._loop, name=name, daemon=True) + + def start(self) -> None: + self._thread.start() + + def stop(self) -> None: + self._stop.set() + if self._thread.is_alive(): + self._thread.join(timeout=5.0) + + def _loop(self) -> None: + while not self._stop.wait(self.period_s): + self._on_tick() + if self._paused(): + continue + try: + self.client.heartbeat(int(self._progress())) + except Exception as error: # noqa: BLE001 - the client's supervisor redials + self.errors.append(f"heartbeat: {error}") + + +class PoolFinished(Exception): + """The syncer finished the run (final cut sent) while this island was + still training or pushing.""" + + def __init__(self, final_outer_version: int, reason: str = "finished") -> None: + super().__init__(f"syncer finished at outer version {final_outer_version} ({reason})") + self.final_outer_version = int(final_outer_version) + self.reason = reason + + +@dataclass(frozen=True) +class JoinAck: + """What the island learned when it joined: the syncer's current version + and whether it joined a run already in progress.""" + + base_version: int + catch_up: bool + + +@dataclass(frozen=True) +class PoolBase: + """A global LoRA vector (flat f32 in layout order) at ``outer_version``.""" + + params: torch.Tensor + outer_version: int + final: bool = False + + +@dataclass +class PoolClientConfig: + syncer_addr: tuple[str, int] + island_id: int + lease_s: float = DEFAULT_ISLAND_LEASE_S + syncer_epoch: int = 0 + session_contract_hash: bytes | None = None + compat_group: str | None = None + wan_streams: int = 4 + join_timeout_s: float = 900.0 + max_rejoin_failures: int = DEFAULT_MAX_REJOIN_FAILURES + + +class PoolIslandClient: + """Heartbeat, PULL bookkeeping, base tracking and finalization for one + island over :class:`SyncerClient`.""" + + FRAGMENT = 0 + + def __init__(self, config: PoolClientConfig, layout: FragmentLayout, *, + on_event: Callable[[dict], None] | None = None, + client_factory: Callable[..., Any] = SyncerClient) -> None: + if layout.num_fragments != 1: + raise ValueError("the island pool syncs the LoRA vector as one fragment") + self.config = config + self.layout = layout + self._on_event = on_event or (lambda ev: None) + self.client = client_factory( + config.syncer_addr, int(config.island_id), layout, dtype=DTYPE_F32, + num_streams=config.wan_streams, session_contract_hash=config.session_contract_hash, + compat_group=config.compat_group, max_reconnects=config.max_rejoin_failures, + syncer_epoch=int(config.syncer_epoch)) + self.latest_base: PoolBase | None = None + # L5.6: (Unix time the base became visible, seconds from receipt to visible) + self.latest_base_timing: tuple[float, float] | None = None + self.inner_step = 0 + self.final_outer_version: int | None = None + self.finished: PoolFinished | None = None + self._permits: dict[int, PullRequest] = {} + self._answered: set[tuple[int, int]] = set() # (global_step, round_attempt) + self._stop = threading.Event() + self._pause_until = 0.0 + self._rejoins_seen = 0 + self._heartbeat = LeaseHeartbeat(self.client, config.lease_s, + progress=lambda: self.inner_step, + paused=self.link_paused, on_tick=self._note_rejoins) + self.errors = self._heartbeat.errors + + # -- lifecycle ----------------------------------------------------------- + def join(self) -> JoinAck: + """Connect (HELLO) and start heartbeating. The base version is the + first broadcast already received (0 before the run starts).""" + self.client.start() + self._heartbeat.start() + self._drain() + version = 0 if self.latest_base is None else self.latest_base.outer_version + return JoinAck(base_version=version, catch_up=version > 0) + + def leave(self, reason: int = LeaveReason.FINISHED) -> None: + """Leave the pool: LEAVE with the reason, then close the connection.""" + self._stop.set() + if not self.client.shutdown.is_set(): + self.client.send_leave(int(reason)) + self.close() + + def close(self) -> None: + self._stop.set() + self.client.close() + self._heartbeat.stop() + + # -- link ---------------------------------------------------------------- + @property + def heartbeat_period_s(self) -> float: + return self.config.lease_s / 3.0 + + def pause_link(self, seconds: float) -> None: + """Test switch: no heartbeats and no pushes for ``seconds``.""" + start = time.time() + self._pause_until = time.monotonic() + float(seconds) + self._on_event({"event": "elastic_debug_pause", "pause_s": float(seconds), "start_unix": start}) + + def end() -> None: + if not self._stop.wait(float(seconds)): + self._on_event({"event": "elastic_debug_pause_end", "pause_s": float(seconds), + "start_unix": start, "end_unix": time.time()}) + threading.Thread(target=end, name="pool-debug-pause", daemon=True).start() + + def link_paused(self) -> bool: + return time.monotonic() < self._pause_until + + def _note_rejoins(self) -> None: + count = int(getattr(self.client, "lease_expiries", 0)) + while self._rejoins_seen < count: + self._rejoins_seen += 1 + self._on_event({"event": "elastic_rejoin", "rejoins": self._rejoins_seen, + "syncer_epoch": int(self.config.syncer_epoch)}) + + # -- bases --------------------------------------------------------------- + def _drain(self) -> None: + self.client.check_health() + self._note_rejoins() + fragment = self.layout.fragments[self.FRAGMENT] + for update in self.client.drain_updates(): + if update.fragment_id != self.FRAGMENT: + raise RuntimeError(f"island pool received fragment {update.fragment_id}") + if self.latest_base is not None and update.version <= self.latest_base.outer_version: + continue + params = unpack_fragment(fragment, update.data, DTYPE_F32) + self._set_base(PoolBase(params, int(update.version)), + max(0.0, time.monotonic() - update.received_at)) + for permit in self.client.drain_pulls(): + if permit.fragment_id != self.FRAGMENT: + raise RuntimeError(f"island pool received a PULL for fragment {permit.fragment_id}") + # A restarted round (higher attempt) needs a fresh answer. + if (permit.global_step, permit.round_attempt) not in self._answered: + previous = self._permits.get(permit.global_step) + if previous is None or permit.round_attempt > previous.round_attempt: + self._permits[permit.global_step] = permit + if self.client.finalizing.is_set() and self.finished is None: + self._take_final_cut() + + def _set_base(self, base: PoolBase, recv_seconds: float) -> None: + self.latest_base = base + self.latest_base_timing = (time.time(), recv_seconds) + + def _take_final_cut(self) -> None: + manifest, fragments = self.client.wait_for_final_fragments() + final = fragments[0] + params = unpack_fragment(self.layout.fragments[self.FRAGMENT], final.data, DTYPE_F32) + self._set_base(PoolBase(params, int(manifest.global_step), final=True), 0.0) + self.client.acknowledge_finalization(manifest) + self.finished = PoolFinished(manifest.global_step) + + def wait_base(self, *, newer_than: int | None, timeout_s: float) -> PoolBase | None: + """The newest base with ``outer_version > newer_than`` (any base when + None), or None after ``timeout_s``. The final cut always qualifies.""" + deadline = time.monotonic() + timeout_s + while True: + self._drain() + base = self.latest_base + if base is not None and (newer_than is None or base.outer_version > newer_than or base.final): + return base + if time.monotonic() >= deadline: + return None + time.sleep(0.05) + + def known_final(self) -> bool: + return self.finished is not None + + def raise_if_finished(self) -> None: + self._drain() + if self.finished is not None: + raise self.finished + + # -- pushes -------------------------------------------------------------- + def elastic_init(self, params: torch.Tensor) -> None: + """Seed the syncer's initial vector (the first one received wins).""" + self.client.send_init(self.FRAGMENT, pack_tensor(params.reshape(-1), DTYPE_F32)) + + def delta_tensor(self, *, base_version: int, c_tokens: int, c_steps: int, update: torch.Tensor, + timeout_s: float = 3600.0) -> int: + """Answer the newest unanswered PULL with ``update`` (local - base at + ``base_version``); returns the round's global step. Raises + :class:`PoolFinished` when the run ends first.""" + deadline = time.monotonic() + timeout_s + while True: + self.raise_if_finished() + if self._permits and not self.link_paused(): + break + if time.monotonic() >= deadline: + raise TimeoutError(f"no PULL from the syncer within {timeout_s:.0f}s") + time.sleep(0.05) + step = max(self._permits) + permit = self._permits.pop(step) + self._permits.clear() # older rounds are closed or will close without this island + self._answered.add((step, permit.round_attempt)) + self.client.push_fragment( + self.FRAGMENT, step, permit.round_attempt, int(base_version), int(self.inner_step), + max(int(c_steps), 1), int(c_tokens), pack_tensor(update.reshape(-1), DTYPE_F32)) + return step + diff --git a/yeto/syncer_profile.py b/yeto/syncer_profile.py index 9ad26f7c1..5b2aa3ec6 100644 --- a/yeto/syncer_profile.py +++ b/yeto/syncer_profile.py @@ -45,18 +45,6 @@ "learner_weight", "require_profile_binding", } -# rl-inter-island-scheduling: optional. Absent or "legacy" = the old encoding -# (byte-identical hash); "elastic" appends the island-scheduling block that the -# syncer's elastic.rs::encode_contract appends. -_ISLAND_SCHEDULING_FIELDS = { - "island_scheduling_mode", - "quorum_theta", - "carry_gamma", - "soft_deadline_s", - "q_min", - "max_carry_lag", -} -_ISLAND_SCHEDULING_DOMAIN = b"yeto-syncer-island-scheduling-v1\0" # Optional. Absent = plain Nesterov (byte-identical hash); a period N >= 2 # appends the block that the syncer's server.rs::semantic_profile_hash appends. _DELAYED_NESTEROV_FIELD = "outer_delayed_nesterov_period" @@ -145,12 +133,6 @@ class SyncerSemanticProfile: max_base_lag: int | None learner_weight: str require_profile_binding: bool - island_scheduling_mode: str = "legacy" - quorum_theta: float | None = None - carry_gamma: float | None = None - soft_deadline_s: int | None = None - q_min: int | None = None - max_carry_lag: int | None = None outer_delayed_nesterov_period: int | None = None expert_token_weights: bool = False @@ -158,36 +140,11 @@ class SyncerSemanticProfile: def from_mapping(cls, raw: object) -> SyncerSemanticProfile: if not isinstance(raw, dict) or not ( _FIELDS <= set(raw) - <= _FIELDS | _ISLAND_SCHEDULING_FIELDS - | {_DELAYED_NESTEROV_FIELD, _EXPERT_TOKEN_WEIGHTS_FIELD} + <= _FIELDS | {_DELAYED_NESTEROV_FIELD, _EXPERT_TOKEN_WEIGHTS_FIELD} ): raise ValueError( "syncer semantic profile fields do not match the v1 schema" ) - mode = raw.get("island_scheduling_mode", "legacy") - params = _ISLAND_SCHEDULING_FIELDS - {"island_scheduling_mode"} - if mode == "legacy": - if params & set(raw): - raise ValueError("island scheduling parameters need island_scheduling_mode=elastic") - scheduling: dict = {} - elif mode == "elastic": - if not params <= set(raw): - raise ValueError("elastic island scheduling needs " + ", ".join(sorted(params))) - theta = _number("quorum_theta", raw["quorum_theta"]) - gamma = _number("carry_gamma", raw["carry_gamma"]) - if not 0.0 < theta <= 1.0 or not 0.0 <= gamma <= 1.0: - raise ValueError("quorum_theta must be in (0,1] and carry_gamma in [0,1]") - scheduling = { - "island_scheduling_mode": "elastic", - "quorum_theta": theta, - "carry_gamma": gamma, - "soft_deadline_s": _integer("soft_deadline_s", raw["soft_deadline_s"], - bits=64, minimum=1), - "q_min": _integer("q_min", raw["q_min"], bits=32, minimum=1), - "max_carry_lag": _integer("max_carry_lag", raw["max_carry_lag"], bits=32), - } - else: - raise ValueError("syncer profile island_scheduling_mode must be legacy|elastic") if raw["schema"] != SYNCER_SEMANTIC_PROFILE_SCHEMA: raise ValueError("unsupported syncer semantic profile schema") delta_correction = raw["delta_correction"] @@ -257,7 +214,6 @@ def from_mapping(cls, raw: object) -> SyncerSemanticProfile: expert_token_weights=_boolean( _EXPERT_TOKEN_WEIGHTS_FIELD, raw.get(_EXPERT_TOKEN_WEIGHTS_FIELD, False) ), - **scheduling, ) profile._validate_general() return profile @@ -336,14 +292,6 @@ def canonical_bytes(self) -> bytes: _encode_optional(encoded, self.max_base_lag, " Date: Thu, 22 Oct 2026 15:00:54 +0000 Subject: [PATCH 139/233] test: fake role streams promise in-order broadcasts like the Miles stream The paired-role fake now declares monotonic_broadcasts, as the Miles role stream does, so the catch-up tests of the asynchronous rounds run against the rule the pairing uses; the one test of a stream without that promise turns it off explicitly. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_full_parameter_policy_versions.py | 3 ++- tests/test_rl_full_parameter_sync.py | 4 ++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_rl_full_parameter_policy_versions.py b/tests/test_rl_full_parameter_policy_versions.py index c2863d17a..0fd0ee0f2 100644 --- a/tests/test_rl_full_parameter_policy_versions.py +++ b/tests/test_rl_full_parameter_policy_versions.py @@ -92,7 +92,8 @@ async def test_rejoin_catch_up_jumps_a_fragment_over_missed_rounds(): @pytest.mark.asyncio async def test_catch_up_waits_without_a_monotonic_stream(): - actor = _PairedRole("actor", [], []) # no ordering promise: version 1 may still arrive + actor = _PairedRole("actor", [], []) + actor.monotonic_broadcasts = False # no ordering promise: version 1 may still arrive pairing = _async_pairing(actor) actor._staged_steps.append(5) diff --git a/tests/test_rl_full_parameter_sync.py b/tests/test_rl_full_parameter_sync.py index 07efb2c09..e7fe47168 100644 --- a/tests/test_rl_full_parameter_sync.py +++ b/tests/test_rl_full_parameter_sync.py @@ -49,6 +49,10 @@ def _steps(progress, role): class _PairedRole: + # Like the Miles role stream: never delivers an older version of a + # fragment after a newer one, so a version jump is a catch-up. + monotonic_broadcasts = True + def __init__( self, role: str, From feb06282721b4ad91bc86f9351dc21c5d6ff31f9 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:08:51 +0000 Subject: [PATCH 140/233] rl: async sync defaults for multi-island RL, strict control mode A launch with two or more islands on strict-avg, decoupled or full-parameter now fills the sync options it leaves unset from --rl-sync-mode (default async): asynchronous full-parameter rounds with at most two fragment rounds in flight, quorum 1 with the adaptive grace window (300 s cap; --grace-gamma/--grace-tau keep their defaults), a base lag of one round, tokens^2/steps island weights, RDA matrix merge for fragmented presets, int4 push (block 128, tensor-aligned blocks) and peer recovery where it applies. The outer step depends on the local rounds per sync H: plain averaging below H=4, Nesterov lr 0.7 momentum 0.9 from H=4. Explicit options always win. --rl-sync-mode strict keeps the previous values; its syncer and island commands are compared byte for byte with the launcher before this change. The syncer records its semantic profile hash next to the checkpoint at every start. The async mode passes --require-profile-record, so a resume whose profile differs, or of a checkpoint written before the record, is refused with the options to pass (--rl-sync-mode strict). Asynchronous full-parameter islands train until the syncer stops them; --rl-lr-schedule auto now gives them a constant LR instead of a linear decay over global_rounds that could reach zero mid-run. Co-Authored-By: Claude Opus 5.5 --- docs/PROTOCOL.md | 40 +++ syncer/src/main.rs | 6 + syncer/src/server.rs | 90 +++++++ tests/test_head_mode.py | 5 +- tests/test_peer_recovery_launch.py | 4 +- tests/test_rl_full_parameter_async.py | 9 +- .../test_rl_full_parameter_policy_versions.py | 4 +- tests/test_rl_launcher.py | 2 + tests/test_rl_launcher_dense_full.py | 9 +- tests/test_rl_miles_adapter_config.py | 12 + tests/test_rl_moe_expert_token_weights.py | 2 +- tests/test_rl_outer_merge_options.py | 9 +- tests/test_rl_outer_merge_real_syncer.py | 9 +- tests/test_rl_sync_membership_options.py | 9 +- tests/test_rl_sync_mode.py | 227 ++++++++++++++++++ tests/test_rl_syncer_shards.py | 2 +- tests/test_run_flags.py | 3 +- yeto/island_overrides.py | 15 +- yeto/launch/rl_args.py | 115 ++++++++- yeto/launch/syncer.py | 11 + yeto/rl/engine/run_config.py | 23 +- yeto/rl/engine/run_flags.py | 89 +++++-- yeto/rl/engine/selection.py | 58 +++++ 23 files changed, 702 insertions(+), 51 deletions(-) create mode 100644 tests/test_rl_sync_mode.py diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md index 22bf8d082..a750c7605 100644 --- a/docs/PROTOCOL.md +++ b/docs/PROTOCOL.md @@ -524,6 +524,46 @@ Every merge appends event-tape fields for protocol/delta semantics, attempt, launch base version, frozen members and generations, responders, staleness, weights, counters, quorum/grace timings, and sync latency. +## RL sync modes (`--rl-sync-mode`) + +A launch with two or more islands on the strict-avg (LoRA), decoupled or +full-parameter preset fills the sync options it leaves unset from the sync +mode. An option given on the command line always wins. One island, +`sao-streaming-full`, `--experimental-rl-sync` and algorithms with a critic +always take the strict values. + +| Option | `async` (default) | `strict` (control) | +| --- | --- | --- | +| `--rl-sync-full-parameter-rounds` (full-parameter) | `async`, at most 2 fragment rounds in flight | `sweep` | +| `--rl-sync-quorum` | 1 | every island | +| `--rl-sync-grace-ms` | 300000: cap of the adaptive grace window | 0 (no grace) | +| `--rl-sync-max-base-lag` | 1 round | 0 | +| `--rl-sync-learner-weight` | `tokens2-over-steps` | `equal` | +| `--rl-sync-matrix-merge` (decoupled, full-parameter) | `rda` (embedding fragments stay averaged) | `avg` | +| `--rl-sync-outer-lr` / `--rl-sync-outer-momentum` | 1.0 / 0 below 4 local rounds per sync (H), 0.7 / 0.9 from H = 4 | the preset's | +| `--rl-sync-push-dtype` (full-parameter) | `int4` (linear, block 128, blocks restart at tensor boundaries) | `f32` | +| `--rl-peer-recovery` (full-parameter) | on with async rounds, a fixed roster, one syncer process, no broadcast deltas and a snapshot store; else off | off | + +The grace window adapts to the islands' measured step time through the +syncer's `--grace-gamma` and `--grace-tau`, which keep their defaults; no +value of gamma is prescribed, the cap only bounds a slow round. The +broadcast overwrites the islands' fragments (merge alpha 0) in both modes. +A quorum of 1 needs a base lag of at least 1 round: otherwise the push of an +island that missed the quorum is a strict failure of a fixed roster. +Asynchronous full-parameter islands train until the syncer stops them, so +`--rl-lr-schedule auto` gives them a constant learning rate. + +`--rl-sync-mode strict` produces the syncer and island commands of the +launcher before the modes existed (tested byte for byte). + +Resume: every syncer start with a checkpoint path writes the semantic +profile hash to `.profile`. The async mode passes +`--require-profile-record`: `--resume` then refuses a checkpoint whose record +differs from the syncer's profile, or a checkpoint without the record: +such a checkpoint was written before the modes, under the strict values, and +must be resumed with `--rl-sync-mode strict` (or the session's original +explicit options). The strict-mode command does not pass the flag. + ## Island pool (dynamic membership) The syncer has one round loop. Without membership options the roster is diff --git a/syncer/src/main.rs b/syncer/src/main.rs index f66470d48..8dc059ac6 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -126,6 +126,11 @@ struct Args { /// Mark the terminal checkpoint as publishable after all merges finish. #[arg(long, default_value_t = false)] mark_final_checkpoint: bool, + /// Resume only a checkpoint whose recorded semantic profile (the + /// `.profile` file written at every start) matches this + /// syncer's; a checkpoint without the record is refused too. + #[arg(long, default_value_t = false)] + require_profile_record: bool, /// Exact local optimizer-step budget per learner (benchmark-only). #[arg(long)] learner_budget_steps: Option, @@ -366,6 +371,7 @@ fn main() -> anyhow::Result<()> { bcast_delta_full_every: args.bcast_delta_full_every, resume: args.resume, mark_final_checkpoint: args.mark_final_checkpoint, + require_profile_record: args.require_profile_record, learner_budget_steps: args.learner_budget_steps, event_tape: args.event_tape, max_base_lag: args.max_base_lag, diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 5e35a6427..8e6c3f4db 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -119,6 +119,9 @@ pub struct Config { pub resume: bool, /// Create the adjacent final marker only for a completed terminal cut. pub mark_final_checkpoint: bool, + /// `--resume` refuses a checkpoint whose `.profile` record is missing or + /// differs from this syncer's semantic profile. Not in the profile hash. + pub require_profile_record: bool, /// Benchmark-only exact local optimizer-step target for every learner. pub learner_budget_steps: Option, /// JSONL event tape: merge records plus sweep ledger reconciliation cuts. @@ -1140,10 +1143,63 @@ fn check_resume_checkpoint(cfg: &Config) -> Result<()> { Ok(()) } +fn profile_record_path(checkpoint: &std::path::Path) -> std::path::PathBuf { + let mut name = checkpoint.as_os_str().to_owned(); + name.push(".profile"); + std::path::PathBuf::from(name) +} + +/// Every start with a checkpoint path records this syncer's semantic profile +/// hash next to the checkpoint. With `--require-profile-record`, `--resume` +/// refuses a checkpoint whose record differs (the merge settings of a session +/// may not change under it) or is missing (written before the record +/// existed). Without the flag nothing is checked: some restarts change the +/// profile on purpose (the learner-budget consolidation restart). +fn check_and_record_profile(cfg: &Config, profile: &[u8; 32]) -> Result<()> { + let Some(checkpoint) = cfg.checkpoint_path.as_ref() else { + return Ok(()); + }; + let record = profile_record_path(checkpoint); + let ours = hex32(profile); + if cfg.resume { + match std::fs::read_to_string(&record) { + Ok(recorded) if recorded.trim() == ours || !cfg.require_profile_record => {} + Ok(recorded) => bail!( + "cannot resume {}: the session was started with other sync settings \ + (semantic profile {}, this syncer {ours}); relaunch with the session's \ + original sync options (a session started before the async defaults: \ + --rl-sync-mode strict)", + checkpoint.display(), + recorded.trim() + ), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + if cfg.require_profile_record { + bail!( + "cannot resume {}: the checkpoint has no semantic profile record, so it \ + was written by an older syncer under the strict defaults; relaunch \ + with --rl-sync-mode strict (or the session's original explicit sync \ + options)", + checkpoint.display() + ); + } + } + Err(error) => { + return Err(error).with_context(|| format!("read {}", record.display())); + } + } + } + let partial = profile_record_path(&checkpoint.with_extension("profile-tmp")); + std::fs::write(&partial, format!("{ours}\n")) + .with_context(|| format!("write {}", partial.display()))?; + std::fs::rename(&partial, &record).with_context(|| format!("write {}", record.display()))?; + Ok(()) +} + pub async fn run(cfg: Config) -> Result<()> { validate_config(&cfg)?; check_resume_checkpoint(&cfg)?; let semantic_profile_hash = cfg.semantic_profile_hash(); + check_and_record_profile(&cfg, &semantic_profile_hash)?; if let Some(steps) = cfg.learner_budget_steps { if !(1..=u32::MAX as u64).contains(&steps) { bail!("--learner-budget-steps must fit a positive u32"); @@ -5981,6 +6037,39 @@ mod tests { registry } + #[test] + fn resume_refuses_a_changed_or_missing_profile_record() { + let dir = std::env::temp_dir().join(format!("yeto-profile-record-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let checkpoint = dir.join("actor.ckpt"); + std::fs::write(&checkpoint, b"x").unwrap(); + let record = profile_record_path(&checkpoint); + let _ = std::fs::remove_file(&record); + let mut config = round_test_config(4); + config.checkpoint_path = Some(checkpoint.clone()); + config.resume = true; + // No requirement: nothing is checked, the record is (re)written. + check_and_record_profile(&config, &[1; 32]).unwrap(); + assert_eq!(std::fs::read_to_string(&record).unwrap().trim(), hex32(&[1; 32])); + check_and_record_profile(&config, &[2; 32]).unwrap(); + check_and_record_profile(&config, &[1; 32]).unwrap(); + // Required: the same profile resumes; another one is refused with the remedy. + config.require_profile_record = true; + check_and_record_profile(&config, &[1; 32]).unwrap(); + let error = check_and_record_profile(&config, &[2; 32]).unwrap_err().to_string(); + assert!(error.contains("other sync settings"), "{error}"); + assert!(error.contains("--rl-sync-mode strict"), "{error}"); + // A missing record is refused when the launcher requires one. + std::fs::remove_file(&record).unwrap(); + let error = check_and_record_profile(&config, &[1; 32]).unwrap_err().to_string(); + assert!(error.contains("no semantic profile record"), "{error}"); + // A fresh start always (re)writes the record. + config.resume = false; + check_and_record_profile(&config, &[3; 32]).unwrap(); + assert_eq!(std::fs::read_to_string(&record).unwrap().trim(), hex32(&[3; 32])); + std::fs::remove_dir_all(&dir).unwrap(); + } + fn round_test_config(total_steps: u64) -> Config { Config { port: 0, @@ -6008,6 +6097,7 @@ mod tests { bcast_delta_full_every: None, resume: false, mark_final_checkpoint: false, + require_profile_record: false, learner_budget_steps: None, event_tape: None, max_base_lag: Some(0), diff --git a/tests/test_head_mode.py b/tests/test_head_mode.py index 4cae7d128..ba88015a5 100644 --- a/tests/test_head_mode.py +++ b/tests/test_head_mode.py @@ -122,12 +122,15 @@ def fake_aws_env(monkeypatch, tmp_path): def test_launch_args_json_roundtrip(): + from yeto.rl.engine.selection import plain + ns = cli.parse_args(LAUNCH_ARGS + ["--cluster-prefix", "rt", "--max-rows", "100"]) payload = json.dumps(cli._serializable_args(ns)) rebuilt = argparse.Namespace(**json.loads(payload)) for key, value in vars(ns).items(): assert getattr(rebuilt, key) == value, key - assert type(getattr(rebuilt, key)) is type(value), key # int/float/bool survive + # int/float/bool survive; a --rl-sync-mode default arrives as its plain value + assert type(getattr(rebuilt, key)) is type(plain(value)), key assert rebuilt.controller == "head" diff --git a/tests/test_peer_recovery_launch.py b/tests/test_peer_recovery_launch.py index b2cec9e27..5523efcc9 100644 --- a/tests/test_peer_recovery_launch.py +++ b/tests/test_peer_recovery_launch.py @@ -7,8 +7,8 @@ import pytest -from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 -from test_rl_full_parameter_async import ASYNC, FULL, TWO_ISLANDS, _flag, _plan +from test_rl_algorithm_provenance import _fake_sky, _no_modal_listing # noqa: F401 +from test_rl_full_parameter_async import ASYNC, FULL, TWO_ISLANDS, _flag, _launcher_args, _plan from yeto.launch.fleet import FleetController from yeto.launch.peer_recovery import PEER_RECOVER_ENV, relaunch_envs diff --git a/tests/test_rl_full_parameter_async.py b/tests/test_rl_full_parameter_async.py index 002f756f1..d160dea59 100644 --- a/tests/test_rl_full_parameter_async.py +++ b/tests/test_rl_full_parameter_async.py @@ -27,7 +27,14 @@ import numpy as np import pytest -from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import _fake_sky, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import _launcher_args as _provenance_launcher_args + + +def _launcher_args(engine, extra=(), gpu="aws:1xa100@us-east-1"): + """These tests pin the strict control (--rl-sync-mode strict); the async + defaults are covered in test_rl_sync_mode.py.""" + return _provenance_launcher_args(engine, ("--rl-sync-mode", "strict", *extra), gpu=gpu) from test_rl_full_parameter_sync import _candidate, _PairedRole, _role_stream_config from yeto.protocol import PullRequest from yeto.rl.engine.full_parameter_sync import FullParameterSyncProfile, RolePairing diff --git a/tests/test_rl_full_parameter_policy_versions.py b/tests/test_rl_full_parameter_policy_versions.py index 0fd0ee0f2..743b97d5e 100644 --- a/tests/test_rl_full_parameter_policy_versions.py +++ b/tests/test_rl_full_parameter_policy_versions.py @@ -260,8 +260,10 @@ def test_full_parameter_synced_trainer_plan(monkeypatch): plan = dry_run_plan(_args(full + ("--rl-sync-full-parameter-rounds", "async"), gpu=gpu, base=synced)) assert plan["syncer"] is not None and len(plan["island_requests"]) == 2 + # Unset rounds default to async in a multi-island run; the strict sweep refuses. + assert dry_run_plan(_args(full, gpu=gpu, base=synced))["syncer"] is not None with pytest.raises(ValueError, match="rounds async"): - _args(full, gpu=gpu, base=synced) + _args(full + ("--rl-sync-mode", "strict"), gpu=gpu, base=synced) # -- the real Rust syncer with a membership lease ------------------------------------- diff --git a/tests/test_rl_launcher.py b/tests/test_rl_launcher.py index cf54afb9a..8add8173f 100644 --- a/tests/test_rl_launcher.py +++ b/tests/test_rl_launcher.py @@ -73,6 +73,8 @@ def _args(extra=()): "--reward-function", "pkg.reward:score", "--trust-remote-code", + "--rl-sync-mode", # these tests pin the strict control + "strict", *extra, ] ) diff --git a/tests/test_rl_launcher_dense_full.py b/tests/test_rl_launcher_dense_full.py index 784ee37ff..781e5da12 100644 --- a/tests/test_rl_launcher_dense_full.py +++ b/tests/test_rl_launcher_dense_full.py @@ -11,7 +11,14 @@ import pytest -from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import _fake_sky, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import _launcher_args as _provenance_launcher_args + + +def _launcher_args(engine, extra=(), gpu="aws:1xa100@us-east-1"): + """These tests pin the strict control (--rl-sync-mode strict); the async + defaults are covered in test_rl_sync_mode.py.""" + return _provenance_launcher_args(engine, ("--rl-sync-mode", "strict", *extra), gpu=gpu) LAYOUT = "e" * 64 TWO_ISLANDS = "modal:2xh100,modal:2xh100" diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index 7cbe30d98..f2073cb19 100644 --- a/tests/test_rl_miles_adapter_config.py +++ b/tests/test_rl_miles_adapter_config.py @@ -684,6 +684,18 @@ def test_elastic_island_lr_schedule_is_constant(): island_scheduling="legacy") == rc.LrSchedule("linear", 6) +def test_async_full_parameter_island_lr_schedule_is_constant(): + """An asynchronous full-parameter island trains until the syncer stops it + (H local rounds per fragment round, more when it is fast under a partial + quorum), so a linear horizon of global_rounds could decay the LR to 0 + while it still trains. The strict sweep keeps the linear schedule.""" + args = _schedule_args(sync_preset="full-parameter", global_rounds=6) + assert rc.resolve_lr_schedule(**args, full_parameter_async=True) == rc.LrSchedule("constant", 6) + assert rc.resolve_lr_schedule(**args) == rc.LrSchedule("linear", 6) + with pytest.raises(ValueError, match="asynchronous full-parameter"): + rc.resolve_lr_schedule(**args, full_parameter_async=True, requested="linear") + + @pytest.mark.parametrize("preset,scheduling", [("strict-avg", "legacy"), ("decoupled", "legacy"), ("strict-avg", "elastic")]) def test_rl_lr_schedule_constant_in_every_mode(preset, scheduling): diff --git a/tests/test_rl_moe_expert_token_weights.py b/tests/test_rl_moe_expert_token_weights.py index 3f0d49c51..6e8214f81 100644 --- a/tests/test_rl_moe_expert_token_weights.py +++ b/tests/test_rl_moe_expert_token_weights.py @@ -313,7 +313,7 @@ def test_launcher_turns_the_option_on_for_syncer_profile_and_islands(): def test_launcher_and_island_refuse_the_option_on_the_sweep(monkeypatch): - from test_rl_algorithm_provenance import _launcher_args + from test_rl_full_parameter_async import _launcher_args # strict control from test_rl_full_parameter_async import FULL, TWO_ISLANDS, _island_args from yeto.launcher import _prepare_rl_args diff --git a/tests/test_rl_outer_merge_options.py b/tests/test_rl_outer_merge_options.py index e671ee6a2..b7e57ef5e 100644 --- a/tests/test_rl_outer_merge_options.py +++ b/tests/test_rl_outer_merge_options.py @@ -12,7 +12,14 @@ import pytest -from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import _fake_sky, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import _launcher_args as _provenance_launcher_args + + +def _launcher_args(engine, extra=(), gpu="aws:1xa100@us-east-1"): + """These tests pin the strict control (--rl-sync-mode strict); the async + defaults are covered in test_rl_sync_mode.py.""" + return _provenance_launcher_args(engine, ("--rl-sync-mode", "strict", *extra), gpu=gpu) TWO_ISLANDS = "modal:2xh100,modal:2xh100" THREE_ISLANDS = "modal:2xh100,modal:2xh100,modal:2xh100" diff --git a/tests/test_rl_outer_merge_real_syncer.py b/tests/test_rl_outer_merge_real_syncer.py index bd33451f0..9b50799be 100644 --- a/tests/test_rl_outer_merge_real_syncer.py +++ b/tests/test_rl_outer_merge_real_syncer.py @@ -24,7 +24,14 @@ import numpy as np import pytest -from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import _fake_sky, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import _launcher_args as _provenance_launcher_args + + +def _launcher_args(engine, extra=(), gpu="aws:1xa100@us-east-1"): + """These tests pin the strict control (--rl-sync-mode strict); the async + defaults are covered in test_rl_sync_mode.py.""" + return _provenance_launcher_args(engine, ("--rl-sync-mode", "strict", *extra), gpu=gpu) from yeto.fragments import MERGE_AVG, MERGE_RDA, Fragment, FragmentLayout from yeto.protocol import DTYPE_F32, SyncerClient from yeto.rl.core import CanonicalTensorSpec, build_rl_fragment_layout diff --git a/tests/test_rl_sync_membership_options.py b/tests/test_rl_sync_membership_options.py index 81067719e..6c11344ad 100644 --- a/tests/test_rl_sync_membership_options.py +++ b/tests/test_rl_sync_membership_options.py @@ -14,7 +14,14 @@ import pytest -from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import _fake_sky, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import _launcher_args as _provenance_launcher_args + + +def _launcher_args(engine, extra=(), gpu="aws:1xa100@us-east-1"): + """These tests pin the strict control (--rl-sync-mode strict); the async + defaults are covered in test_rl_sync_mode.py.""" + return _provenance_launcher_args(engine, ("--rl-sync-mode", "strict", *extra), gpu=gpu) TWO_ISLANDS = "modal:2xh100,modal:2xh100" THREE_ISLANDS = "modal:2xh100,modal:2xh100,modal:2xh100" diff --git a/tests/test_rl_sync_mode.py b/tests/test_rl_sync_mode.py new file mode 100644 index 000000000..4b0d89ed3 --- /dev/null +++ b/tests/test_rl_sync_mode.py @@ -0,0 +1,227 @@ +"""--rl-sync-mode: the async defaults of a multi-island RL run and the strict +control. No Ray, no syncer process. + +The strict mode must produce the commands the launcher produced before the +modes existed; the comparison runs the launcher of ``BASE`` (checked out with +``git archive``) and of this tree on the same command lines. +""" + +from __future__ import annotations + +import json +import os +import shlex +import subprocess +import sys +import tarfile +from io import BytesIO +from pathlib import Path + +import pytest + +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_full_parameter_async import FULL, TWO_ISLANDS + +ROOT = Path(__file__).resolve().parents[1] +BASE = "a7589467" # the last commit before --rl-sync-mode + +DECOUPLED = ("--rl-sync-preset", "decoupled", "--local-rl-rounds-per-sync", "2", + "--fragments", "2", "--fragment-pattern", "binpack") +POOL = ("--rl-sync-quorum-theta", "0.75", "--rl-sync-max-base-lag", "2") +FULL_ASYNC = FULL + ("--rl-sync-full-parameter-rounds", "async") +# Command lines whose strict-mode plan must equal the plan before the modes. +CONTROL_CASES = { + "strict-avg": (), + "strict-avg-pool": POOL, + "strict-avg-k1": ("--rl-sync-quorum", "1", "--rl-sync-max-base-lag", "1"), + "decoupled": DECOUPLED, + "full-parameter-sweep": FULL, + "full-parameter-async": FULL_ASYNC, + "full-parameter-int8": FULL + ("--rl-sync-push-dtype", "int8"), +} + +_PLAN_SCRIPT = r""" +import json, sys, types +sys.path[:0] = [sys.argv[1] + "/tests", sys.argv[1]] +from test_rl_launcher import _Resources, _Storage, _StorageMode, _Task +sys.modules["sky"] = types.SimpleNamespace( + Task=_Task, Resources=_Resources, Storage=_Storage, StorageMode=_StorageMode) +import yeto.launch.modal_islands as modal_islands +modal_islands._list_modal_apps = lambda: [] +from test_rl_algorithm_provenance import _launcher_args +from yeto.launcher import _prepare_rl_args, dry_run_plan +plans = {} +for name, extra in json.loads(sys.argv[2]).items(): + args = _launcher_args("ports", tuple(extra), gpu=sys.argv[3]) + args.controller = "local" + _prepare_rl_args(args) + plans[name] = dry_run_plan(args) +print(json.dumps(plans, sort_keys=True, default=str).replace(sys.argv[1], "")) +""" + + +def _plans(root: Path, cases: dict) -> dict: + env = {**os.environ, "YETO_ISLAND_HMAC_KEY": "k3y", "PYTHONDONTWRITEBYTECODE": "1"} + result = subprocess.run( + [sys.executable, "-c", _PLAN_SCRIPT, str(root), json.dumps(cases), TWO_ISLANDS], + cwd=root, env=env, capture_output=True, text=True, timeout=600) + assert result.returncode == 0, result.stderr[-4000:] + return json.loads(result.stdout.splitlines()[-1]) + + +@pytest.fixture(scope="module") +def base_tree(tmp_path_factory): + archive = subprocess.run(["git", "-C", str(ROOT), "archive", BASE], + capture_output=True) + if archive.returncode != 0: + pytest.skip(f"git archive {BASE} is unavailable here") + root = tmp_path_factory.mktemp("base") + with tarfile.open(fileobj=BytesIO(archive.stdout)) as tar: + tar.extractall(root, filter="data") + return root + + +def test_strict_mode_commands_equal_the_commands_before_the_modes(base_tree): + before = _plans(base_tree, CONTROL_CASES) + strict = _plans(ROOT, {name: (*extra, "--rl-sync-mode", "strict") + for name, extra in CONTROL_CASES.items()}) + for name in CONTROL_CASES: + assert strict[name] == before[name], name + + +# -- the async defaults ------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _hmac_key(monkeypatch): + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") + + +def _plan(extra=(), gpu=TWO_ISLANDS): + from yeto.launcher import _prepare_rl_args, dry_run_plan + + args = _launcher_args("ports", tuple(extra), gpu=gpu) + args.controller = "local" + _prepare_rl_args(args) + return args, dry_run_plan(args) + + +def _flags(command: str, *names: str) -> dict[str, str | None]: + tokens = shlex.split(command.replace("$(", "").replace(")", "")) + return {name: tokens[tokens.index(name) + 1] if name in tokens else None for name in names} + + +SYNCER_FLAGS = ("--quorum", "--grace-ms", "--grace-gamma", "--grace-tau", "--pipeline", + "--outer-lr", "--outer-momentum", "--max-base-lag", "--learner-weight") + + +def test_full_parameter_defaults_to_asynchronous_rounds(): + args, plan = _plan(FULL) + assert args.rl_sync_mode_effective == "async" + command = plan["role_syncers"][0]["command"] + assert "--policy-sweep-fragments" not in command + assert "--require-profile-record" in command + assert _flags(command, *SYNCER_FLAGS) == { + "--quorum": "1", "--grace-ms": "300000", "--grace-gamma": "0.8", "--grace-tau": "2.0", + "--pipeline": "2", "--outer-lr": "1.0", "--outer-momentum": "0.0", + "--max-base-lag": str(args.fragments), "--learner-weight": "tokens2-over-steps"} + for island in plan["island_requests"]: + assert _flags(island["learner_command"], "--full-parameter-rounds", "--sweep-push-dtype", + "--rl-sync-matrix-merge") == { + "--full-parameter-rounds": "async", "--sweep-push-dtype": "int4", + "--rl-sync-matrix-merge": "rda"} + + +@pytest.mark.parametrize(("rounds", "lr", "momentum"), [("3", "1.0", "0.0"), ("4", "0.7", "0.9")]) +def test_async_outer_step_depends_on_the_local_rounds_per_sync(rounds, lr, momentum): + _, plan = _plan(FULL + ("--local-rl-rounds-per-sync", rounds)) + command = plan["role_syncers"][0]["command"] + assert _flags(command, "--outer-lr", "--outer-momentum") == { + "--outer-lr": lr, "--outer-momentum": momentum} + + +def test_explicit_options_override_the_async_defaults(): + args, plan = _plan(FULL + ( + "--rl-sync-quorum", "2", "--rl-sync-grace-ms", "0", "--rl-sync-learner-weight", "equal", + "--rl-sync-max-base-lag", "0", "--rl-sync-outer-lr", "0.5", "--rl-sync-outer-momentum", + "0.5", "--rl-sync-matrix-merge", "avg", "--rl-sync-push-dtype", "bf16", + "--rl-sync-fragment-pipeline", "1", "--local-rl-rounds-per-sync", "8")) + command = plan["role_syncers"][0]["command"] + assert _flags(command, "--quorum", "--grace-ms", "--pipeline", "--outer-lr", + "--outer-momentum", "--max-base-lag", "--learner-weight") == { + "--quorum": "2", "--grace-ms": "0", "--pipeline": "1", "--outer-lr": "0.5", + "--outer-momentum": "0.5", "--max-base-lag": "0", "--learner-weight": "equal"} + island = plan["island_requests"][0]["learner_command"] + assert _flags(island, "--sweep-push-dtype")["--sweep-push-dtype"] == "bf16" + assert "--rl-sync-matrix-merge" not in island + + +def test_explicit_sweep_keeps_the_strict_control(): + args, default = _plan(FULL + ("--rl-sync-full-parameter-rounds", "sweep")) + _, strict = _plan(FULL + ("--rl-sync-mode", "strict")) + assert args.rl_sync_mode_effective == "strict" + assert default == strict + + +def test_sweep_only_option_names_the_strict_mode(): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", FULL + ("--rl-sync-async-checkpoint",), gpu=TWO_ISLANDS) + with pytest.raises(ValueError, match="--rl-sync-mode strict"): + _prepare_rl_args(args) + + +def test_lora_defaults_to_quorum_one_with_grace_and_token_weights(): + args, plan = _plan() + command = plan["role_syncers"][0]["command"] + assert _flags(command, *SYNCER_FLAGS) == { + "--quorum": "1", "--grace-ms": "300000", "--grace-gamma": "0.8", "--grace-tau": "2.0", + "--pipeline": "1", "--outer-lr": "1.0", "--outer-momentum": "0.0", + "--max-base-lag": "1", "--learner-weight": "tokens2-over-steps"} + # strict-avg merges one whole-adapter fragment; RDA is for fragmented presets. + assert args.rl_sync_matrix_merge == "avg" + + +def test_decoupled_defaults_take_rda_and_nesterov_from_four_rounds(): + args, plan = _plan(DECOUPLED[:2] + ("--local-rl-rounds-per-sync", "4") + DECOUPLED[4:]) + command = plan["role_syncers"][0]["command"] + assert _flags(command, "--quorum", "--outer-lr", "--outer-momentum", "--learner-weight") == { + "--quorum": "1", "--outer-lr": "0.7", "--outer-momentum": "0.9", + "--learner-weight": "tokens2-over-steps"} + assert args.rl_sync_matrix_merge == "rda" + + +def test_capacity_quorum_replaces_the_count_quorum_default(): + args, _ = _plan(POOL) + assert args.rl_sync_quorum is None and args.quorum == 1 + + +@pytest.mark.parametrize("extra", [(), FULL]) +def test_one_island_keeps_the_strict_values(extra): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", extra, gpu="modal:2xh100") + args.controller = "local" + try: + _prepare_rl_args(args) + except ValueError: + pass # a one-island preset may be refused for other reasons + assert args.rl_sync_mode_effective == "strict" + assert args.rl_sync_full_parameter_rounds == "sweep" + assert args.rl_sync_push_dtype == "f32" + + +def test_peer_recovery_defaults_on_where_it_applies(): + from yeto.launcher import _prepare_rl_args + + store = ("--rl-checkpoint-store", "/shared/run-store") + args, _ = _plan(FULL + store) + assert args.rl_peer_recovery is True + args, _ = _plan(FULL) # no snapshot store: off + assert args.rl_peer_recovery is False + # Off: the store alone is then refused, which shows recovery stayed off. + for extra in (("--no-rl-peer-recovery",), ("--rl-sync-mode", "strict"), POOL): + args = _launcher_args("ports", FULL + store + extra, gpu=TWO_ISLANDS) + with pytest.raises(ValueError, match="or --rl-peer-recovery"): + _prepare_rl_args(args) + assert args.rl_peer_recovery is False diff --git a/tests/test_rl_syncer_shards.py b/tests/test_rl_syncer_shards.py index 797bf5f59..a2c1d9d83 100644 --- a/tests/test_rl_syncer_shards.py +++ b/tests/test_rl_syncer_shards.py @@ -263,7 +263,7 @@ def test_launcher_runs_one_actor_syncer_per_shard(): ]) def test_launcher_refuses_shards_outside_asynchronous_rounds(extra, message, monkeypatch): from yeto.launcher import _prepare_rl_args - from test_rl_algorithm_provenance import _launcher_args + from test_rl_full_parameter_async import _launcher_args # strict control from test_rl_full_parameter_async import TWO_ISLANDS monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k3y") diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index c7152be8e..8068db135 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -96,6 +96,7 @@ def test_island_divergences_are_the_declared_ones(): "--reward-function": {"required": True}, "--over-sampling-batch-size": {"required": True}, "--rl-syncer-epoch": {"default": 0}, + "--rl-sync-matrix-merge": {"default": "avg"}, "--rollout-max-response-len": {"default": None, "required": True}, } aliases = {f.option: f.cli_aliases for f in rl_flags() if f.cli_aliases} @@ -250,7 +251,7 @@ def test_forward_tables_name_flags_of_both_layers(): # flag renamed in the schema but still spelled the old way in the launcher fails. _NON_RL_LITERALS = frozenset({ "--abbrev-ref", "--adapter-dir", "--add", "--address", "--allow-unattested-legacy-adapter", - "--allow-unauthenticated-islands", "--async-checkpoint", "--bcast-delta-full-every", "--bind", "--cap-add", "--carry-gamma", + "--allow-unauthenticated-islands", "--require-profile-record", "--async-checkpoint", "--bcast-delta-full-every", "--bind", "--cap-add", "--carry-gamma", "--checkpoint-every", "--checkpoint-path", "--codex-harness-binary", "--config-json", "--create", "--dashboard-host", "--max-learners", "--dearmor", "--depth", "--detach", "--dtype", "--entrypoint", "--expected-island-contract", "--expert-token-weights", "--final-ack-timeout-s", "--final-grace-s", "--follow", "--force", diff --git a/yeto/island_overrides.py b/yeto/island_overrides.py index 6e67729a5..a6c33cdd3 100644 --- a/yeto/island_overrides.py +++ b/yeto/island_overrides.py @@ -134,9 +134,10 @@ def _old_value(args, key: str): if key == "rl_lr_schedule": value = value or "auto" if value == "auto": # resolved as run_config.resolve_lr_schedule does - from .launch.syncer import island_pool + from .launch.syncer import island_pool, rl_full_parameter_async - run_until_stop = getattr(args, "sync_preset", None) == "decoupled" or island_pool(args) + run_until_stop = (getattr(args, "sync_preset", None) == "decoupled" + or island_pool(args) or rl_full_parameter_async(args)) return "constant" if run_until_stop else "linear" return value if key == "rl_max_policy_age": @@ -155,11 +156,13 @@ def check_island_args(args) -> None: sched = getattr(args, "rl_lr_schedule", None) or "auto" _lr(sched) if sched == "linear": - from .launcher import island_pool + from .launch.syncer import island_pool, rl_full_parameter_async - if getattr(args, "sync_preset", None) == "decoupled" or island_pool(args): - raise ValueError("--rl-lr-schedule linear needs a known local step count; decoupled and " - "island pool islands run until the syncer stops them (use auto or constant)") + if (getattr(args, "sync_preset", None) == "decoupled" or island_pool(args) + or rl_full_parameter_async(args)): + raise ValueError("--rl-lr-schedule linear needs a known local step count; decoupled, " + "island pool and asynchronous full-parameter islands run until the " + "syncer stops them (use auto or constant)") from .launcher import _rl_backend_module from .rl.engine.policy_age import validate_limit diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 7fa3d9a7c..d3e8ad90d 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -9,7 +9,14 @@ from pathlib import Path from ..gpu_spec import parse_gpu_spec -from ..rl.engine.selection import FULL_PARAMETER_SYNC_PRESET, SAO_STREAMING_SYNC_PRESET +from ..rl.engine.selection import ( + ASYNC_GRACE_CAP_MS, + ASYNC_NESTEROV_MIN_H, + FULL_PARAMETER_SYNC_PRESET, + SAO_STREAMING_SYNC_PRESET, + is_sync_mode_default, + plain, +) from . import REPO_ROOT from .rl_backend import ( _rl_backend_image_ok, _rl_backend_module, _rl_callable, _rl_miles_function, @@ -278,7 +285,12 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: max_policy_age=getattr(args, "rl_max_policy_age", 0) or 0, fragment_pipeline=fragment_pipeline, fragments=args.fragments) if problems: - raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + "; ".join(problems)) + hint = (" (a multi-island run defaults to asynchronous rounds; --rl-sync-mode strict " + "or --rl-sync-full-parameter-rounds sweep selects the strict sweep)" + if async_rounds and getattr(args, "rl_sync_mode_effective", "strict") == "async" + else "") + raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + + "; ".join(problems) + hint) args.quorum = num_islands args.grace_ms = 0 args.pipeline = fragment_pipeline @@ -299,6 +311,83 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: args.outer_momentum = 0.0 +# The options --rl-sync-mode fills when the command leaves them unset: dest +# -> the async-mode value. None-default options (quorum, grace, outer step) +# are unset when None; the others carry a marked default (selection.py). +_ASYNC_MODE_VALUES = { + "rl_sync_quorum": 1, + "rl_sync_grace_ms": ASYNC_GRACE_CAP_MS, + "rl_sync_max_base_lag": 1, + "rl_sync_learner_weight": "tokens2-over-steps", +} +_ASYNC_MODE_FULL_PARAMETER_VALUES = { + "rl_sync_full_parameter_rounds": "async", + "rl_sync_push_dtype": "int4", +} +# Marked defaults; the strict mode keeps their (strict) value. +_SYNC_MODE_MARKED = ( + "rl_sync_full_parameter_rounds", "rl_sync_fragment_pipeline", "rl_sync_push_dtype", + "rl_sync_learner_weight", "rl_sync_max_base_lag", "rl_sync_matrix_merge", +) + + +def rl_sync_mode_applies(args) -> bool: + """The async mode fills unset options only for a run with two or more + islands on strict-avg, decoupled or full-parameter, without a critic or + --experimental-rl-sync; every other run keeps the strict values.""" + return (getattr(args, "rl_sync_preset", "strict-avg") in ( + "strict-avg", "decoupled", FULL_PARAMETER_SYNC_PRESET) + and rl_island_count(args) >= 2 + and not getattr(args, "experimental_rl_sync", False) + and not rl_needs_critic(args)) + + +def _apply_rl_sync_mode(args) -> None: + """Resolve ``--rl-sync-mode`` into concrete option values. + + Only options the command left unset change; an explicit option always + wins. Afterwards ``args.rl_sync_mode_effective`` is the mode that took + effect ("strict" when the async mode does not apply) and no option keeps + its mark. Args never resolved (built in code, or parsed without the RL + checks) read as strict: every marked default equals its strict value.""" + mode = getattr(args, "rl_sync_mode", "strict") + explicit_sweep = (getattr(args, "rl_sync_preset", None) == FULL_PARAMETER_SYNC_PRESET + and getattr(args, "rl_sync_full_parameter_rounds", "sweep") == "sweep" + and not is_sync_mode_default(args.rl_sync_full_parameter_rounds)) + # An explicit full-parameter sweep is the strict control profile. + if mode == "async" and rl_sync_mode_applies(args) and not explicit_sweep: + values = dict(_ASYNC_MODE_VALUES) + nesterov = int(getattr(args, "local_rl_rounds_per_sync", 1)) >= ASYNC_NESTEROV_MIN_H + values["rl_sync_outer_lr"] = 0.7 if nesterov else 1.0 + values["rl_sync_outer_momentum"] = 0.9 if nesterov else 0.0 + if args.rl_sync_preset == FULL_PARAMETER_SYNC_PRESET: + values.update(_ASYNC_MODE_FULL_PARAMETER_VALUES) + # Quantization blocks restart at tensor boundaries by default + # (--rl-sync-push-no-tensor-blocks turns that off). + if args.rl_sync_preset in ("decoupled", FULL_PARAMETER_SYNC_PRESET): + values["rl_sync_matrix_merge"] = "rda" + # Peer recovery: _default_peer_recovery, right after this function. + if getattr(args, "rl_sync_quorum_theta", None) is not None: + values.pop("rl_sync_quorum") # the capacity quorum replaces the count + for dest, value in values.items(): + current = getattr(args, dest, None) + if current is None or is_sync_mode_default(current): + setattr(args, dest, value) + if (args.rl_sync_preset == FULL_PARAMETER_SYNC_PRESET + and args.rl_sync_full_parameter_rounds == "async" + and is_sync_mode_default(getattr(args, "rl_sync_fragment_pipeline", None))): + # At most two fragment rounds in flight; the fragment count is + # checked (and may be 1) in _prepare_dense_full_sweep. + args.rl_sync_fragment_pipeline = min(FULL_PARAMETER_ASYNC_MAX_PIPELINE, + max(int(getattr(args, "fragments", 1) or 1), 1)) + else: + mode = "strict" + for dest in _SYNC_MODE_MARKED: + if hasattr(args, dest): + setattr(args, dest, plain(getattr(args, dest))) + args.rl_sync_mode_effective = mode + + # S27 P1: (dest, option, preset-keeping default) of the outer-merge options. _RL_OUTER_MERGE_OPTIONS = ( ("rl_sync_quorum", "--rl-sync-quorum", None), @@ -311,6 +400,26 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: ) +def _default_peer_recovery(args) -> None: + """``--rl-peer-recovery`` left unset: on in the async mode where it applies + (full-parameter async rounds, a fixed roster, one syncer process, no + broadcast deltas, a snapshot store), off otherwise; the strict mode and + ``--no-rl-peer-recovery`` keep it off.""" + if getattr(args, "rl_peer_recovery", None) is not None: + return + args.rl_peer_recovery = False + if (getattr(args, "rl_sync_mode_effective", "strict") != "async" + or not rl_full_parameter_async(args) or island_pool(args) + or syncer_shard_count(args) > 1 + or getattr(args, "rl_sync_bcast_delta_full_every", 0)): + return + from .peer_recovery import peer_recovery_store + + args.rl_peer_recovery = True + if peer_recovery_store(args) is None: + args.rl_peer_recovery = False + + def _apply_rl_outer_merge_options(args, *, num_islands: int) -> None: """S27 P1: override the outer-merge settings the preset just fixed. @@ -435,6 +544,8 @@ def _prepare_rl_args( return _prepare_node_inventory(args) + _apply_rl_sync_mode(args) + _default_peer_recovery(args) from ..adapter_lifecycle import directory_sha256, selected_parent diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index e3b178066..ae4ffb43a 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -549,6 +549,16 @@ def _rl_expert_token_weights(args) -> bool: return bool(getattr(args, "rl_sync_expert_token_weights", False)) +def _syncer_profile_record(args) -> str: + """--rl-sync-mode async: the syncer resumes a checkpoint only when the + checkpoint's recorded semantic profile matches its own, so a session + started with other sync settings (e.g. before the async defaults) is + refused with the flags to pass instead of resuming under new ones. The + strict mode omits the flag, which keeps its command unchanged.""" + async_mode = getattr(args, "rl_sync_mode_effective", "strict") == "async" + return " --require-profile-record" if async_mode else "" + + def _rl_role_syncer_command(args, num_learners: int, binary: str, syncer: RoleSyncer, shard_plan: tuple[int, int] | None = None) -> str: @@ -581,6 +591,7 @@ def _rl_role_syncer_command(args, num_learners: int, binary: str, f"{' --expert-token-weights' if _rl_expert_token_weights(args) else ''}" f" --checkpoint-path {syncer.checkpoint}" f" --checkpoint-every {_syncer_checkpoint_every(args)}" + f"{_syncer_profile_record(args)}" f" {_resume_if_exists(syncer.checkpoint)}" f" --event-tape {syncer.event_tape}" f"{_syncer_weighted_fold(args)}" diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index b0f8940ed..c0a3d4b09 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -353,14 +353,18 @@ def resolve_lr_schedule( global_batch: int, island_scheduling: str = "legacy", requested: str = "auto", + full_parameter_async: bool = False, ) -> LrSchedule | None: """The island's LR schedule (design D1-D3), ``--rl-lr-schedule`` (S17 N16). - ``auto`` (default) keeps the mode rule: decoupled and elastic islands run - until the syncer stops them (an elastic island that re-JOINs needs local - rounds past ``global_rounds``; S17 N5 a73ab1b2), so any finite linear - horizon can reach zero mid-run -> constant; strict-avg / full-parameter / - single island -> linear over ``global_rounds * optimizer_steps``. + ``auto`` (default) keeps the mode rule: decoupled, island-pool and + asynchronous full-parameter islands run until the syncer stops them (an + island that re-JOINs needs local rounds past ``global_rounds``; S17 N5 + a73ab1b2; an asynchronous full-parameter island trains H local rounds per + fragment round and, with a partial quorum, as many rounds as it is fast), + so any finite linear horizon can reach zero mid-run -> constant; + strict-avg / full-parameter sweep / single island -> linear over + ``global_rounds * optimizer_steps``. ``constant`` holds the LR in every mode (e.g. a single island compared with a fixed-LR baseline; S17 N4 had to approximate it with a 20000-step linear horizon). ``linear`` is refused where the horizon is unknown. @@ -373,11 +377,13 @@ def resolve_lr_schedule( if eval_only: return None horizon = global_rounds * optimizer_steps - run_until_stop = sync_preset == "decoupled" or island_scheduling == "elastic" + run_until_stop = (sync_preset == "decoupled" or island_scheduling == "elastic" + or full_parameter_async) if requested == "linear" and run_until_stop: raise ValueError( - "--rl-lr-schedule linear needs a known local step count; decoupled and " - "elastic islands run until the syncer stops them (use auto or constant)") + "--rl-lr-schedule linear needs a known local step count; decoupled, " + "island-pool and asynchronous full-parameter islands run until the syncer " + "stops them (use auto or constant)") if requested == "constant" or (requested == "auto" and run_until_stop): # decay_iters only satisfies Megatron's ``lr_decay_steps > 0``; a # constant schedule never reads it. @@ -1036,6 +1042,7 @@ def resolve_rl_run_config( global_batch=global_batch, island_scheduling="elastic" if getattr(args, "rl_island_pool", False) else "legacy", requested=getattr(args, "rl_lr_schedule", None) or "auto", + full_parameter_async=getattr(args, "full_parameter_rounds", "sweep") == "async", ), seed=args.seed, critic=resolve_critic_run_config(args), diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 6c93732f5..dad8e9859 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -37,8 +37,12 @@ from yeto.rl.engine.selection import ( FULL_PARAMETER_ROUNDS, RL_MATRIX_MERGES, + RL_SYNC_MODES, SYNC_PRESETS, + ASYNC_GRACE_CAP_MS, + ASYNC_NESTEROV_MIN_H, canonical_sync_preset, + sync_mode_default, ) CLI = "cli" @@ -644,10 +648,11 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "not support the requested limit is refused before launch"), # S17 N16: LR schedule choice (yeto.rl.engine.run_config.resolve_lr_schedule) shared_flag("--rl-lr-schedule", choices=("auto", "linear", "constant"), default="auto", - help="island optimizer LR schedule: auto (default) = constant for decoupled " - "and island-pool islands, linear decay over global_rounds x optimizer_steps " + help="island optimizer LR schedule: auto (default) = constant for decoupled, " + "island-pool and asynchronous full-parameter islands (they train until the " + "syncer stops them), linear decay over global_rounds x optimizer_steps " "otherwise; constant = fixed LR in every mode (single island included); " - "linear is refused for decoupled and island pools"), + "linear is refused where islands train until the syncer stops them"), cli_flag("--rl-sync-quorum-theta", type=float, default=None, metavar="F", help="membership: merge a round once the arrived islands hold F of the " "capacity the round expects (replaces --rl-sync-quorum). Any membership " @@ -691,10 +696,13 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "--syncer-epoch and to every island's HELLO (default: read from the running " "syncer, else 0)", island={"default": 0}), - cli_flag("--rl-peer-recovery", action="store_true", default=False, + cli_flag("--rl-peer-recovery", action=argparse.BooleanOptionalAction, default=None, help="full-parameter async: an island relaunched after it lost its trainer " "state restores the inner optimizer state from a peer island's snapshot, " - "then takes every fragment from the syncer"), + "then takes every fragment from the syncer. Default: on in --rl-sync-mode " + "async where it applies (full-parameter async rounds, a fixed island roster, " + "one syncer process, no broadcast deltas, a snapshot store from " + "--rl-peer-recovery-store or --rl-checkpoint-store), else off"), cli_flag("--rl-recovery-bound-steps", type=int, default=None, metavar="H", help="--rl-peer-recovery: bound on one recovery in syncer steps; the syncer " "records a recovery past it as missed (default: measured, not checked)"), @@ -878,6 +886,27 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. _SYNC = ( cli_flag("--local-rl-rounds-per-sync", type=int, default=1), + cli_flag( + "--rl-sync-mode", + choices=RL_SYNC_MODES, + default="async", + help="how a run with two or more islands fills the sync options left unset " + "(strict-avg, decoupled and full-parameter; one island, SAO, a critic and " + "--experimental-rl-sync keep the strict values). async (default): " + "full-parameter asynchronous sharded rounds (--rl-sync-full-parameter-rounds " + "async, at most 2 fragment rounds in flight), --rl-sync-quorum 1 with the " + "adaptive grace window (--rl-sync-grace-ms %d cap; --grace-gamma/--grace-tau keep " + "their defaults, no gamma value is prescribed), --rl-sync-max-base-lag 1, " + "tokens2-over-steps island weights, rda matrix merge (decoupled, full-parameter), " + "outer step by local rounds per sync H: plain averaging (lr 1.0, momentum 0) " + "below H=%d, Nesterov lr 0.7 momentum 0.9 from H=%d on, broadcast overwrite " + "(merge alpha 0), " + "int4 push (full-parameter, block 128). strict: the control, every island " + "per round, equal weights, plain average, f32; its commands equal the " + "commands from before the modes. An option given explicitly always wins. A " + "session started in one mode resumes only with the same settings" + % (ASYNC_GRACE_CAP_MS, ASYNC_NESTEROV_MIN_H, ASYNC_NESTEROV_MIN_H), + ), cli_flag( "--rl-sync-preset", choices=SYNC_PRESETS, @@ -893,21 +922,23 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. cli_flag( "--rl-sync-fragment-pipeline", type=int, - default=1, + default=sync_mode_default(1), help="full-parameter only: fragment rounds the syncer gathers at once " - "(1..--fragments). Above 1, fragment k+1 uploads while fragment k merges " - "and downloads; merge results are unchanged (default 1 = serial)", + "(1..--fragments; asynchronous rounds: 1 or 2). Above 1, fragment k+1 uploads " + "while fragment k merges and downloads (default: min(2, --fragments) with " + "asynchronous rounds, else 1 = serial)", ), cli_flag( "--rl-sync-push-dtype", choices=("f32", "bf16", "int8", "int4", "int4-log"), - default="f32", + default=sync_mode_default("f32"), help="full-parameter only: wire dtype of the islands' pushed outer deltas " "(pseudo-gradients). bf16 halves upload bytes, int8 (block absmax scale, see " "--rl-sync-push-int8-block) quarters them, int4 (linear, 15 levels per block " "absmax scale, see --rl-sync-push-int4-block) and int4-log (the SFT q4 codec: " "256-value blocks, log2 levels) cut them to ~1/7.5-1/7.8; the syncer keeps " - "f32 masters, momentum and f32 FINAL (default f32 = exact)", + "f32 masters, momentum and f32 FINAL (default: int4 in --rl-sync-mode async, " + "f32 = exact in --rl-sync-mode strict)", ), cli_flag( "--rl-sync-overlap", @@ -983,9 +1014,10 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. cli_flag( "--rl-sync-full-parameter-rounds", choices=FULL_PARAMETER_ROUNDS, - default="sweep", + default=sync_mode_default("sweep"), help="full-parameter only (S27 P2): sweep = strict policy sweeps (every " - "island per fragment round, plain average; the control, default); async = " + "island per fragment round, plain average; the control, default in " + "--rl-sync-mode strict); async = (default in --rl-sync-mode async) " "asynchronous fragment rounds: H = --local-rl-rounds-per-sync, at most " "--rl-sync-fragment-pipeline (1 or 2) fragment rounds in flight, the " "decoupled preset's outer Nesterov (0.7, 0.9), and the --rl-sync-quorum " @@ -997,7 +1029,7 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. default=None, metavar="K", help="strict-avg / decoupled / full-parameter async: islands whose pushes close a round " - "(1..islands; default every island)", + "(1..islands; default 1 in --rl-sync-mode async, every island in strict)", ), cli_flag( "--rl-sync-grace-ms", @@ -1005,14 +1037,17 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. default=None, metavar="MS", help="strict-avg / decoupled / full-parameter async: cap of the adaptive post-quorum grace window " - "(--grace-gamma, --grace-tau); default 0 = no grace", + "(--grace-gamma, --grace-tau; the window adapts to the islands' step time, " + "the cap only bounds it); default %d in --rl-sync-mode async, 0 = no grace in strict" + % ASYNC_GRACE_CAP_MS, ), cli_flag( "--rl-sync-learner-weight", choices=("equal", "tokens2-over-steps"), - default="equal", + default=sync_mode_default("equal"), help="strict-avg / decoupled / full-parameter async: island weight in the outer merge; " - "tokens2-over-steps weighs by c_tokens^2 / c_steps (default equal)", + "tokens2-over-steps weighs by c_tokens^2 / c_steps (default tokens2-over-steps in " + "--rl-sync-mode async, equal in strict)", ), cli_flag( "--rl-sync-expert-token-weights", @@ -1026,10 +1061,12 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. cli_flag( "--rl-sync-max-base-lag", type=int, - default=0, + default=sync_mode_default(0), metavar="N", help="strict-avg / decoupled / full-parameter async: how many rounds a push's base " - "may lag the fragment being merged (default 0 = late pushes are refused). One round " + "may lag the fragment being merged (default 1 in --rl-sync-mode async, so a delta " + "that misses the quorum merges into the next round; 0 = late pushes are refused, " + "the strict default). One round " "merges every fragment once; the launcher passes N x fragments to the syncer, whose " "--max-base-lag counts versions", ), @@ -1038,22 +1075,28 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. type=float, default=None, help="strict-avg / decoupled / full-parameter async: outer Nesterov learning rate " - "(default: the preset's, 1.0 strict-avg, 0.7 decoupled and full-parameter async)", + "(default in --rl-sync-mode async: 1.0 below %d local rounds per sync, 0.7 " + "from there; strict: the preset's, 1.0 strict-avg, 0.7 decoupled and " + "full-parameter async)" % ASYNC_NESTEROV_MIN_H, ), cli_flag( "--rl-sync-outer-momentum", type=float, default=None, help="strict-avg / decoupled / full-parameter async: outer Nesterov momentum in [0, 1) " - "(default: the preset's, 0.0 strict-avg, 0.9 decoupled and full-parameter async)", + "(default in --rl-sync-mode async: 0 below %d local rounds per sync, 0.9 " + "from there; strict: the preset's, 0.0 strict-avg, 0.9 decoupled and " + "full-parameter async)" % ASYNC_NESTEROV_MIN_H, ), shared_flag( "--rl-sync-matrix-merge", choices=RL_MATRIX_MERGES, - default="avg", + default=sync_mode_default("avg"), + island={"default": "avg"}, help="decoupled and full-parameter async only: avg merges every fragment by weighted averaging " - "(default); rda merges fragments without embedding tensors by " - "radial-directional averaging, embedding fragments stay averaged", + "(default in --rl-sync-mode strict); rda (default in --rl-sync-mode async) merges " + "fragments without embedding tensors by radial-directional averaging, embedding " + "fragments stay averaged", ), cli_flag( "--rl-initial-adapter", diff --git a/yeto/rl/engine/selection.py b/yeto/rl/engine/selection.py index 3a4d27372..b19316f1d 100644 --- a/yeto/rl/engine/selection.py +++ b/yeto/rl/engine/selection.py @@ -34,6 +34,64 @@ # is the strict policy sweep (every island per round, plain average, the # control); "async" runs asynchronous sharded fragment rounds. FULL_PARAMETER_ROUNDS = ("sweep", "async") +# How a multi-island RL run fills the sync options the command leaves unset. +# "async" (default): asynchronous sharded rounds, quorum 1 with the adaptive +# grace window, tokens^2/steps island weights, RDA matrix merge, outer +# Nesterov, int4 push. "strict": the strict-roster control (every island per +# round, plain average, f32), i.e. the values each option had before the +# modes existed. An option given on the command line always wins. +RL_SYNC_MODES = ("async", "strict") +# The adaptive grace window's cap in the async mode. The window itself is +# computed from the islands' measured step time (--grace-gamma, --grace-tau); +# the cap only bounds a slow round, it does not fix the wait. +ASYNC_GRACE_CAP_MS = 300_000 +# The async mode's outer step needs enough local progress per merge: with +# fewer than this many local RL rounds per sync (H), the default outer step is +# plain averaging (lr 1.0, momentum 0); from H on it is Nesterov (lr 0.7, +# momentum 0.9). On GPU, 2 full-parameter islands with H = 1 diverged under +# Nesterov 0.7/0.9 and trained normally with plain averaging. +ASYNC_NESTEROV_MIN_H = 4 + + +class SyncModeDefault: + """Marks an argparse default that ``--rl-sync-mode`` may replace. + + The marked value equals the strict-mode value, so code that reads the + option before the mode is resolved sees the strict default.""" + + __slots__ = () + + +class _SyncModeDefaultStr(SyncModeDefault, str): + __slots__ = () + + +class _SyncModeDefaultInt(SyncModeDefault, int): + __slots__ = () + + +def sync_mode_default(value: str | int) -> str | int: + """``value`` marked as a default the sync mode resolves.""" + + if isinstance(value, bool) or not isinstance(value, (str, int)): + raise TypeError(f"no sync-mode default for {value!r}") + return (_SyncModeDefaultStr if isinstance(value, str) else _SyncModeDefaultInt)(value) + + +def is_sync_mode_default(value: object) -> bool: + """True when ``value`` is still the unset (marked) default.""" + + return isinstance(value, SyncModeDefault) + + +def plain(value: object) -> object: + """``value`` without its sync-mode mark.""" + + if isinstance(value, str): + return str(value) + if isinstance(value, int) and not isinstance(value, bool): + return int(value) + return value def canonical_sync_preset(value: str) -> str: From a08f1781f72e203777e9b7066f0e714f39a616ba Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:43:45 +0000 Subject: [PATCH 141/233] rl: a restarted full-parameter island joins a running async pool How it is wired: - Miles role stream: in an island pool with asynchronous rounds (may_rejoin: lease + partial quorum) a fresh process accepts the syncer's catch-up as its initial cut. It keeps the newest version of each fragment, overwrites the checkpoint weights with those values and takes them as anchors, so its first push is measured from the pool's current cut. Only island 0 sends INIT; the running syncer admits the new connection generation under the membership rules (supersede or re-JOIN) and resends every fragment. Full-parameter islands may be droppable with async rounds. - _apply_authoritative(optimizer_proof=False): a fresh Megatron trainer has no optimizer-state proof before its first local step (it records the last rollout ID only there). Two applies skip the comparison: the initial catch-up of a restarted island, and the first apply after a peer snapshot restore (the restore is checked against the donor's snapshot digest instead). The restore path previously asked Megatron for proofs and failed with "last rollout ID is unavailable". - A restarted island that reaches the pool after its final cut stops at once instead of waiting for broadcasts until the timeout. - Test-only fault: exit_at_round / exit_islands make the island process exit with 137 before that rollout, only in its first Modal container (the runner exports YETO_MODAL_INCARNATION from the tape volume markers), so Modal's re-run is a fresh process that must rejoin. The keys combine with round_delay_s / round_delay_islands. Tests: real syncer with raw clients; real syncer with the Miles role stream and chunked adapter (old process closed or silent); fault keys and first-container exit; peer recovery with a trainer that raises like Megatron before its first local step. Co-Authored-By: Claude Opus 5.5 --- docs/PROTOCOL.md | 11 +- tests/test_peer_recovery_miles.py | 13 +- tests/test_rl_driver_profiles.py | 35 +++ .../test_rl_full_parameter_restart_rejoin.py | 118 ++++++++ tests/test_rl_miles_restart_rejoin.py | 264 ++++++++++++++++++ tests/test_spot_droppable.py | 7 +- yeto/cloud/droppable.py | 10 +- yeto/modal_runner.py | 9 +- .../miles/full_parameter/role_stream.py | 101 +++++-- .../miles/models/chunked_full_parameter.py | 12 +- yeto/rl/engine/driver.py | 33 ++- 11 files changed, 575 insertions(+), 38 deletions(-) create mode 100644 tests/test_rl_full_parameter_restart_rejoin.py create mode 100644 tests/test_rl_miles_restart_rejoin.py diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md index a750c7605..211e0b769 100644 --- a/docs/PROTOCOL.md +++ b/docs/PROTOCOL.md @@ -605,7 +605,16 @@ to the current versions. Per preset: - full-parameter, asynchronous rounds (`MilesFullParameterRoleStream`, `FullParameterSync`) overwrites a fragment with a broadcast that skips versions (catch-up) and releases its retained pushes of the skipped rounds. - A fresh full-parameter island cannot join a pool after its initial cut yet. At the end the island ACKs the final cut like a strict + A restarted island process (new connection generation, same island id) + joins a running pool: the syncer admits it (it supersedes a live older + generation, or re-JOINs after a lease expiry) and resends every fragment at + its current version. The island takes the newest version of each fragment + as its initial cut: the trainer weights are overwritten with it and it + becomes the anchor of the island's next push. Only island 0 sends INIT; a + restarted island 0 sends it again and the running syncer ignores it. The + restarted island starts with a fresh optimizer state (peer recovery, + `--rl-peer-recovery`, restores it where available). At the end the island + ACKs the final cut like a strict island; a push that races the end ends the island normally (`elastic_finished` on its tape). diff --git a/tests/test_peer_recovery_miles.py b/tests/test_peer_recovery_miles.py index 45071cb3b..5b8d2a0aa 100644 --- a/tests/test_peer_recovery_miles.py +++ b/tests/test_peer_recovery_miles.py @@ -87,6 +87,14 @@ async def full_parameter_optimizer_states(self): for s in self.states) +class FreshGroup(FakeGroup): + """Like Megatron in a fresh process: no optimizer-state proof before the + island's first local step, even after a snapshot restore (s27-gd-rec-a).""" + + async def full_parameter_optimizer_states(self): + raise RuntimeError("full-parameter last rollout ID is unavailable") + + class FakeAdapter: def __init__(self): self.layout = SimpleNamespace(layout_hash=LAYOUT_HASH, fragments=SimpleNamespace(num_fragments=P)) @@ -209,7 +217,8 @@ def donor_stream(store_root, group, client): return stream -def test_donor_snapshot_and_recovering_restore_through_the_role_stream(tmp_path, stubs): +@pytest.mark.parametrize("group_type", [FakeGroup, FreshGroup]) +def test_donor_snapshot_and_recovering_restore_through_the_role_stream(tmp_path, stubs, group_type): store_root = tmp_path / "store" donor_group = FakeGroup(adam_step=17) request = rp.SnapshotRequest(1, 9, 5, (3, 4)) @@ -226,7 +235,7 @@ def test_donor_snapshot_and_recovering_restore_through_the_role_stream(tmp_path, rp.RecoverySource(rp.SOURCE_PEER, 0, versions, sha256, locator), rp.RecoveryDone(5, 6, 1, (5, 6))], updates=resent) stubs.clients.append(client) - group = FakeGroup(adam_step=0) + group = group_type(adam_step=0) stream = MilesFullParameterRoleStream(config(store_root, recover=True)) asyncio.run(stream.initialize(group)) diff --git a/tests/test_rl_driver_profiles.py b/tests/test_rl_driver_profiles.py index 077b4352a..07404d4f2 100644 --- a/tests/test_rl_driver_profiles.py +++ b/tests/test_rl_driver_profiles.py @@ -341,3 +341,38 @@ def test_round_delay_applies_only_to_named_islands(tmp_path, monkeypatch): bad.write_text('{"round_delay_s": 1, "round_delay_islands": "1"}') with pytest.raises(ValueError): drv.load_fault_injection({drv.FAULT_INJECTION_ENV: str(bad)}) + + +def test_fault_injection_process_exit_keys(tmp_path): + from yeto.rl.engine.driver import FAULT_INJECTION_ENV, load_fault_injection + + path = tmp_path / "f.json" + path.write_text(json.dumps({"exit_at_round": 3, "exit_islands": [1, 1]})) + assert load_fault_injection({FAULT_INJECTION_ENV: str(path)}) == { + "exit_at_round": 3, "exit_islands": [1]} + for bad in ({"exit_at_round": 0, "exit_islands": [1]}, {"exit_at_round": 2}, + {"exit_islands": [1]}): + path.write_text(json.dumps(bad)) + with pytest.raises(ValueError): + load_fault_injection({FAULT_INJECTION_ENV: str(path)}) + + +@pytest.mark.parametrize(("learner", "incarnation", "exits"), + [(1, None, True), (1, "0", True), (1, "1", False), (0, None, False)]) +def test_fault_injection_exits_only_the_first_container(monkeypatch, learner, incarnation, exits): + from types import SimpleNamespace + + from yeto.rl.engine import driver as module + + exited = [] + monkeypatch.setattr("os._exit", exited.append) + if incarnation is None: + monkeypatch.delenv(module.MODAL_INCARNATION_ENV, raising=False) + else: + monkeypatch.setenv(module.MODAL_INCARNATION_ENV, incarnation) + events = [] + fake = SimpleNamespace(fault_injection={"exit_at_round": 2, "exit_islands": [1]}, + learner_id=learner, emit=lambda *a, **k: events.append(a)) + module.IslandDriver._inject_exit(fake, 1) + module.IslandDriver._inject_exit(fake, 2) + assert exited == ([137] if exits else []) diff --git a/tests/test_rl_full_parameter_restart_rejoin.py b/tests/test_rl_full_parameter_restart_rejoin.py new file mode 100644 index 000000000..a6c93c849 --- /dev/null +++ b/tests/test_rl_full_parameter_restart_rejoin.py @@ -0,0 +1,118 @@ +"""S27 E-4c: a killed full-parameter island process restarts and re-joins a +running asynchronous pool. Real Rust syncer, no Ray. + +The restarted process dials with a new connection generation. The syncer +admits it (superseding or re-joining under the membership rules) and resends +every fragment at its current version; the island takes those values as its +anchors, keeps training, and acknowledges the final cut. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_full_parameter_async import ( # noqa: F401 + P, + NUMEL, + _client, + _delta, + _f32, + _hmac_key, + _initial, + _start_async_syncer, + _wait, + needs_cargo, + syncer_binary, +) + +LEASE = ("--rl-sync-quorum", "1", "--rl-sync-island-lease-s", "3") + + +def _run_restart(binary, tmp_path, *, total_steps: int, kill_after: int): + """Learner 1 dies once fragment versions reach ``kill_after``; a new + process for learner 1 joins, catches up and both islands finish.""" + + process, port, profile_hash = _start_async_syncer(binary, tmp_path, *LEASE, + total_steps=total_steps) + clients = {0: _client(0, port, profile_hash), 1: _client(1, port, profile_hash)} + versions = {0: [0] * P, 1: [0] * P} + seen = {0: set(), 1: set()} + restarted = {"done": False, "catch_up": None, "pushed_after": 0} + + def drain(learner): + client = clients[learner] + client.check_health() + for update in client.drain_updates(): + fid = update.fragment_id + versions[learner][fid] = max(versions[learner][fid], update.version) + for permit in client.drain_pulls(): + key = (permit.global_step, permit.round_attempt) + if key in seen[learner]: + continue + seen[learner].add(key) + base = versions[learner][permit.fragment_id] + client.push_fragment(permit.fragment_id, permit.global_step, permit.round_attempt, + base, permit.global_step, 1, 7, + _f32(_delta(learner, permit.global_step))) + if learner == 1 and restarted["done"]: + restarted["pushed_after"] += 1 + + def restart(): + clients[1].close() # the process dies: its sockets close, no LEAVE + fresh = _client(1, port, profile_hash) + fresh.start() + caught: dict[int, int] = {} + + def catch_up(): + fresh.check_health() + for update in fresh.drain_updates(): + caught[update.fragment_id] = max(caught.get(update.fragment_id, 0), update.version) + return len(caught) == P + + _wait(catch_up, "restarted island got no catch-up", timeout=30) + clients[1] = fresh + versions[1] = [caught[f] for f in range(P)] + seen[1].clear() + restarted.update(done=True, catch_up=dict(caught)) + + def poll(): + drain(0) + if not restarted["done"] and min(versions[0]) >= kill_after: + restart() + drain(1) + return all(client.finalizing.is_set() for client in clients.values()) + + try: + for client in clients.values(): + client.start() + for fid in range(P): + clients[0].send_init(fid, _f32(_initial(fid))) + _wait(poll, "syncer never finalized", timeout=90) + finals = [] + for client in clients.values(): + manifest, fragments = client.wait_for_final_fragments(timeout=30) + finals.append({f.fragment_id: np.frombuffer(bytes(f.data), dtype=np.float32)[-NUMEL:] + for f in fragments}) + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + finally: + for client in clients.values(): + client.close() + if process.poll() is None: + process.terminate() + log, _ = process.communicate(timeout=10) + for fid in range(P): + assert np.array_equal(finals[0][fid], finals[1][fid]) + return restarted, log, (tmp_path / "tape.jsonl").read_text() + + +@needs_cargo +@pytest.mark.usefixtures("_hmac_key") +def test_a_restarted_island_catches_up_and_finishes(syncer_binary, tmp_path): + restarted, log, tape = _run_restart(syncer_binary, tmp_path, total_steps=6, kill_after=2) + assert restarted["done"] + assert min(restarted["catch_up"].values()) >= 2 + assert restarted["pushed_after"] > 0, "the restarted island never contributed" + assert '"catch_up":true' in tape diff --git a/tests/test_rl_miles_restart_rejoin.py b/tests/test_rl_miles_restart_rejoin.py new file mode 100644 index 000000000..5985d0416 --- /dev/null +++ b/tests/test_rl_miles_restart_rejoin.py @@ -0,0 +1,264 @@ +"""S27 E-4c: a restarted full-parameter island joins a running pool through +the real Miles role stream. Real Rust syncer (asynchronous rounds, island +lease); the Megatron trainer group and Ray are CPU fakes. + +Learner 1 trains a few rounds as a raw client, then its process dies (the +sockets close, no LEAVE). A new process for learner 1 starts the real +``MilesFullParameterRoleStream`` from the checkpoint weights. It must not +send INIT, must take the syncer's catch-up (every fragment at its current +version) as its trainer weights and anchors, push deltas measured from those +anchors that the syncer accepts, and acknowledge the final cut. + +Needs Miles (``importorskip``) and cargo. +""" + +from __future__ import annotations + +# ruff: noqa: I001 -- the Miles test helpers install a bridge stub before import. + +import asyncio +import json +import threading +import time +from dataclasses import replace + +import numpy as np +import pytest +import torch + +pytest.importorskip("miles") + +from test_e3_overlap_miles_role_stream import SIZES, _Driver, _Trainer, _adapter # noqa: E402 +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: E402,F401 +from test_rl_full_parameter_async import ( # noqa: E402,F401 + P, + _hmac_key, + _start_async_syncer, + _wait, + needs_cargo, + syncer_binary, +) +from test_rl_full_parameter_sync import _role_stream_config # noqa: E402 +from test_rl_miles_chunked_full_parameter import FakeRay, _evidence, component # noqa: E402 + +from yeto.protocol import DTYPE_F32, SyncerClient # noqa: E402 +from yeto.rl.adapters.miles.full_parameter import role_stream as role_stream_module # noqa: E402 +from yeto.rl.adapters.miles.full_parameter.role_stream import ( # noqa: E402 + MilesFullParameterRoleStream, + role_stream_session_contract_hash, +) +from yeto.rl.adapters.miles.models.chunked_full_parameter import ( # noqa: E402 + MilesChunkedFullParameterAdapter, +) +from yeto.rl.engine.full_parameter_sync import ( # noqa: E402 + FullParameterSync, + FullParameterSyncProfile, + RolePairing, +) + +LEASE_S = 3.0 +POOL = ("--rl-sync-quorum", "1", "--rl-sync-island-lease-s", str(LEASE_S)) +ROUNDS = 8 # outer versions per fragment: total steps + + +def _initial() -> dict[int, torch.Tensor]: + return {f: torch.arange(n, dtype=torch.float32) + float(f) for f, n in enumerate(SIZES)} + + +class _FreshTrainer(_Trainer): + """Like Megatron: no optimizer-state proof before the first local step + (s27-rj-a: "full-parameter last rollout ID is unavailable").""" + + async def full_parameter_optimizer_states(self): + if not self.trained_from: + raise RuntimeError("full-parameter last rollout ID is unavailable") + return await super().full_parameter_optimizer_states() + + +def _raw_delta(learner: int, step: int, fid: int) -> np.ndarray: + return np.full(SIZES[fid], 0.01 * (learner + 1) + 0.001 * step, dtype=np.float32) + + +class _RawIsland(threading.Thread): + """A raw-client island: answers every PULL against its newest version.""" + + def __init__(self, client: SyncerClient, learner: int) -> None: + super().__init__(daemon=True) + self.client, self.learner = client, learner + self.versions = [0] * P + self.seen: set = set() + self.halt = threading.Event() + self.error: BaseException | None = None + + def step(self) -> None: + self.client.check_health() + for update in self.client.drain_updates(): + self.versions[update.fragment_id] = max(self.versions[update.fragment_id], update.version) + for permit in self.client.drain_pulls(): + key = (permit.global_step, permit.round_attempt) + if key in self.seen: + continue + self.seen.add(key) + self.client.push_fragment( + permit.fragment_id, permit.global_step, permit.round_attempt, + self.versions[permit.fragment_id], permit.global_step, 1, 7, + _raw_delta(self.learner, permit.global_step, permit.fragment_id).tobytes()) + + def run(self) -> None: + try: + while not self.halt.is_set() and not self.client.finalizing.is_set(): + self.step() + self.halt.wait(0.005) + except BaseException as error: # noqa: BLE001 - reported by the test + self.error = error + + +def _stream_config(port: int, profile_hash: str, adapter): + return replace( + _role_stream_config("actor", port), + component=component(), learner_id=1, learner_generation=0, learner_generations=(0, 0), + total_fragment_steps=ROUNDS * P, expected_fragments=P, + expected_layout_hash=adapter.layout.layout_hash, syncer_profile_hash=profile_hash, + pipeline_depth=2, wait_timeout=60.0, poll_seconds=0.005, + max_fragment_bytes=64, max_chunk_bytes=8, partial_quorum=True, island_lease_s=LEASE_S, + ) + + +def _raw_client(learner, port, profile_hash, adapter, contract) -> SyncerClient: + return SyncerClient(("127.0.0.1", port), learner, adapter.layout.fragments, dtype=DTYPE_F32, + num_streams=1, connect_timeout=10, max_reconnects=0, finalization_timeout=60, + session_contract_hash=contract, syncer_profile_hash=bytes.fromhex(profile_hash)) + + +@needs_cargo +@pytest.mark.usefixtures("_hmac_key") +@pytest.mark.parametrize("death", ["closed", "silent"]) +def test_restarted_miles_island_catches_up_trains_and_acknowledges(syncer_binary, tmp_path, monkeypatch, + death): + """``closed``: the old sockets close (the syncer sees the departure). + ``silent``: the old connection stays open and mute (a lost host); the new + generation supersedes it before its lease runs out.""" + process, port, profile_hash = _start_async_syncer(syncer_binary, tmp_path, *POOL, + total_steps=ROUNDS) + adapter = _adapter() + config = _stream_config(port, profile_hash, adapter) + contract = role_stream_session_contract_hash(adapter.layout.layout_hash, config) + initial = _initial() + islands = {learner: _RawIsland(_raw_client(learner, port, profile_hash, adapter, contract), learner) + for learner in (0, 1)} + stream = dead = None + try: + for island in islands.values(): + island.client.start() + for fid in range(P): + islands[0].client.send_init(fid, initial[fid].numpy().tobytes()) + for island in islands.values(): + island.start() + _wait(lambda: min(islands[0].versions) >= 2, "the pool never reached version 2", timeout=30) + # Learner 1's process dies without a LEAVE. + islands[1].halt.set() + islands[1].join(timeout=5) + dead = islands[1].client + if death == "closed": + dead.close() + died_at = list(islands[0].versions) + + # The new learner-1 process starts the Miles role stream from the + # checkpoint weights (the run's initial weights). + ray = FakeRay() + trainer = _FreshTrainer(adapter, ray, initial, lambda _k: {f: torch.full((n,), 0.05) + for f, n in enumerate(SIZES)}) + + async def initialize(cls, group, *, policy_version, **_kwargs): + cut = await group.export_full_parameter_chunked_cut(policy_version, 0, + max_chunk_bytes=adapter.max_chunk_bytes) + return adapter, adapter.validate_cut(cut, policy_version, 0) + + monkeypatch.setattr(MilesChunkedFullParameterAdapter, "initialize", classmethod(initialize)) + sent_init = [] + real_client = role_stream_module.SyncerClient + + def client_factory(*args, **kwargs): + client = real_client(*args, **kwargs) + original = client.send_initial_fragments + client.send_initial_fragments = lambda parts: sent_init.append(1) or original(parts) + return client + + monkeypatch.setattr(role_stream_module, "SyncerClient", client_factory) + stream = MilesFullParameterRoleStream(config, ray_module=ray) + profile = FullParameterSyncProfile.from_configs({"actor": config}) + session = FullParameterSync( + RolePairing({"actor": stream}, profile), groups={"actor": trainer}, run=asyncio.run, + trajectories=lambda rollout_id: _evidence(rollout_id=rollout_id, + behavior_policy_hash=stream.current.policy_hash)) + driver = _Driver() + session.start(driver) + + caught = stream.fragment_versions + assert not sent_init, "a non-seed island must not send INIT" + assert all(v >= d for v, d in zip(caught, died_at, strict=True)) and min(caught) >= 2 + for fid in range(P): + # The trainer holds the pool's cut, not the checkpoint weights. + assert not torch.equal(trainer.values[fid], initial[fid]) + assert stream._anchors[fid].version == caught[fid] + + stop = False + deadline = time.monotonic() + 60 + rollout_id = -1 + while time.monotonic() < deadline: + rollout_id += 1 + boundary = session.boundary(driver, rollout_id=rollout_id, stats=None) + if boundary.stop: + stop = True + session.published(driver, rollout_id=rollout_id, + policy_hash=boundary.state.policy_tensor_hash()) + break + assert stop, "the restarted island never reached the final cut" + session.finish(driver) + assert stream.final_applied and stream.final_acknowledged + islands[0].join(timeout=30) + assert islands[0].error is None, islands[0].error + manifest, finals = islands[0].client.wait_for_final_fragments(timeout=30) + islands[0].client.acknowledge_finalization(manifest) + for fragment in finals: + final = np.frombuffer(bytes(fragment.data), dtype=np.float32)[-SIZES[fragment.fragment_id]:] + np.testing.assert_array_equal(trainer.values[fragment.fragment_id].numpy(), final) + assert process.wait(timeout=20) == 0 + finally: + for island in islands.values(): + island.halt.set() + island.client.close() + if stream is not None: + stream.close() + if dead is not None: + dead.close() + if process.poll() is None: + process.terminate() + log, _ = process.communicate(timeout=10) + records = [json.loads(line) for line in (tmp_path / "tape.jsonl").read_text().splitlines() if line] + island_1 = [r for r in records if r.get("island_id") == 1] + rejoin = next(i for i, r in enumerate(island_1) if r.get("kind") == "pool_join" and r.get("catch_up")) + after = island_1[rejoin + 1:] + assert any(r["kind"] == "delta_accepted" for r in after), "the restarted island never contributed" + # Quorum 1 closes some rounds before the island answers (out_of_round); + # nothing is refused for its base after the catch-up. + assert {r.get("reason") for r in after if r["kind"] == "delta_rejected"} <= {"out_of_round"}, after + assert "predates the re-JOIN catch-up" not in log + + +def test_initial_wait_stops_when_the_pool_finished_first(monkeypatch): + """s27-rj-a: a restarted island that dials after the pool's final cut + must not wait for broadcasts until the 900 s timeout.""" + from test_rl_full_parameter_sync import _role_stream_config as config_of + from types import SimpleNamespace + + from yeto.rl.adapters.miles.full_parameter.role_stream import PoolFinishedBeforeJoin + + finalizing = threading.Event() + finalizing.set() + client = SimpleNamespace(check_health=lambda: None, drain_updates=lambda: [], + finalizing=finalizing, send_progress=lambda: None) + config = replace(config_of("actor", 30100), partial_quorum=True, island_lease_s=60.0) + adapter = SimpleNamespace(layout=SimpleNamespace(fragments=SimpleNamespace(num_fragments=2))) + with pytest.raises(PoolFinishedBeforeJoin): + MilesFullParameterRoleStream(config)._wait_for_initial_fragments(client, adapter) diff --git a/tests/test_spot_droppable.py b/tests/test_spot_droppable.py index 1fe276cac..05913cb30 100644 --- a/tests/test_spot_droppable.py +++ b/tests/test_spot_droppable.py @@ -41,9 +41,12 @@ def test_fixed_roster_refused_with_spec_reason(): dr.admit(_args(rl_sync_island_lease_s=None), ["nebius", "aws", "aws"]) -def test_full_parameter_islands_cannot_be_droppable(): - with pytest.raises(ValueError, match="full-parameter island cannot join a running pool"): +def test_full_parameter_islands_are_droppable_only_with_async_rounds(): + with pytest.raises(ValueError, match="joins a running pool only"): dr.admit(_args(rl_sync_preset="full-parameter"), ["nebius", "aws", "aws"]) + roles = dr.admit(_args(rl_sync_preset="full-parameter", rl_sync_full_parameter_rounds="async"), + ["nebius", "aws", "aws"]) + assert roles == {0: "anchor", 1: "droppable", 2: "droppable"} def test_needs_an_anchor(): diff --git a/yeto/cloud/droppable.py b/yeto/cloud/droppable.py index 4415334cd..600d9adc1 100644 --- a/yeto/cloud/droppable.py +++ b/yeto/cloud/droppable.py @@ -73,10 +73,12 @@ def admit(args, clouds: list[str]) -> dict[int, str]: if not island_pool(args): raise ValueError(LEGACY_REFUSED) - if getattr(args, "rl_sync_preset", None) == "full-parameter": - # A droppable island comes back as a fresh process after a reclaim. - raise ValueError("--rl-island-role: a full-parameter island cannot join a running pool " - "after the initial cut yet, so it cannot be droppable") + if (getattr(args, "rl_sync_preset", None) == "full-parameter" + and getattr(args, "rl_sync_full_parameter_rounds", "sweep") != "async"): + # A droppable island comes back as a fresh process after a reclaim; + # only asynchronous rounds let a fresh process catch up and join. + raise ValueError("--rl-island-role: a full-parameter island joins a running pool only " + "with --rl-sync-full-parameter-rounds async, so it cannot be droppable") roles = parse_roles(items, len(clouds)) anchors = [i for i, r in roles.items() if r == ANCHOR] if not anchors: diff --git a/yeto/modal_runner.py b/yeto/modal_runner.py index 761b53c9a..3169616bc 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -449,7 +449,7 @@ def __init__(self, src: str, dst: str, commit, interval_s: float = TAPE_SYNC_INT # names its copies ".inc" (k = 0 keeps the old names). self.incarnation: int | None = None - def _claim_incarnation(self) -> int: + def claim_incarnation(self) -> int: if self.incarnation is None: self.dst.mkdir(parents=True, exist_ok=True) taken = sorted(self.dst.glob("INCARNATION-*.txt")) @@ -462,7 +462,7 @@ def _claim_incarnation(self) -> int: return self.incarnation def _target_name(self, name: str) -> str: - k = self._claim_incarnation() + k = self.claim_incarnation() if k == 0: return name for suffix in TAPE_SUFFIXES: @@ -472,7 +472,7 @@ def _target_name(self, name: str) -> str: def sync_once(self) -> int: self.dst.mkdir(parents=True, exist_ok=True) - self._claim_incarnation() + self.claim_incarnation() for name, text in list(self.extra.items()): (self.dst / self._target_name(name)).write_text(text, encoding="utf-8") del self.extra[name] @@ -872,6 +872,9 @@ def island_main(cfg_json: str) -> int: TAPE_SOURCE_DIR, f"{TAPE_MOUNT}/{cfg.tape_subdir}/rank{rank}", _volume_commit(cfg.tape_volume_name), mirror_globs=TAPE_MIRROR_GLOBS, ) + # 0 for the island's first container, k for its k-th re-run (test-only + # fault injection exits only in the first one). + env["YETO_MODAL_INCARNATION"] = str(tape.claim_incarnation()) hostmem = None interval = hostmem_interval_s(cfg) if interval is not None: diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 2340513d4..80e088b76 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -384,6 +384,14 @@ def role_stream_session_contract_hash( return hashlib.sha256(b"yeto-role-stream-session-v3\0" + encoded).digest() +class PoolFinishedBeforeJoin(RuntimeError): + """A (re)started island reached the pool after its final cut.""" + + +def _is_set(event) -> bool: + return bool(event is not None and callable(getattr(event, "is_set", None)) and event.is_set()) + + class MilesFullParameterRoleStream: """Reference-backed full-parameter transport for one trainable role.""" @@ -487,6 +495,13 @@ def _check_holds(self, state: TrainerResidentState) -> None: if self.current is None or self.current.policy_hash != state.policy_tensor_hash(): raise WeightsChanged(f"{self.config.role} trainer no longer holds the published cut") + @property + def may_rejoin(self) -> bool: + """An island pool with asynchronous rounds: a (re)started island may + join a running pool, whose initial cut is the current one.""" + + return self.config.island_lease_s is not None and self.config.partial_quorum + @property def finalizing(self) -> bool: event = getattr(self.client, "finalizing", None) @@ -566,7 +581,7 @@ def init_parts(fragment_id: int): authoritative = await asyncio.to_thread( self._wait_for_initial_fragments, client, adapter, init_parts ) - adapter.verify_initial_fragments(initial, authoritative) + adapter.verify_initial_fragments(initial, authoritative, allow_newer=self.may_rejoin) except BaseException: if heartbeat is not None: heartbeat.stop() @@ -586,6 +601,22 @@ def init_parts(fragment_id: int): self._rounds_at_anchor = [0] * count self._tokens_at_anchor = [0] * count self._routed_at_anchor = [RoutedTokens()] * count + newer = {stored.fragment_id: stored for stored in authoritative if stored.version > 0} + if newer: + # A restarted island joined a running pool: the syncer resent + # every fragment at its current version. Those values replace the + # checkpoint weights and become this island's anchors, so its + # first push is measured from the pool's current cut. + # Before the first local step Megatron has no optimizer-state + # proof (no rollout trained yet), so none is compared. + await self._apply_authoritative(newer, optimizer_proof=False) + print( + f"[yeto-initial-cut] learner {self.config.learner_id} caught up to a running " + f"pool: {len(newer)}/{count} fragment(s) newer than version 0, " + f"versions {list(self.fragment_versions)}", + file=sys.stderr, + flush=True, + ) async def _install_routed_token_counter(self, group, adapter) -> None: if check_expert_partition(adapter.layout.specs) is None: @@ -667,7 +698,7 @@ async def _initialize_recovering(self, group, adapter, initial, client, sink) -> self._tokens_at_anchor = [0] * count self._routed_at_anchor = [RoutedTokens()] * count resent, island.resent = dict(island.resent), {} - await self._apply_authoritative(resent, recovered=True) + await self._apply_authoritative(resent, optimizer_proof=False) except BaseException: island.release() client.close() @@ -714,6 +745,14 @@ def _wait_for_initial_fragments( self._report_initial_cut(client, observed, count, started, "waiting") while len(observed) < count: client.check_health() + if self.may_rejoin and _is_set(getattr(client, "finalizing", None)): + # The pool reached its final cut before this island caught up; + # waiting for broadcasts would only run into the timeout. + self._report_initial_cut(client, observed, count, started, "pool finished") + raise PoolFinishedBeforeJoin( + f"{self.config.algorithm} {self.config.role}: the island pool reached " + "its final cut before this island caught up" + ) if init_parts is not None and client.init_replay_due(): # S26 E3b (s26-e3b-p06-C32a): the group carrying INIT died # before the syncer applied it; send it again on this group. @@ -721,13 +760,25 @@ def _wait_for_initial_fragments( client.send_initial_fragments(init_parts) for update in client.drain_updates(): stored = self._stored_update(update) - if update.version != 0 or stored.fragment_id in observed: - if self.config.island_lease_s is not None and update.version != 0: - raise RuntimeError( - f"{self.config.algorithm} {self.config.role}: the island pool is past " - "its initial cut; a full-parameter island cannot join a running pool " - "yet (it can re-JOIN after a lease expiry while its process runs)" - ) + previous = observed.get(stored.fragment_id) + if self.may_rejoin: + # A running pool sends its catch-up (every fragment at its + # current version) and may broadcast newer rounds before + # the catch-up completes: keep the newest version. + if previous is not None and stored.version <= previous.version: + if ( + stored.version == previous.version + and stored.wire_payload_hash != previous.wire_payload_hash + ): + raise RuntimeError( + f"{self.config.algorithm} {self.config.role} catch-up fragment " + f"{stored.fragment_id} version {stored.version} changed bytes" + ) + stored.release() + continue + if previous is not None: + previous.release() + elif update.version != 0 or previous is not None: raise RuntimeError( f"{self.config.algorithm} {self.config.role} initial broadcast changed" ) @@ -1001,12 +1052,16 @@ async def _apply_authoritative( *, anchors: Mapping[int, StoredAuthoritativeFragment] | None = None, counters: tuple[int, int, RoutedTokens] | None = None, - recovered: bool = False, + optimizer_proof: bool = True, ) -> None: """Apply ``selected`` fragments; ``anchors``/``counters`` override the recorded anchor and its round/token counters (overlap rebase). - ``recovered``: the first apply after a peer snapshot restore, whose - optimizer records carry the donor's policy version and generation.""" + + ``optimizer_proof=False``: the trainer has not trained yet (an initial + catch-up, or the first apply after a peer snapshot restore), so Megatron + has no optimizer-state proof to compare: it records the last rollout ID + only at a local step. A restore is checked against the donor's snapshot + digest instead.""" if self.adapter is None or self.current is None or self.group is None: raise RuntimeError(f"{self.config.algorithm} {self.config.role} stream is unavailable") target_version = self.transport_version + 1 @@ -1015,21 +1070,23 @@ async def _apply_authoritative( target_policy_version=target_version, authoritative_fragments=tuple(selected.values()), ) - before = tuple(await self.group.full_parameter_optimizer_states()) + before = (tuple(await self.group.full_parameter_optimizer_states()) + if optimizer_proof else None) await self.adapter.apply( self.group, target, commit_token=self._apply_token(target_version, selected), ) - after = tuple(await self.group.full_parameter_optimizer_states()) - validate_optimizer_apply_transition( - before, - after, - role=self.config.role, - base_policy_version=None if recovered else self.transport_version, - base_local_generation=None if recovered else self.local_generation, - target_policy_version=target_version, - ) + if before is not None: + after = tuple(await self.group.full_parameter_optimizer_states()) + validate_optimizer_apply_transition( + before, + after, + role=self.config.role, + base_policy_version=self.transport_version, + base_local_generation=self.local_generation, + target_policy_version=target_version, + ) previous_cut = self.current self.current = target self.transport_version = target_version diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index 15881cce1..ad4c8e273 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -684,8 +684,14 @@ def verify_initial_fragments( self, reference: ReferencedPolicyCut, fragments: Sequence[StoredAuthoritativeFragment], + *, + allow_newer: bool = False, ) -> None: - """Require the syncer's authoritative v0 bytes to equal the checkpoint.""" + """Require the syncer's authoritative v0 bytes to equal the checkpoint. + + ``allow_newer``: an island joining a running pool receives some + fragments at newer versions; those must match the layout, topology and + plan, but their content is the pool's, not the checkpoint's.""" expected = { descriptor.fragment_id: (shard, descriptor) @@ -698,11 +704,11 @@ def verify_initial_fragments( if ( row is None or stored.fragment_id in observed - or stored.version != 0 + or (stored.version != 0 and not allow_newer) or stored.parameter_layout_hash != self.layout.layout_hash or stored.topology != row[0].topology or stored.plan_hash != row[0].plan_hash - or stored.descriptor != row[1] + or (stored.version == 0 and stored.descriptor != row[1]) ): raise RuntimeError("dense syncer version-zero fragment changed") observed[stored.fragment_id] = stored diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index 2dff41301..8ef75437d 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -61,6 +61,7 @@ import json from contextlib import contextmanager import math +import sys import time from collections.abc import Callable, Mapping from dataclasses import dataclass @@ -132,6 +133,8 @@ FAULT_INJECTION_ENV = "YETO_RL_FAULT_INJECTION" FAULT_INJECTION_FILE = "yeto-rl-fault-injection.json" # at the repo root +# Set by the Modal runner: 0 for an island's first container, k for its k-th re-run. +MODAL_INCARNATION_ENV = "YETO_MODAL_INCARNATION" def load_fault_injection(environ: Mapping[str, str] | None = None) -> dict[str, Any]: @@ -142,6 +145,9 @@ def load_fault_injection(environ: Mapping[str, str] | None = None) -> dict[str, frozen code snapshot). Keys: ``publish_delay_s``; ``round_delay_s`` (sleep before each round's rollout, to emulate a slow island) with optional ``round_delay_islands`` (learner ids it applies to; default every island). + ``exit_at_round`` and ``exit_islands`` (the island process exits with + code 137 before that rollout, only in its first Modal container, so a + re-run of the island is a fresh process that rejoins). """ import os @@ -151,7 +157,8 @@ def load_fault_injection(environ: Mapping[str, str] | None = None) -> dict[str, if not candidate.is_file(): return {} raw = json.loads(candidate.read_text(encoding="utf-8")) - unknown = sorted(set(raw) - {"publish_delay_s", "round_delay_s", "round_delay_islands"}) + unknown = sorted(set(raw) - {"publish_delay_s", "round_delay_s", "round_delay_islands", + "exit_at_round", "exit_islands"}) if unknown: raise ValueError(f"unknown fault injection keys {unknown}") out: dict[str, Any] = {} @@ -167,6 +174,16 @@ def load_fault_injection(environ: Mapping[str, str] | None = None) -> dict[str, raise ValueError("round_delay_islands must be a list of learner ids") if "round_delay_s" in out: out["round_delay_islands"] = sorted(set(islands)) + if "exit_at_round" in raw: + round_id, islands = raw["exit_at_round"], raw.get("exit_islands") + if type(round_id) is not int or round_id < 1: + raise ValueError("exit_at_round must be a positive rollout id") + if not isinstance(islands, list) or not islands or not all(type(i) is int for i in islands): + raise ValueError("exit_at_round needs exit_islands, a list of learner ids") + out["exit_at_round"] = round_id + out["exit_islands"] = sorted(set(islands)) + elif "exit_islands" in raw: + raise ValueError("exit_islands needs exit_at_round") return out @@ -834,6 +851,19 @@ def _with_policy_labels(self, batch: RolloutBatchHandle) -> RolloutBatchHandle: ) return dataclasses.replace(batch, groups=groups) + def _inject_exit(self, rollout_id: int) -> None: + """Test-only: kill this island process (first Modal container only).""" + import os + + if (self.fault_injection.get("exit_at_round") != rollout_id + or self.learner_id not in self.fault_injection.get("exit_islands", ()) + or os.environ.get(MODAL_INCARNATION_ENV, "0") != "0"): + return + self.emit("rl_fault_injected", kind="process_exit", rollout_id=rollout_id) + sys.stdout.flush() + sys.stderr.flush() + os._exit(137) + def _refuse_if_recovery_required(self) -> None: """3.7 (recovery-design.md F13): an island that opened in RECOVERY_REQUIRED (restart recovery impossible or failed) never publishes, generates or @@ -1972,6 +2002,7 @@ def _run(self) -> TrainableState: if self.max_rollouts is not None and rollout_id >= self.max_rollouts: raise DriverError("run reached max_rollouts without a stop") self.safe_point(rollout_id) + self._inject_exit(rollout_id) boundary = self.run_round(rollout_id) state = boundary.state rollout_id += 1 From 75eac8c1a658e7af8a3ec90051a685c697936971 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:50:07 +0000 Subject: [PATCH 142/233] syncer: finalize as soon as the last round commits The scheduler loop committed the last round and then fell through to the event wait. Since the lease term joined that wait's deadline set, an empty pipeline slept until the next lease expiry. Idle learners send no traffic after the last pull, so their leases lapsed during that sleep and finalization failed with every learner absent. How it is wired: after the commit and launch phases, the loop re-checks its own exit condition (no rounds left to launch, nothing in flight, computing, ready or publishing) and continues to the loop head instead of waiting. The restart/rejoin full-parameter test covers this path; it failed on most runs before this change. Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 8e6c3f4db..e050bef0d 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -2940,6 +2940,17 @@ async fn scheduler( if submitted_any { continue; // refill the pipeline before waiting again } + // The commit above may have drained the last round. Re-check the loop + // condition now: waiting here would sleep until the next lease expiry, + // and the leases of idle learners then lapse before the final cut. + if next_t > last_step + && inflight.is_empty() + && computing.is_empty() + && ready.is_empty() + && publishing.is_none() + { + continue; + } // Wait for the next event, the earliest in-flight deadline, or the // launch throttle opening (whichever comes first). Without the // throttle term an empty pipeline would spin; without in-flight From 057a44a9695af58debee4a3a8e9528ac3d61966c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:11:46 +0000 Subject: [PATCH 143/233] modal: accept a sharded head's comma-separated syncer address list in the reachability check How it is wired: the runner's head reachability check splits the syncer address on commas and probes every shard address (a sharded head passes one address per shard); a single address is unchanged. Co-Authored-By: Claude Opus 5.5 --- tests/test_modal_runner.py | 10 ++++++++++ yeto/modal_runner.py | 5 +++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/test_modal_runner.py b/tests/test_modal_runner.py index 91a79b141..610f48835 100644 --- a/tests/test_modal_runner.py +++ b/tests/test_modal_runner.py @@ -635,3 +635,13 @@ def test_container_change_exit_code_is_registered_and_unique(): codes = [launcher.ISLAND_FAILED_EXIT, launcher.TEARDOWN_UNVERIFIED_EXIT, launcher.RUN_STALLED_EXIT, launcher.NO_SYNC_INCOMPLETE_EXIT] assert launcher.CONTAINER_CHANGED_EXIT == 7 and 7 not in codes + + +def test_resolve_syncer_for_modal_sharded_head_list(): + from yeto.modal_runner import resolve_syncer_for_modal + + assert resolve_syncer_for_modal("3.231.26.103:29400,3.231.26.103:29401", None) == \ + "3.231.26.103:29400,3.231.26.103:29401" + import pytest + with pytest.raises(ValueError): + resolve_syncer_for_modal("3.231.26.103:29400,10.0.0.1:29401", None) diff --git a/yeto/modal_runner.py b/yeto/modal_runner.py index 3169616bc..161b35fef 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -168,8 +168,9 @@ def resolve_syncer_for_modal(syncer_addr: str, public_override: str | None) -> s reachable from Modal and no `--syncer-public-addr` was given.""" if public_override: return public_override - host = syncer_addr.rsplit(":", 1)[0].strip("[]") - if is_public_address(host): + # a sharded head (--syncer-shards N) is a comma-separated list, one address per shard + hosts = [addr.rsplit(":", 1)[0].strip("[]") for addr in syncer_addr.split(",")] + if all(is_public_address(host) for host in hosts): return syncer_addr raise ValueError( f"syncer at {syncer_addr} is not reachable from Modal (private or unresolvable " From 2ab9b6f23ef4552c3c512cca3d49155df4c96416 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:20:51 +0000 Subject: [PATCH 144/233] miles config: full-parameter sync accepts optimizer CPU offload with a distopt-aware Miles How it is wired: --rl-optimizer-offload-fraction with the full-parameter sync keeps Megatron's distributed + precision-aware optimizer (OPTIMIZER_OFFLOAD_ARGV) when the pinned Miles reads DP=1 distributed-optimizer masters (pins.MILES_DISTOPT_FULL_PARAMETER_COMMIT, fork branch s27-distopt-fp). Until the image is re-pinned to that commit the config refuses, with a message naming the needed commit. Without offload the non-distributed optimizer stays (no change). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_miles_adapter_config.py | 21 ++++++++++-- yeto/rl/adapters/miles/config.py | 46 ++++++++++++++++++++------- yeto/rl/adapters/miles/pins.py | 6 ++++ 3 files changed, 60 insertions(+), 13 deletions(-) diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index f2073cb19..6408dfdb4 100644 --- a/tests/test_rl_miles_adapter_config.py +++ b/tests/test_rl_miles_adapter_config.py @@ -275,13 +275,30 @@ def test_no_sync_full_parameter_trainer_keeps_the_distributed_optimizer(): mc.validate_parsed_args(_parsed(use_distributed_optimizer=True), launch, num_cells=lambda a: 1) -def test_optimizer_offload_with_the_full_parameter_sync_is_refused(): +def test_optimizer_offload_with_the_full_parameter_sync_needs_the_distopt_reader(monkeypatch): + # S27 G-B: Miles reads DP=1 distributed-optimizer masters (incl. the + # HybridDeviceOptimizer's offloaded FP32 copies) from 56190a5a8 on; an + # older pinned image is refused before launch. cfg = sub(make_config(), "trainable", parameter_mode="full", target_modules=()) cfg = dataclasses.replace(cfg, train_memory=rc.TrainMemoryConfig(optimizer_offload_fraction=0.5)) - with pytest.raises(mc.MilesConfigError, match="needs the distributed optimizer"): + with pytest.raises(mc.MilesConfigError, match="56190a5a8 or later"): mc.translate_run_config(cfg, AlgorithmSpec()) +def test_optimizer_offload_with_the_full_parameter_sync_keeps_the_distributed_optimizer(monkeypatch): + cfg = sub(make_config(), "trainable", parameter_mode="full", target_modules=()) + cfg = dataclasses.replace(cfg, train_memory=rc.TrainMemoryConfig(optimizer_offload_fraction=1.0)) + from yeto.rl.adapters.miles import pins + + monkeypatch.setattr(pins, "MILES_FORK_LINEAGE", pins.MILES_FORK_LINEAGE + (mc.MILES_DISTOPT_FULL_PARAMETER_COMMIT,)) + monkeypatch.setattr(mc, "MILES_NEXT_COMMIT", mc.MILES_DISTOPT_FULL_PARAMETER_COMMIT) + launch = mc.translate_run_config(cfg, AlgorithmSpec()) + assert "use_distributed_optimizer" not in launch.runtime_attrs + assert ("--use-distributed-optimizer --use-precision-aware-optimizer --optimizer-cpu-offload " + "--overlap-cpu-optimizer-d2h-h2d --optimizer-offload-fraction 1.0") in " ".join(launch.argv) + mc.validate_parsed_args(_parsed(use_distributed_optimizer=True), launch, num_cells=lambda a: 1) + + def test_build_ports_launch_marks_a_no_sync_trainer_island(monkeypatch): from yeto.rl.adapters.miles import island_entry diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 1af82e71c..eb787d272 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -30,6 +30,11 @@ from dataclasses import dataclass, field from typing import Any +from yeto.rl.adapters.miles.pins import ( + MILES_DISTOPT_FULL_PARAMETER_COMMIT, + MILES_NEXT_COMMIT, + miles_fork_lineage_from, +) from yeto.rl.engine.algorithm import ( AlgorithmSpecError, BOUNDED_NONZERO_STD_FILTER, @@ -1242,16 +1247,25 @@ def sample_buffer_argv(role: str, *, max_policy_age: int) -> tuple[str, ...]: def _full_parameter_optimizer_attrs(trainable: Any, memory: Any, *, full_parameter_sync: bool) -> dict[str, Any]: - """Keep every FP32 master whole when the full-parameter sync reads them. + """Pick the optimizer layout the full-parameter sync reads its FP32 masters from. The ports full-parameter sync (presets ``full-parameter`` / - ``sao-streaming-full``) builds its shard manifest from whole masters (Miles - ``full_parameter_state._master_parameter`` refuses the flat owned-range - slice that ``param.main_param`` is under the distributed optimizer, even at - DP=1). Miles' ``set_default_megatron_args`` turns the distributed optimizer - on for Adam regardless of argv, so the attribute is applied in - ``preprocess_args``, after Miles' defaults and before Megatron's - ``validate_args``. + ``sao-streaming-full``) reads one whole FP32 master per parameter on each + rank (islands run DP=1). + + * No optimizer offload: the non-distributed optimizer, whose masters are + whole tensors. Miles' ``set_default_megatron_args`` turns the + distributed optimizer on for Adam regardless of argv, so the attribute + is applied in ``preprocess_args``, after Miles' defaults and before + Megatron's ``validate_args``. + * ``--rl-optimizer-offload-fraction`` > 0: Megatron's distributed + + precision-aware optimizer with a HybridDeviceOptimizer + (``OPTIMIZER_OFFLOAD_ARGV``). At DP=1 each parameter's FP32 master is + the hybrid optimizer's whole FP32 copy (host-pinned when offloaded); + Miles reads it and, on a cut apply, writes both that copy and the BF16 + model view (``full_parameter_state._propagate_master_writes``). Needs + a Miles image at or after :data:`MILES_DISTOPT_FULL_PARAMETER_COMMIT`; + DP>1 is still refused by Miles at the first manifest. A no-sync trainer island (samples from the store, L4.6) has no such reader: it keeps Miles' distributed optimizer, so DP replicas shard the optimizer @@ -1261,12 +1275,22 @@ def _full_parameter_optimizer_attrs(trainable: Any, memory: Any, *, full_paramet if trainable.parameter_mode != "full" or not full_parameter_sync: return {} if float(getattr(memory, "optimizer_offload_fraction", 0.0) or 0.0) > 0: - raise MilesConfigError( - "--rl-optimizer-offload-fraction needs the distributed optimizer; the full-parameter sync " - "keeps whole FP32 masters (use a no-sync trainer island)") + if not _miles_reads_distributed_optimizer_masters(): + raise MilesConfigError( + "--rl-optimizer-offload-fraction with the full-parameter sync needs Miles " + f"{MILES_DISTOPT_FULL_PARAMETER_COMMIT[:9]} or later (reads distributed-optimizer masters); " + f"the pinned image ships {MILES_NEXT_COMMIT[:9]} (or use a no-sync trainer island)") + return {} return {"use_distributed_optimizer": False} +def _miles_reads_distributed_optimizer_masters() -> bool: + from yeto.rl.adapters.miles import pins + + return (MILES_DISTOPT_FULL_PARAMETER_COMMIT in pins.MILES_FORK_LINEAGE + and MILES_NEXT_COMMIT in miles_fork_lineage_from(MILES_DISTOPT_FULL_PARAMETER_COMMIT)) + + def _resolve_tito_parsers(model: str) -> tuple[str | None, str | None]: from miles.utils.chat_template_utils import resolve_reasoning_and_tool_call_parser diff --git a/yeto/rl/adapters/miles/pins.py b/yeto/rl/adapters/miles/pins.py index 57173faa5..494be16b7 100644 --- a/yeto/rl/adapters/miles/pins.py +++ b/yeto/rl/adapters/miles/pins.py @@ -114,6 +114,12 @@ class MilesRevisionPins(NamedTuple): "ed5f64fd7a183cccc8ab7ca3da90fba570108c2b", ) +# First Miles commit whose full-parameter sync reads distributed-optimizer +# masters (and the HybridDeviceOptimizer's offloaded FP32 copies) at DP=1 +# (S27 G-B; local fork branch s27-distopt-fp on ed5f64fd7, no image yet). +# Append it to MILES_FORK_LINEAGE together with the image re-pin. +MILES_DISTOPT_FULL_PARAMETER_COMMIT = "56190a5a86d14f1eb96d592038d76a9ad8651b0c" + def miles_fork_lineage_from(commit: str) -> frozenset[str]: """``commit`` and every later commit of :data:`MILES_FORK_LINEAGE`.""" From 8e0fc69321492eb49fd77ec4b65d26c554d52097 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:53:45 +0000 Subject: [PATCH 145/233] syncer: carry late pushes into the next round in async mode; pace grace from every push How it is wired: - In asynchronous rounds (a positive --max-base-lag and no fixed roster) a push that arrives after its round closed is no longer dropped. It joins the next round of the same fragment (queued until that round launches), is merged against its own base with its own token weight and the carry discount, and is recorded in the merge record under "carried" (step it missed, base, lag). It takes no quorum slot. Pushes beyond the lag bound are still refused. A late push from an anchor its own round already merged is not carried, and after a carry the island may not answer again from the same or an older anchor of that fragment (anchor_reused), so no delta is counted twice. Strict runs (--max-base-lag 0) and unbounded runs keep the previous behavior. - Adaptive grace: StepRates times a local step from its first report (pushes and heartbeats repeat the same step count; re-stamping on each shrank the window to almost nothing), and every unfenced push feeds the estimate, including refused ones, so a peer whose pushes keep missing their round is waited for. The computed window and its inputs are logged when a round arms it. - wandb: carried late deltas are logged from the syncer tape. Tests: CPU reproduction of the GPU timing scaled 1:100 (the slower island now answers every round directly); carry, lag bound and anchor_reused cases; free-run participation checked loosely with no-loss as the exact check. Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 551 +++++++++++++++++++++++- tests/test_rl_late_carry_real_syncer.py | 359 +++++++++++++++ tests/test_wandb_tape.py | 3 +- yeto/wandb_tape.py | 7 + 4 files changed, 908 insertions(+), 12 deletions(-) create mode 100644 tests/test_rl_late_carry_real_syncer.py diff --git a/syncer/src/server.rs b/syncer/src/server.rs index e050bef0d..aaba34d72 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -786,16 +786,25 @@ const STEP_EMA_ALPHA: f64 = 0.5; impl StepRates { fn note(&mut self, member: Member, local_step: u64, now: Instant) { let entry = self.0.entry(member).or_insert((None, 0, None)); - if let Some(prev) = entry.0 { - if local_step > entry.1 { + // One local step produces a push per fragment (and heartbeats): only + // the first report of a new step marks its start. Re-stamping on + // every report of the same step would measure only the gap between + // the last report of one step and the first of the next. + match entry.0 { + Some(prev) if local_step > entry.1 => { let secs = now.duration_since(prev).as_secs_f64() / (local_step - entry.1) as f64; entry.2 = Some(match entry.2 { Some(ema) => STEP_EMA_ALPHA * secs + (1.0 - STEP_EMA_ALPHA) * ema, None => secs, // first sample seeds the EMA }); + *entry = (Some(now), local_step, entry.2); } + // A step counter that went backwards (a restarted learner) + // starts a new baseline; a repeat of the current step keeps it. + Some(_) if local_step < entry.1 => *entry = (Some(now), local_step, entry.2), + Some(_) => {} + None => *entry = (Some(now), local_step, None), } - *entry = (Some(now), local_step, entry.2); } /// Slowest estimated step time among one frozen/current member set. @@ -2645,6 +2654,7 @@ async fn scheduler( let mut computing = tokio::task::JoinSet::>::new(); let mut ready: BTreeMap = BTreeMap::new(); let mut busy_fragments: HashSet = HashSet::new(); + let mut late_carry = LateCarry::new(cfg.max_base_lag, fixed_roster); // --async-checkpoint: the committed round whose checkpoint is in flight. let mut publishing: Option = None; let mut cutoff_established = false; @@ -2809,20 +2819,40 @@ async fn scheduler( grace_ms: None, fold, pushes: HashMap::new(), + carried: Vec::new(), }); + if let Some(carry) = late_carry.as_mut() { + let round = inflight.last_mut().expect("round just launched"); + for (member, missed, outcome) in carry.attach_queued(round) { + log_carry(&cfg, dynamic, member, missed, t, outcome); + } + } } // Arm the grace window of any round that just reached quorum. for r in inflight.iter_mut() { if r.grace_deadline.is_none() && r.quorum_reached() { + let step_secs = step_rates.max_step_secs_for(&r.expected_members); + let quorum_secs = r.started.elapsed().as_secs_f64(); let grace = adaptive_grace( cfg.grace_tau, cfg.grace_gamma, - step_rates.max_step_secs_for(&r.expected_members), - r.started.elapsed().as_secs_f64(), + step_secs, + quorum_secs, last_sync_secs, Duration::from_millis(cfg.grace_ms), ); + if r.pushes.len() < r.expected_members.len() { + info!( + step = r.t, + fragment = r.p, + step_secs = ?step_secs, + quorum_secs, + sync_secs = last_sync_secs, + grace_ms = grace.as_millis() as u64, + "grace window armed" + ); + } // The soft deadline is an upper bound on the grace window too. let now = Instant::now(); let grace = r @@ -2920,6 +2950,11 @@ async fn scheduler( r.base_version = st.versions[r.p]; r.attempt += 1; r.pull = encode_pull(r.p, r.t, r.attempt); + if let Some(carry) = late_carry.as_mut() { + for (member, push) in &r.pushes { + carry.release_anchor(*member, push); + } + } r.pushes.clear(); r.fold = round_fold(&cfg, &st, r.p, &r.expected_members, r.quorum_size); r.started = Instant::now(); @@ -3097,7 +3132,59 @@ async fn scheduler( rejected("fenced"); continue; } + // Pace from every unfenced push, refused ones included: a + // peer whose pushes keep missing their round is exactly + // the one the grace window must be sized for. + step_rates.note(member, local_step, Instant::now()); let base_version = push.base_version; + if let Some(carry) = late_carry.as_mut() { + if carry.reuses_anchor(member, &push) { + warn!( + learner_id, + generation, + step = global_step, + fragment = fragment_id, + base = base_version, + "push reuses an anchor whose delta was already taken" + ); + rejected("anchor_reused"); + continue; + } + let open = inflight.iter().any(|round| { + round.t == global_step + && round.p == fragment_id as usize + && round.attempt == push.round_attempt + }); + if !open { + let current = st.versions[fragment_id as usize]; + let outcome = carry.carry(&mut inflight, member, push, next_t, current); + if !matches!(outcome, CarryOutcome::Rejected(_)) { + carry.spend_anchor(member, fragment_id as usize, base_version); + } + match outcome { + CarryOutcome::Attached { step, lag } => { + log_carry(&cfg, dynamic, member, global_step, step, Ok(lag)) + } + CarryOutcome::Queued => info!( + learner_id, + generation, + step = global_step, + fragment = fragment_id, + base = base_version, + "late push queued for the fragment's next round" + ), + CarryOutcome::Rejected(disposition) => log_carry( + &cfg, + dynamic, + member, + global_step, + global_step, + Err(disposition), + ), + } + continue; + } + } let round_base = inflight .iter() .find(|round| round.t == global_step && round.p == fragment_id as usize) @@ -3106,7 +3193,9 @@ async fn scheduler( route_push(&mut inflight, member, push, cfg.max_base_lag, fixed_roster); match disposition { PushDisposition::Accepted => { - step_rates.note(member, local_step, Instant::now()); + if let Some(carry) = late_carry.as_mut() { + carry.note_gathered(member, fragment_id as usize, base_version); + } let lag = round_base.map_or(0, |base| base.saturating_sub(base_version)); if lag == 0 { pool("delta_accepted", format!("\"island_id\":{learner_id}")); @@ -3812,6 +3901,10 @@ struct Round { /// keeps every push's gradient until the buffered merge. fold: Option, pushes: HashMap, + /// Late deltas of earlier rounds of this fragment (async rounds, see + /// `LateCarry`): merged here against their own base, in arrival order. + /// They never count toward quorum and stay buffered until the merge. + carried: Vec<(Member, Push)>, } impl Round { @@ -4086,6 +4179,213 @@ fn route_push( PushDisposition::Accepted } +/// Late-push carry for asynchronous rounds. +/// +/// A push that reaches the syncer after its round closed still holds the +/// island's work since its anchor (`local - raw anchor at base_version`). +/// Instead of dropping it, the syncer merges it into the next round of the +/// same fragment against its own base, weighted by its counters (and the +/// carry discount), as long as its base lag stays within `--max-base-lag`. +/// +/// A carried delta counts once: after the syncer carries a delta from an +/// anchor (fragment version), a later push of that island connection from +/// the same or an older anchor of that fragment is refused +/// (`anchor_reused`), and a late push from an anchor whose delta was +/// already merged is not carried. Pushes that never involve a carry keep +/// the existing stale-base rules. +struct LateCarry { + limit: u64, + /// Late pushes waiting for their fragment's next round. + queued: HashMap>, + /// Newest base version of a carried delta per island connection and + /// fragment. + anchors: HashMap<(Member, usize), u64>, + /// Newest base version of a delta gathered by its own round. + gathered: HashMap<(Member, usize), u64>, +} + +#[derive(Debug, Eq, PartialEq)] +enum CarryOutcome { + /// Added to the open round of this fragment launched at `step`. + Attached { step: u64, lag: u64 }, + /// Held until the fragment's next round launches. + Queued, + Rejected(PushDisposition), +} + +impl LateCarry { + /// Carry is on for asynchronous rounds only: a positive base-lag bound + /// and no fixed roster. Strict (`--max-base-lag 0`) and unbounded runs + /// keep refusing late pushes. + fn new(max_base_lag: Option, fixed_roster: bool) -> Option { + match max_base_lag { + Some(limit) if limit > 0 && !fixed_roster => Some(Self { + limit, + queued: HashMap::new(), + anchors: HashMap::new(), + gathered: HashMap::new(), + }), + _ => None, + } + } + + /// True when `push` reuses an anchor this connection already spent. + fn reuses_anchor(&self, member: Member, push: &Push) -> bool { + self.anchors + .get(&(member, push.fragment_id as usize)) + .is_some_and(|&spent| push.base_version <= spent) + } + + fn spend_anchor(&mut self, member: Member, fragment: usize, base_version: u64) { + let spent = self.anchors.entry((member, fragment)).or_insert(0); + *spent = (*spent).max(base_version); + } + + fn note_gathered(&mut self, member: Member, fragment: usize, base_version: u64) { + let taken = self.gathered.entry((member, fragment)).or_insert(0); + *taken = (*taken).max(base_version); + } + + /// A restarted round drops its gathered pushes; their islands answer the + /// new attempt from the same anchor, so the anchor is free again. + fn release_anchor(&mut self, member: Member, push: &Push) { + let key = (member, push.fragment_id as usize); + if self.gathered.get(&key) == Some(&push.base_version) { + if push.base_version == 0 { + self.gathered.remove(&key); + } else { + self.gathered.insert(key, push.base_version - 1); + } + } + } + + /// Route a push that `route_push` found no open round for. + /// `next_t` is the next step to launch; `current_version` the fragment's + /// committed version. + fn carry( + &mut self, + rounds: &mut [Round], + member: Member, + push: Push, + next_t: u64, + current_version: u64, + ) -> CarryOutcome { + let p = push.fragment_id as usize; + // A push for a round still open under another attempt, or for a + // round never launched, is not late. + if push.global_step >= next_t + || rounds.iter().any(|round| round.t == push.global_step && round.p == p) + { + return CarryOutcome::Rejected(PushDisposition::OutOfRound); + } + if push.base_version > current_version { + return CarryOutcome::Rejected(PushDisposition::FutureBase); + } + if self + .gathered + .get(&(member, p)) + .is_some_and(|&taken| push.base_version <= taken) + { + return CarryOutcome::Rejected(PushDisposition::Duplicate); + } + if current_version - push.base_version > self.limit { + return CarryOutcome::Rejected(PushDisposition::StaleBase); + } + if let Some(round) = rounds + .iter_mut() + .find(|round| round.p == p && round.t > push.global_step) + { + return match attach_carry(round, member, push, self.limit) { + Ok(lag) => CarryOutcome::Attached { step: round.t, lag }, + Err(disposition) => CarryOutcome::Rejected(disposition), + }; + } + self.queued.entry(p).or_default().push((member, push)); + CarryOutcome::Queued + } + + /// Move the late pushes queued for `round`'s fragment into it. + fn attach_queued(&mut self, round: &mut Round) -> Vec<(Member, u64, Result)> { + let Some(queued) = self.queued.remove(&round.p) else { + return Vec::new(); + }; + queued + .into_iter() + .map(|(member, push)| { + let missed = push.global_step; + (member, missed, attach_carry(round, member, push, self.limit)) + }) + .collect() + } +} + +/// Log (and with dynamic membership, tape) the outcome of a late push. +fn log_carry( + cfg: &Config, + dynamic: bool, + member: Member, + missed_step: u64, + into_step: u64, + outcome: Result, +) { + let island = member.learner_id; + match outcome { + Ok(lag) => { + info!(island, missed_step, into_step, lag, "late push carried into a later round"); + if dynamic { + append_pool_record( + cfg.event_tape.as_deref(), + "delta_carried_over", + cfg.syncer_epoch, + &format!( + "\"island_id\":{island},\"missed_step\":{missed_step},\"step\":{into_step},\"lag\":{lag},\"discount\":{:?}", + cfg.membership.carry_discount(lag) + ), + ); + } + } + Err(disposition) => { + warn!(island, missed_step, reason = ?disposition, "late push not carried"); + if dynamic { + let reason = match disposition { + PushDisposition::StaleBase => "carry_lag_exceeded", + PushDisposition::FutureBase => "future_base", + PushDisposition::Duplicate => "anchor_reused", + _ => "out_of_round", + }; + append_pool_record( + cfg.event_tape.as_deref(), + "delta_rejected", + cfg.syncer_epoch, + &format!("\"island_id\":{island},\"reason\":\"{reason}\""), + ); + } + } + } +} + +/// Add a late delta to an open round; returns its base lag. +fn attach_carry( + round: &mut Round, + member: Member, + push: Push, + limit: u64, +) -> Result { + // The exact fold reproduces a fixed responder set bit for bit. + if matches!(round.fold, Some(RoundFold::Exact(_))) { + return Err(PushDisposition::OutOfRound); + } + if push.base_version > round.base_version { + return Err(PushDisposition::FutureBase); + } + let lag = round.base_version - push.base_version; + if lag > limit { + return Err(PushDisposition::StaleBase); + } + round.carried.push((member, push)); + Ok(lag) +} + /// The streaming reduction for a new round attempt, or `None` to keep every /// push's gradient until the buffered merge. /// @@ -4181,6 +4481,9 @@ fn prepare_round_compute( } let mut ordered_pushes: Vec<_> = round.pushes.iter().collect(); ordered_pushes.sort_unstable_by_key(|(member, _)| **member); + let responders = ordered_pushes.len(); + // Carried late deltas follow this round's own pushes, in arrival order. + ordered_pushes.extend(round.carried.iter().map(|(member, push)| (member, push))); let mut outer_gradients = Vec::with_capacity(ordered_pushes.len()); let mut weights = Vec::with_capacity(ordered_pushes.len()); let mut tensor_scales = Vec::with_capacity(ordered_pushes.len()); @@ -4209,7 +4512,7 @@ fn prepare_round_compute( if !cfg.membership.is_static() { append_outer_step_record(cfg, round, &ordered_pushes, &weights); } - let Some(fold) = round.fold.take() else { + let Some(mut fold) = round.fold.take() else { let tensor_scales = cfg.expert_token_weights.then_some(tensor_scales.as_slice()); return st.prepare_merge_scaled( round.p, @@ -4222,7 +4525,8 @@ fn prepare_round_compute( // Every responder must have been folded exactly once. The exact fold // reduced in learner-id order, which equals the buffered merge's Member // order when ids are unique; the streaming fold in arrival order. - let responders: Vec = ordered_pushes.iter().map(|(m, _)| m.learner_id).collect(); + let responders: Vec = + ordered_pushes[..responders].iter().map(|(m, _)| m.learner_id).collect(); let mut folded = fold.folded_learners().to_vec(); folded.sort_unstable(); if folded != responders { @@ -4234,6 +4538,17 @@ fn prepare_round_compute( responders ); } + // Carried deltas stay buffered until now (only a streaming fold admits + // them, see `attach_carry`). + if !round.carried.is_empty() { + let RoundFold::Streaming(streaming) = &mut fold else { + bail!("round t={} fragment {}: carried deltas need a streaming fold", round.t, round.p); + }; + for (index, (member, push)) in round.carried.iter_mut().enumerate() { + let weight = weights[responders.len() + index]; + streaming.fold(member.learner_id, std::mem::take(&mut push.outer_gradient), weight)?; + } + } let delta = fold .finish() .with_context(|| format!("round t={} fragment {}", round.t, round.p))?; @@ -4253,6 +4568,7 @@ struct CommittedRound { quorum_ms: Option, grace_ms: Option, pushes: HashMap, + carried: Vec<(Member, Push)>, gnorm: f64, compute_secs: f64, sync_start: Instant, @@ -4305,6 +4621,7 @@ fn apply_round_commit( quorum_ms, grace_ms, pushes, + carried, .. } = round; let expected_t = st.schedule().step_after(st.global_step); @@ -4323,7 +4640,9 @@ fn apply_round_commit( st.global_step = t; let mut ordered_pushes: Vec<_> = pushes.iter().collect(); ordered_pushes.sort_unstable_by_key(|(member, _)| **member); - for (_, push) in ordered_pushes { + // A carried delta was never counted by the round it missed. + let merged = ordered_pushes.into_iter().chain(carried.iter().map(|(member, push)| (member, push))); + for (_, push) in merged { if let Some(fragments) = cfg.policy_sweep_fragments { st.record_fragment_merge( push.learner_id, @@ -4348,6 +4667,7 @@ fn apply_round_commit( quorum_ms, grace_ms, pushes, + carried, gnorm, compute_secs, sync_start, @@ -4377,6 +4697,7 @@ async fn publish_round( quorum_ms, grace_ms, pushes, + carried, gnorm, compute_secs, sync_start, @@ -4469,6 +4790,7 @@ async fn publish_round( sync_ms, checkpoint_ms, &pushes, + &carried, gnorm, ms, &st.layout_fingerprint, @@ -5541,6 +5863,7 @@ fn append_tape( sync_ms: u64, checkpoint_ms: Option, pushes: &HashMap, + carried: &[(Member, Push)], gnorm: f64, ms: u64, layout_fingerprint: &[u8; 32], @@ -5575,7 +5898,37 @@ fn append_tape( crate::merge::learner_weight(push.c_tokens, u64::from(push.c_steps)) } }; - let weight_sum: f64 = pushes.values().map(event_weight).sum(); + let weight_sum: f64 = pushes + .values() + .chain(carried.iter().map(|(_, push)| push)) + .map(event_weight) + .sum(); + // Late deltas of earlier rounds merged here (async rounds only; the field + // is absent when there are none). + let carried_field = if carried.is_empty() { + String::new() + } else { + let entries: Vec = carried + .iter() + .map(|(member, p)| { + let weight = event_weight(p); + let contribution = if weight_sum > 0.0 { weight / weight_sum } else { 0.0 }; + format!( + "{{\"id\":{},\"generation\":{},\"missed_step\":{},\"base_version\":{},\"lag\":{},\"c_steps\":{},\"c_tokens\":{},\"weight\":{},\"contribution\":{}}}", + p.learner_id, + member.generation, + p.global_step, + p.base_version, + launch_base_version.saturating_sub(p.base_version), + p.c_steps, + p.c_tokens, + weight, + contribution + ) + }) + .collect(); + format!(",\"carried\":[{}]", entries.join(",")) + }; let mut responders: Vec = pushes .iter() .map(|(member, p)| { @@ -5619,7 +5972,7 @@ fn append_tape( let checkpoint_field = checkpoint_tape_field(checkpoint_ms); let sweep_fields = policy_sweep_event_fields(step, policy_sweep_fragments); let line = format!( - "{{\"protocol_version\":{PROTOCOL_VERSION},\"delta_semantics\":\"local_minus_raw_anchor\",\"sync/layout_hash\":\"{layout_hash}\",\"sync/base_version\":{launch_base_version},\"sync/responders\":{},\"sync/quorum\":{quorum},\"sync/rejected_stale_updates\":0,\"sync/merge_seconds\":{merge_seconds},\"sync/global_delta_norm\":{gnorm},\"step\":{step},\"fragment\":{fragment}{sweep_fields},\"launch_base_version\":{launch_base_version},\"attempt\":{attempt},\"gnorm\":{gnorm},\"ms\":{ms},\"quorum\":{quorum},\"expected\":{},\"expected_members\":{},\"responded\":{},\"responded_members\":{},\"missed_grace\":{},\"missed_members\":{},\"quorum_ms\":{quorum_ms},\"grace_ms\":{grace_ms},\"sync_ms\":{sync_ms}{checkpoint_field},\"time_unix\":{:?},\"responders\":[{}]}}\n", + "{{\"protocol_version\":{PROTOCOL_VERSION},\"delta_semantics\":\"local_minus_raw_anchor\",\"sync/layout_hash\":\"{layout_hash}\",\"sync/base_version\":{launch_base_version},\"sync/responders\":{},\"sync/quorum\":{quorum},\"sync/rejected_stale_updates\":0,\"sync/merge_seconds\":{merge_seconds},\"sync/global_delta_norm\":{gnorm},\"step\":{step},\"fragment\":{fragment}{sweep_fields},\"launch_base_version\":{launch_base_version},\"attempt\":{attempt},\"gnorm\":{gnorm},\"ms\":{ms},\"quorum\":{quorum},\"expected\":{},\"expected_members\":{},\"responded\":{},\"responded_members\":{},\"missed_grace\":{},\"missed_members\":{},\"quorum_ms\":{quorum_ms},\"grace_ms\":{grace_ms},\"sync_ms\":{sync_ms}{checkpoint_field}{carried_field},\"time_unix\":{:?},\"responders\":[{}]}}\n", responded.len(), json_ids(&expected_learners), json_members(expected_members), @@ -6572,6 +6925,7 @@ mod tests { quorum_ms: Some(0), grace_ms: Some(0), fold: None, + carried: Vec::new(), pushes: HashMap::from([( member, Push { @@ -6654,6 +7008,7 @@ mod tests { quorum_ms: Some(0), grace_ms: Some(0), fold: None, + carried: Vec::new(), pushes: HashMap::from([( member, Push { @@ -6716,6 +7071,7 @@ mod tests { quorum_ms: Some(0), grace_ms: Some(0), fold: None, + carried: Vec::new(), pushes, } } @@ -7785,6 +8141,174 @@ mod tests { assert!((rates.max_step_secs_for(&[first]).unwrap() - 0.5).abs() < 1e-9); } + #[test] + fn late_carry_is_off_for_strict_unbounded_and_fixed_roster_runs() { + assert!(LateCarry::new(Some(0), false).is_none()); + assert!(LateCarry::new(None, false).is_none()); + assert!(LateCarry::new(Some(2), true).is_none()); + assert!(LateCarry::new(Some(2), false).is_some()); + } + + #[test] + fn late_push_queues_then_joins_the_next_round_of_its_fragment() { + let mut carry = LateCarry::new(Some(2), false).unwrap(); + let island = member(1, 10); + // Round t=7 of fragment 1 (base 5) closed without island 1; the + // fragment is now at version 7 and its next round is not open yet. + let mut late = test_push_for(1, 5); + late.global_step = 7; + let mut rounds: Vec = Vec::new(); + assert_eq!(carry.carry(&mut rounds, island, late, 8, 7), CarryOutcome::Queued); + let mut next = test_round(vec![member(0, 10), island]); + next.t = 9; + next.base_version = 7; + let attached = carry.attach_queued(&mut next); + assert_eq!(attached.len(), 1); + assert_eq!(attached[0].2, Ok(2)); + assert_eq!(next.carried.len(), 1); + assert!(next.pushes.is_empty(), "a carried delta takes no quorum slot"); + // The island still answers the new round from its new anchor. + let mut fresh = test_push_for(1, 7); + fresh.global_step = 9; + assert_eq!( + route_push(&mut [next], island, fresh, Some(2), false), + PushDisposition::Accepted + ); + } + + #[test] + fn late_push_joins_an_already_open_later_round() { + let mut carry = LateCarry::new(Some(2), false).unwrap(); + let mut open = test_round(vec![member(0, 10)]); + open.t = 9; + open.base_version = 7; + let mut late = test_push_for(1, 5); + late.global_step = 7; + let mut rounds = vec![open]; + assert_eq!( + carry.carry(&mut rounds, member(1, 10), late, 10, 7), + CarryOutcome::Attached { step: 9, lag: 2 } + ); + assert_eq!(rounds[0].carried.len(), 1); + } + + #[test] + fn late_push_beyond_the_lag_bound_or_for_an_open_attempt_is_refused() { + let mut carry = LateCarry::new(Some(1), false).unwrap(); + let mut late = test_push_for(1, 5); + late.global_step = 7; + assert_eq!( + carry.carry(&mut Vec::new(), member(1, 10), late, 8, 7), + CarryOutcome::Rejected(PushDisposition::StaleBase) + ); + // Same step and fragment still open under attempt 2: not late. + let mut retried = test_round(vec![member(1, 10)]); + retried.attempt = 2; + let mut stale_attempt = test_push_for(1, 5); + stale_attempt.global_step = 7; + assert_eq!( + carry.carry(&mut vec![retried], member(1, 10), stale_attempt, 8, 5), + CarryOutcome::Rejected(PushDisposition::OutOfRound) + ); + // A never-launched step is not late either. + let mut future = test_push_for(1, 5); + future.global_step = 9; + assert_eq!( + carry.carry(&mut Vec::new(), member(1, 10), future, 8, 7), + CarryOutcome::Rejected(PushDisposition::OutOfRound) + ); + assert!(carry.queued.is_empty()); + } + + #[test] + fn a_carried_anchor_yields_one_delta() { + let mut carry = LateCarry::new(Some(2), false).unwrap(); + let island = member(1, 10); + let push = test_push_for(1, 5); + assert!(!carry.reuses_anchor(island, &push)); + carry.spend_anchor(island, 1, 5); + assert!(carry.reuses_anchor(island, &push)); + assert!(carry.reuses_anchor(island, &test_push_for(1, 3))); + assert!(!carry.reuses_anchor(island, &test_push_for(1, 7))); + // Another connection or fragment is unaffected. + assert!(!carry.reuses_anchor(member(1, 11), &push)); + // A late push from an anchor already merged by its own round is not + // carried; a restart clears the gathered push and frees the anchor. + let other = member(2, 10); + carry.note_gathered(other, 1, 5); + let mut late = test_push_for(2, 5); + late.global_step = 7; + assert_eq!( + carry.carry(&mut Vec::new(), other, late, 8, 7), + CarryOutcome::Rejected(PushDisposition::Duplicate) + ); + carry.release_anchor(other, &test_push_for(2, 5)); + let mut late = test_push_for(2, 5); + late.global_step = 7; + assert_eq!(carry.carry(&mut Vec::new(), other, late, 8, 7), CarryOutcome::Queued); + } + + fn carried_merge(cfg: &Config) -> f32 { + let mut st = two_fragment_streaming_state(); + st.versions[0] = 4; + let island = member(0, 1); + let late = member(1, 1); + let mut round = test_round(vec![island, late]); + round.p = 0; + round.t = 6; + round.base_version = 4; + round.fold = round_fold(cfg, &st, 0, &round.expected_members, 1); + let mut own = test_push_for(0, 4); + own.fragment_id = 0; + own.global_step = 6; + own.outer_gradient = vec![1.0]; + assert_eq!( + route_push(std::slice::from_mut(&mut round), island, own, Some(2), false), + PushDisposition::Accepted + ); + let mut carried = test_push_for(1, 2); + carried.fragment_id = 0; + carried.c_tokens = 40; // tokens^2/steps: 4x the own push's weight + carried.outer_gradient = vec![3.0]; + assert_eq!(attach_carry(&mut round, late, carried, 2), Ok(2)); + let merge = prepare_round_compute(cfg, &st, &mut round).unwrap().finish_inline().unwrap(); + st.commit_merge(merge).unwrap(); + st.params[0][0] + } + + #[test] + fn carried_delta_merges_against_its_own_base_with_its_token_weight() { + // Outer lr 1, no momentum: the fragment moves by minus the weighted + // mean outer gradient, (1·1 + 4·3) / 5 = 2.6. + let mut cfg = round_test_config(1); + cfg.max_base_lag = Some(2); + cfg.learner_weight = LearnerWeight::Tokens2OverSteps; + assert!((carried_merge(&cfg) + 2.6).abs() < 1e-6); + let mut streaming = cfg.clone(); + streaming.weighted_fold = WeightedFoldMode::Streaming; + assert!((carried_merge(&streaming) + 2.6).abs() < 1e-6); + // Carry discount 0.5 per version, lag 2: weight 4 · 0.25 = 1. + let mut discounted = cfg.clone(); + discounted.membership.carry_gamma = 0.5; + assert!((carried_merge(&discounted) + 2.0).abs() < 1e-6); + } + + #[test] + fn step_rates_time_a_step_from_its_first_fragment_push() { + // One local step pushes every fragment; the repeats must not move the + // step start, or the estimate shrinks to the gap between steps. + let mut rates = StepRates::default(); + let t0 = Instant::now(); + let island = member(1, 10); + for (step, start) in [(1u64, 0u64), (2, 10)] { + for fragment in 0..4 { + rates.note(island, step, t0 + Duration::from_secs(start + 2 * fragment)); + } + } + let est = rates.max_step_secs_for(&[island]).unwrap(); + assert!((est - 10.0).abs() < 1e-9, "step start to step start is 10 s, got {est}"); + } + #[test] fn step_rates_smooth_with_ema() { let mut rates = StepRates::default(); @@ -7832,6 +8356,7 @@ mod tests { quorum_ms: None, grace_ms: None, fold: None, + carried: Vec::new(), pushes: HashMap::new(), } } @@ -8266,6 +8791,7 @@ mod tests { 33, None, &pushes, + &[], 0.5, 44, &[7; 32], @@ -8315,6 +8841,7 @@ mod tests { 2, None, &pushes, + &[], 0.5, 3, &[7; 32], @@ -8344,6 +8871,7 @@ mod tests { 2, None, &pushes, + &[], 0.5, 3, &[7; 32], @@ -8454,6 +8982,7 @@ mod tests { 2, None, &pushes, + &[], 0.5, 3, &[7; 32], diff --git a/tests/test_rl_late_carry_real_syncer.py b/tests/test_rl_late_carry_real_syncer.py new file mode 100644 index 000000000..251830509 --- /dev/null +++ b/tests/test_rl_late_carry_real_syncer.py @@ -0,0 +1,359 @@ +"""Late-push carry against the real Rust syncer (CPU, no Ray). + +With quorum K below the island count, a round may close before a slightly +slower island's push arrives. In asynchronous rounds (a positive base-lag +bound) the syncer merges that late delta into the next round of the same +fragment, against its own base and with its own token weight, instead of +dropping it. Each island anchor yields one delta: a push from an anchor +whose delta was already taken is refused, so nothing is counted twice. + +The rig is the launcher-built syncer of ``test_rl_outer_merge_real_syncer`` +(P=2 fragments, pipeline 1): round ``t`` merges fragment ``(t - 1) % 2``. +""" + +from __future__ import annotations + +import json +import threading +import time +from pathlib import Path + +import numpy as np +import pytest + +from test_rl_outer_merge_real_syncer import (DEFAULT_LR, DEFAULT_MU, NUMEL, Reference, # noqa: F401 + _check, _deltas, _f32, _from, _initial, make_rig, + syncer_binary) +from test_rl_algorithm_provenance import _fake_sky, _no_modal_listing # noqa: F401 + +ASYNC_K1 = ("--rl-sync-quorum", "1", "--rl-sync-max-base-lag", "1", + "--rl-sync-learner-weight", "tokens2-over-steps") + + +@pytest.fixture(autouse=True) +def _hmac_key(monkeypatch): + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "late-carry-test-key") + + +def _tape(rig): + path = Path(rig.command.split("--event-tape ")[1].split()[0]) + if not path.exists(): + return [] + with path.open() as tape: + return [json.loads(line) for line in tape if line.strip()] + + +def _merge_records(rig): + return {r["step"]: r for r in _tape(rig) if "responders" in r and "step" in r} + + +def _weight(c_steps, c_tokens): + return c_tokens * c_tokens / c_steps + + +def test_late_push_is_merged_once_into_the_next_round_of_its_fragment(make_rig): + rig = make_rig(2, *ASYNC_K1, "--rl-sync-grace-ms", "0", total_steps=4) + assert "--quorum 1" in rig.command and "--max-base-lag 2" in rig.command + initial = _initial(rig.layout) + rig.start(initial) + ref = Reference(rig.layout, initial, DEFAULT_LR, DEFAULT_MU) + d = _deltas(rig.layout, 0, 7, 4) + + # Round 1 (fragment 0, base 0) closes on island 0 alone. + rig.push(0, 1, d[0], base=0, c_steps=1, c_tokens=4) + _check(rig.bcast(0, 0, 1), ref.step(0, [d[0]], [1.0])) + # Island 1's delta for round 1 arrives late; fragment 1 has round 2 open. + rig.push(1, 1, d[1], base=0, c_steps=2, c_tokens=6) + d1 = _deltas(rig.layout, 1, 8, 1)[0] + rig.push(0, 2, d1, base=0) + _check(rig.bcast(0, 1, 2), ref.step(1, [d1], [1.0])) + + # Round 3 (fragment 0, base 1). Island 1 answers from its old anchor 0 + # again: that anchor already gave the carried delta, so it is refused. + rig.push(1, 3, d[2], base=0, c_steps=3, c_tokens=9) + with pytest.raises(TimeoutError): + rig.bcast(0, 0, 3, timeout=1.0) + rig.push(0, 3, d[3], base=1, c_steps=1, c_tokens=4) + want = ref.step(0, [d[3], d[1]], [_weight(1, 4), _weight(2, 6)]) + _check(rig.bcast(0, 0, 3), want) + + records = _merge_records(rig) + assert [r["id"] for r in records[3]["responders"]] == [0] + (carried,) = records[3]["carried"] + assert (carried["id"], carried["missed_step"], carried["base_version"], carried["lag"]) == (1, 1, 0, 1) + assert "carried" not in records[1] and "carried" not in records[2] + + +def test_late_push_beyond_the_base_lag_bound_is_dropped(make_rig): + rig = make_rig(2, *ASYNC_K1, "--rl-sync-grace-ms", "0", total_steps=6) + initial = _initial(rig.layout) + rig.start(initial) + ref = Reference(rig.layout, initial, DEFAULT_LR, DEFAULT_MU) + for step in (1, 2, 3): + fid = (step - 1) % 2 + delta = _deltas(rig.layout, fid, step, 1)[0] + rig.push(0, step, delta, base=max(step - 2, 0)) + _check(rig.bcast(0, fid, step), ref.step(fid, [delta], [1.0])) + # Fragment 0 is at version 3: a delta from base 0 lags 3 > 2 versions. + rig.push(1, 1, np.ones(NUMEL), base=0) + delta = _deltas(rig.layout, 1, 4, 1)[0] + rig.push(0, 4, delta, base=2) + _check(rig.bcast(0, 1, 4), ref.step(1, [delta], [1.0])) + delta = _deltas(rig.layout, 0, 5, 1)[0] + rig.push(0, 5, delta, base=3) + _check(rig.bcast(0, 0, 5), ref.step(0, [delta], [1.0])) + assert all("carried" not in r for r in _merge_records(rig).values()) + assert rig.process.poll() is None + + +def test_strict_control_still_drops_a_late_push(make_rig): + rig = make_rig(2, "--rl-sync-quorum", "1", "--rl-sync-grace-ms", "0", total_steps=4) + assert "--max-base-lag 0" in rig.command + initial = _initial(rig.layout) + rig.start(initial) + ref = Reference(rig.layout, initial, DEFAULT_LR, DEFAULT_MU) + d = _deltas(rig.layout, 0, 3, 2) + rig.push(0, 1, d[0], base=0) + _check(rig.bcast(0, 0, 1), ref.step(0, [d[0]], [1.0])) + rig.push(1, 1, d[1], base=0) # late: dropped + d1 = _deltas(rig.layout, 1, 4, 1)[0] + rig.push(0, 2, d1, base=0) + _check(rig.bcast(0, 1, 2), ref.step(1, [d1], [1.0])) + d3 = _deltas(rig.layout, 0, 5, 1)[0] + rig.push(0, 3, d3, base=1) + _check(rig.bcast(0, 0, 3), ref.step(0, [d3], [1.0])) + assert all("carried" not in r for r in _merge_records(rig).values()) + + +# ---- many rounds with free-running islands ----------------------------------------- + +class Island(threading.Thread): + """A fake island that trains in fixed-length local steps. + + Each step adds a fixed per-island increment to every fragment. At a step + end the island answers every PULL it holds with ``local - anchor`` of + that fragment, then keeps training; a broadcast becomes the fragment's + anchor at a step end and the progress not yet pushed is kept on top of + it (``theta <- G + unpushed progress``). A broadcast older than an open + PULL is applied before answering it; one that closes the PULL's own + round is applied after the push, so that push may arrive late. Like the + pool client, it heartbeats several times per step with its step count. + Every delta sent is recorded by (step, fragment, base).""" + + def __init__(self, rig, index, step_s, stop): + super().__init__(daemon=True) + self.rig, self.index, self.step_s, self.stop = rig, index, step_s, stop + self.client = rig.clients[index] + fragments = rig.layout.fragments + rng = np.random.default_rng(100 + index) + self.increment = [rng.normal(size=f.numel).astype(np.float32).astype(np.float64) * 0.01 + for f in fragments] + self.version = [0] * len(fragments) + self.unpushed = [0] * len(fragments) + self.staged = {} + self.pending = {} + self.sent = {} + self.updates = {} + self.steps = 0 + self.error = None + + def _apply(self, fid, below): + for (f, version) in sorted(self.staged): + if f == fid and version < below: + self.staged.pop((f, version)) + self.version[fid] = max(self.version[fid], version) + + def run(self): + try: + while not self.stop.is_set(): + for _ in range(4): # heartbeats repeat the current step count + time.sleep(self.step_s / 4) + self.client.heartbeat(self.steps) + self.steps += 1 + self.client.check_health() + for fid in range(len(self.unpushed)): + self.unpushed[fid] += 1 + for update in self.client.drain_updates(): + self.updates[(update.fragment_id, update.version)] = _from(update.data) + self.staged[(update.fragment_id, update.version)] = True + for pull in self.client.drain_pulls(): + self.pending[pull.fragment_id] = pull + for fid, pull in sorted(self.pending.items()): + self._apply(fid, pull.global_step) + steps = self.unpushed[fid] + delta = (self.increment[fid] * steps).astype(np.float32) + base = self.version[fid] + self.sent[(pull.global_step, fid, base)] = (delta.astype(np.float64), steps) + self.client.push_fragment(fid, pull.global_step, pull.round_attempt, base, + self.steps, steps, 10 * steps, delta.tobytes()) + self.unpushed[fid] = 0 + self.pending.clear() + for fid in range(len(self.unpushed)): + self._apply(fid, float("inf")) + except Exception as error: # noqa: BLE001 - surfaced by the test + if not self.stop.is_set(): + self.error = error + + +def _run_free(make_rig, step_s, rounds, grace_ms): + rig = make_rig(2, *ASYNC_K1, "--rl-sync-grace-ms", str(grace_ms), total_steps=rounds + 2) + initial = _initial(rig.layout) + rig.start(initial) + stop = threading.Event() + islands = [Island(rig, i, s, stop) for i, s in enumerate(step_s)] + for island in islands: + island.start() + deadline = time.monotonic() + 60 + rounds * max(step_s) * 4 + try: + while time.monotonic() < deadline: + records = _merge_records(rig) + if all(t in records for t in range(1, rounds + 1)): + break + for island in islands: + assert island.error is None, island.error + time.sleep(0.1) + else: + raise TimeoutError(f"syncer merged only {sorted(_merge_records(rig))}") + time.sleep(2 * max(step_s)) # let island 0 drain the last broadcasts + finally: + stop.set() + for island in islands: + island.join(timeout=5) + records = _merge_records(rig) + return rig, initial, islands, [records[t] for t in range(1, rounds + 1)] + + +def check_free_run(rig, initial, islands, records): + """Every merge equals the numpy reference over the deltas its tape record + names (carried ones included); no delta is taken twice. Returns the + number of rounds each island contributed to, and its late deltas.""" + ref = Reference(rig.layout, initial, DEFAULT_LR, DEFAULT_MU) + taken = set() + participated = [0] * len(islands) + carried = [0] * len(islands) + for record in records: + fid, step = record["fragment"], record["step"] + entries = ([(r, step) for r in record["responders"]] + + [(c, c["missed_step"]) for c in record.get("carried", [])]) + deltas, weights, seen = [], [], set() + for entry, sent_step in entries: + island = entry["id"] + key = (island, sent_step, fid, entry["base_version"]) + assert key not in taken, f"delta {key} merged twice" + taken.add(key) + seen.add(island) + delta, steps = islands[island].sent[key[1:]] + assert entry["c_steps"] == steps + deltas.append(delta) + weights.append(_weight(entry["c_steps"], entry["c_tokens"])) + for island in seen: + participated[island] += 1 + for entry in record.get("carried", []): + carried[entry["id"]] += 1 + want = ref.step(fid, deltas, weights) + got = islands[0].updates.get((fid, step)) + if got is not None: + _check(got, want) + return taken, participated, carried + + +def _summary(rounds, records, participated, carried): + direct = [sum(any(r["id"] == i for r in rec["responders"]) for rec in records) for i in (0, 1)] + return f"rounds {rounds}: merged-in {participated}, direct {direct}, carried {carried}" + + +@pytest.mark.parametrize("step_s", [(0.25, 0.25), (0.25, 0.28)], ids=["equal", "slower"]) +def test_free_running_islands_lose_no_delta_and_count_none_twice(make_rig, step_s): + rounds = 24 + rig, initial, islands, records = _run_free(make_rig, step_s, rounds, grace_ms=0) + taken, participated, carried = check_free_run(rig, initial, islands, records) + print(_summary(rounds, records, participated, carried)) + # Every delta sent for an early enough round was merged exactly once. + for index, island in enumerate(islands): + lost = [key for key in island.sent + if key[0] <= rounds - 2 and (index, *key) not in taken] + assert not lost, f"island {index} lost deltas {lost}" + # Both islands contribute to many rounds. The count depends on load: an + # island cannot answer a round it never held a PULL for (a later PULL of + # the same fragment in the same local step replaces it). The no-loss + # check above is the exact one. + assert min(participated) >= rounds // 2, participated + + +def test_adaptive_grace_waits_for_a_slightly_slower_island(make_rig): + rounds = 24 + rig, initial, islands, records = _run_free(make_rig, (0.25, 0.28), rounds, grace_ms=2000) + taken, participated, carried = check_free_run(rig, initial, islands, records) + print(_summary(rounds, records, participated, carried)) + direct = sum(any(r["id"] == 1 for r in rec["responders"]) for rec in records) + assert direct >= rounds * 3 // 4, _summary(rounds, records, participated, carried) + + +class ReactiveIsland(threading.Thread): + """An island that starts a local round on each PULL and answers it + ``work_s`` later from the newest broadcast, heartbeating four times per + round with its round count (like the pool client). Two of them with + different ``work_s`` keep a fixed lateness between their pushes.""" + + def __init__(self, rig, index, work_s, stop): + super().__init__(daemon=True) + self.client, self.index, self.work_s, self.stop = rig.clients[index], index, work_s, stop + self.version = [0] * rig.layout.num_fragments + self.rounds = 0 + self.error = None + + def run(self): + try: + while not self.stop.is_set(): + pulls = self.client.drain_pulls() + if not pulls: + time.sleep(0.005) + continue + for _ in range(4): + time.sleep(self.work_s / 4) + self.client.heartbeat(self.rounds) + for update in self.client.drain_updates(): + fid = update.fragment_id + self.version[fid] = max(self.version[fid], update.version) + self.rounds += 1 + for pull in pulls: + fid = pull.fragment_id + self.client.push_fragment(fid, pull.global_step, pull.round_attempt, + self.version[fid], self.rounds, 1, 10, + _f32(np.full(NUMEL, 0.01 * (self.index + 1)))) + except Exception as error: # noqa: BLE001 - surfaced by the test + if not self.stop.is_set(): + self.error = error + + +def test_grace_window_waits_for_a_peer_that_is_late_by_less_than_the_cap(make_rig): + """GPU run G-A scaled 1:100: a round of ~100 s, island 1 always ~20 s + behind island 0, grace cap 30 s. The grace window must wait for it, so + island 1 answers most rounds directly (it used to answer 1 of 24).""" + rounds = 24 + rig = make_rig(2, *ASYNC_K1, "--rl-sync-grace-ms", "300", total_steps=rounds + 2) + rig.start(_initial(rig.layout)) + stop = threading.Event() + islands = [ReactiveIsland(rig, 0, 0.8, stop), ReactiveIsland(rig, 1, 1.0, stop)] + for island in islands: + island.start() + deadline = time.monotonic() + 300 + try: + while time.monotonic() < deadline: + records = _merge_records(rig) + if all(t in records for t in range(1, rounds + 1)): + break + for island in islands: + assert island.error is None, island.error + time.sleep(0.1) + else: + raise TimeoutError(f"syncer merged only {sorted(_merge_records(rig))}") + finally: + stop.set() + for island in islands: + island.join(timeout=5) + records = [_merge_records(rig)[t] for t in range(1, rounds + 1)] + direct = sum(any(r["id"] == 1 for r in rec["responders"]) for rec in records) + grace = [rec["grace_ms"] for rec in records] + print(f"island 1 direct in {direct}/{rounds} rounds; grace_ms {grace}") + assert direct >= rounds - 3, (direct, grace) diff --git a/tests/test_wandb_tape.py b/tests/test_wandb_tape.py index 2305d7fcd..f1434a1ac 100644 --- a/tests/test_wandb_tape.py +++ b/tests/test_wandb_tape.py @@ -184,7 +184,7 @@ def test_every_rust_tape_field_is_consumed_or_deliberately_skipped(): consumed_top = { "step", "fragment", "attempt", "gnorm", "ms", "sync_ms", "quorum", "expected", "responded", "missed_grace", "launch_base_version", - "quorum_ms", "grace_ms", "responders", "checkpoint_ms", + "quorum_ms", "grace_ms", "responders", "checkpoint_ms", "carried", } # Constants and the member-identity mirrors of the id lists: nothing a # time series can be drawn from. @@ -200,6 +200,7 @@ def test_every_rust_tape_field_is_consumed_or_deliberately_skipped(): "id", "staleness", "contribution", "weight", "c_steps", "c_tokens", "accounted_c_steps", "accounted_c_tokens", "base_version", "recv_seconds", + "carried", "missed_step", "lag", # the "carried" list, built before the record } # arrived_unix: wall clock of the push, for offline timing analysis. skipped_responder = {"generation", "arrived_unix"} diff --git a/yeto/wandb_tape.py b/yeto/wandb_tape.py index 59c4e584a..0a3cdb70b 100644 --- a/yeto/wandb_tape.py +++ b/yeto/wandb_tape.py @@ -121,6 +121,13 @@ def tape_metrics(rec: dict) -> dict: if value is not None: metrics[f"sync/{optional}"] = float(value) + # Late deltas of earlier rounds merged here (async rounds). + carried = rec.get("carried") or [] + metrics["sync/carried"] = float(len(carried)) + if carried: + metrics["sync/carried_lag_max"] = max(float(c.get("lag", 0)) for c in carried) + metrics["sync/carried_missed_step_min"] = min(float(c.get("missed_step", 0)) for c in carried) + stalenesses = [float(r.get("staleness", 0)) for r in responders] if stalenesses: metrics["sync/staleness_max"] = max(stalenesses) From fce646b119fe00d2eaba6d90e3e37854f22b8e11 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:59:21 +0000 Subject: [PATCH 146/233] syncer, rl: heartbeats keep an island's lease from HELLO through startup The lease runs from HELLO, but while the syncer waited for INIT and the start quorum it dropped heartbeats, so an island that connected first lost its lease while another island was still warming up (rounds=0) and re-JOINed once the pool started. How it is wired: - syncer: heartbeats received before the initial cut renew the sender's lease, as they do in the round loop. - Miles full-parameter role stream: the recovering path (peer restore) starts the island's lease heartbeat right after HELLO and keeps it, as the normal initial-cut path already does; it is stopped on failure. Test: real syncer, one island heartbeating, the INIT sent after 2.5 leases; no lease expires (it expired before this change). Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 7 +++- .../test_rl_full_parameter_restart_rejoin.py | 38 +++++++++++++++++++ .../miles/full_parameter/role_stream.py | 9 ++++- 3 files changed, 51 insertions(+), 3 deletions(-) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index aaba34d72..31a093653 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -2552,7 +2552,12 @@ async fn scheduler( } } Event::Push { .. } => warn!("push before initialization; dropped"), - Event::Heartbeat { .. } | Event::Leave { .. } => {} + // The lease runs from HELLO: an island that waits for the other + // islands' INIT (or warms up its engine) keeps it by heartbeat. + Event::Heartbeat { member, .. } => { + registry.lock().unwrap().membership.renew(member, Instant::now()); + } + Event::Leave { .. } => {} Event::FinalAck { member, .. } => warn!( learner_id = member.learner_id, generation = member.generation, diff --git a/tests/test_rl_full_parameter_restart_rejoin.py b/tests/test_rl_full_parameter_restart_rejoin.py index a6c93c849..dc7b38055 100644 --- a/tests/test_rl_full_parameter_restart_rejoin.py +++ b/tests/test_rl_full_parameter_restart_rejoin.py @@ -116,3 +116,41 @@ def test_a_restarted_island_catches_up_and_finishes(syncer_binary, tmp_path): assert min(restarted["catch_up"].values()) >= 2 assert restarted["pushed_after"] > 0, "the restarted island never contributed" assert '"catch_up":true' in tape + + +@needs_cargo +@pytest.mark.usefixtures("_hmac_key") +def test_heartbeats_keep_the_lease_while_the_pool_waits_for_init(syncer_binary, tmp_path): + """s27-x2-pause-a: an island that connected first lost its lease while + the other island was still starting (no INIT yet), because the syncer + ignored heartbeats before the initial cut. Here the INIT comes after + 2.5 leases; the first island's heartbeats keep it in the pool.""" + import time + + from yeto.rl.pool_client import LeaseHeartbeat + + lease_s = 3.0 + log_file = tmp_path / "syncer.log" + process, port, profile_hash = _start_async_syncer(syncer_binary, tmp_path, *LEASE, + total_steps=2, log_file=log_file) + early = _client(1, port, profile_hash) + late = _client(0, port, profile_hash) + heartbeat = LeaseHeartbeat(early, lease_s) + try: + early.start() + heartbeat.start() + time.sleep(2.5 * lease_s) # island 0 warms up its engine + late.start() + for fid in range(P): + late.send_init(fid, _f32(_initial(fid))) + _wait(lambda: "global parameters initialized" in log_file.read_text(), "no INIT") + time.sleep(0.5) + log = log_file.read_text() + assert "lease expired" not in log, log + assert heartbeat.errors == [] + finally: + heartbeat.stop() + early.close() + late.close() + process.terminate() + process.wait(timeout=10) diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 80e088b76..0e4072f97 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -557,7 +557,7 @@ async def initialize(self, group) -> None: sink = adapter.fragment_sink(initial, ray_module=self.ray_module) client.install_inbound_chunk_sink(MilesRayInboundChunkSink(sink)) if self._recovering: - await self._initialize_recovering(group, adapter, initial, client, sink) + await self._initialize_recovering(group, adapter, initial, client, sink, heartbeat) return try: client.start() @@ -668,7 +668,7 @@ def _recovery_island(self, group, adapter, *, recovering: bool): shards = MilesRankShards(group, adapter.layout.layout_hash) return MilesPeerRecoveryIsland(self, shards, asyncio.get_running_loop(), recovering=recovering) - async def _initialize_recovering(self, group, adapter, initial, client, sink) -> None: + async def _initialize_recovering(self, group, adapter, initial, client, sink, heartbeat) -> None: """S27 E-3: this trainer is fresh after the island lost its state. No INIT is sent (the syncer holds the run's fragments). A donor's @@ -682,6 +682,8 @@ async def _initialize_recovering(self, group, adapter, initial, client, sink) -> count = adapter.layout.fragments.num_fragments try: client.start() + if heartbeat is not None: + heartbeat.start() # a peer restore may take longer than the lease outcome = await asyncio.to_thread( recover, client, island, self._snapshot_store(), timeout=options.timeout_s) client.mark_initial_cut_received() @@ -690,6 +692,7 @@ async def _initialize_recovering(self, group, adapter, initial, client, sink) -> self.group = group self.adapter = adapter self.client = client + self.heartbeat = heartbeat self.sink = sink self.current = initial self._anchors = [None] * count @@ -700,6 +703,8 @@ async def _initialize_recovering(self, group, adapter, initial, client, sink) -> resent, island.resent = dict(island.resent), {} await self._apply_authoritative(resent, optimizer_proof=False) except BaseException: + if heartbeat is not None: + heartbeat.stop() island.release() client.close() sink.release_all() From 5634a338fe9e3f0214ec6299c8853d596f0489d1 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:00:39 +0000 Subject: [PATCH 147/233] test: real syncer control for tensor-aligned push blocks How it is wired: the real-syncer sweep pushes a 70-value tensor (not a multiple of the quantization block) next to different neighbours in two fragments. With tensor blocks the two decoded copies are bit-identical (existing test); the new control runs the same sweep without the flag and checks that the copies differ, for int8, int4 and int4-log, so the identical-copy check has power. Co-Authored-By: Claude Opus 5.5 --- tests/test_s27_tensor_blocks.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_s27_tensor_blocks.py b/tests/test_s27_tensor_blocks.py index c03715442..063d03db1 100644 --- a/tests/test_s27_tensor_blocks.py +++ b/tests/test_s27_tensor_blocks.py @@ -369,6 +369,19 @@ def test_real_syncer_decodes_tensor_blocks_and_keeps_the_replica_identical( finals[1][start[1]:start[1] + TENSOR.size]) +@pytest.mark.parametrize(("codec", "base"), (("int8", DTYPE_F32_PUSH_I8), ("int4", DTYPE_F32_PUSH_I4), + ("int4-log", DTYPE_F32_PUSH_Q4))) +def test_real_syncer_without_tensor_blocks_lets_the_replica_drift( + syncer_binary, tmp_path, codec, base): # noqa: F811 + """Control for the test above: TENSOR (70 values, not a multiple of the + block) shares a block with different neighbours in each fragment, so + without the flag the syncer's two decoded copies differ.""" + _, finals = _sweep_through_real_syncer(syncer_binary, tmp_path, base) + start = (LOUD.size, QUIET.size) + assert not np.array_equal(finals[0][start[0]:start[0] + TENSOR.size], + finals[1][start[1]:start[1] + TENSOR.size]) + + def test_real_syncer_refuses_the_flag_on_an_unquantized_session(syncer_binary, tmp_path): # noqa: F811 process, port, profile_hash = _start_syncer(syncer_binary, tmp_path) layout = FragmentLayout([Fragment(MERGE_RDA, [(f"w{i}", 3)]) for i in range(FRAGMENTS)]) From c6d68c194e65e1e1ebe3664f578a90cd5b7a50c4 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:30:19 +0000 Subject: [PATCH 148/233] Ports image: Miles cc693c693 (full-parameter sync with distributed optimizer + CPU offload) How it is wired: the ports image and pins move to Miles cc693c693 (ghcr.io/michaellchung/yeto-miles-ports:cc693c6-2fa8801 @sha256:2011cbf9bf6e2f472a0819d77685d6cec56f4941957ae5652f09f3ccb7604e08). 56190a5a8 (full-parameter sync reads DP=1 distributed-optimizer masters, incl. CPU offload) and cc693c693 (the optimizer-state proof leaves out the HybridDeviceOptimizer's master_param mirror, which changed on every global apply with CPU offload) are appended to MILES_FORK_LINEAGE. MILES_DISTOPT_FULL_PARAMETER_COMMIT is 56190a5a8, an ancestor of the pin, so the config accepts optimizer offload with the full-parameter sync. The 56190a5-2fa8801 image is not pinned. Co-Authored-By: Claude Opus 5.5 --- docker/miles-ports/Dockerfile | 2 +- .../build-record.json | 12 +++++ .../image-manifest.json | 50 +++++++++++++++++++ tests/test_rl_miles_adapter_config.py | 7 +-- tests/test_rl_ports_image.py | 2 +- yeto/rl/adapters/miles/pins.py | 19 +++++-- 6 files changed, 80 insertions(+), 12 deletions(-) create mode 100644 openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/build-record.json create mode 100644 openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/image-manifest.json diff --git a/docker/miles-ports/Dockerfile b/docker/miles-ports/Dockerfile index 11aeb800e..2f4e7bf56 100644 --- a/docker/miles-ports/Dockerfile +++ b/docker/miles-ports/Dockerfile @@ -19,7 +19,7 @@ FROM docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c ARG MILES_REPOSITORY=https://github.com/michaellchung/miles -ARG MILES_COMMIT=ed5f64fd7a183cccc8ab7ca3da90fba570108c2b +ARG MILES_COMMIT=cc693c6930be178f411f711f55c3d334a087c097 ARG SGLANG_REPOSITORY=https://github.com/michaellchung/sglang ARG SGLANG_COMMIT=2fa880182eefdb31e64d8ea70317be27f9f32121 # setuptools-scm version of SGLANG_COMMIT (describe v0.5.20-69-g4e4148f1b); diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/build-record.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/build-record.json new file mode 100644 index 000000000..51682b438 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/build-record.json @@ -0,0 +1,12 @@ +{ + "tag": "ghcr.io/michaellchung/yeto-miles-ports:cc693c6-2fa8801", + "digest": "sha256:2011cbf9bf6e2f472a0819d77685d6cec56f4941957ae5652f09f3ccb7604e08", + "layer_tar_sha256": "c0e0273ca2f74ff912296730a79a699c6fb603e61147fde41d56bbbf58c8671a", + "base": "docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "base_index": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "miles_commit": "cc693c6930be178f411f711f55c3d334a087c097", + "sglang_commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "sglang_version": "0.5.21.dev70+g2fa8801", + "peft_version": "0.20.0", + "peft_wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/image-manifest.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/image-manifest.json new file mode 100644 index 000000000..ce15ff5a1 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/image-manifest.json @@ -0,0 +1,50 @@ +{ + "base": { + "index_digest": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "manifest_digest": "sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "miles_commit": "9e4260de047a704208535c0e90c531929879ab40", + "platform": "linux/amd64", + "repository": "docker.io/radixark/miles", + "sglang_commit": "e5bba1f9c433da5352e521a8bb826bab960e7836" + }, + "builder": "scripts/build_miles_ports_image.sh (crane append); docker/miles-ports/Dockerfile equivalent", + "image": "ghcr.io/michaellchung/yeto-miles-ports", + "megatron_bridge": { + "commit": "8cd3466d14d2337c8492827b3712482c2b3e4866", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_bridge_init.py)" + ] + }, + "megatron_lm": { + "commit": "a84b105473eca9c9c6e49dbf1c8aa460c5cf796e", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_lm_lazy_modelopt.patch)" + ], + "path": "/root/Megatron-LM", + "repository": "https://github.com/radixark/Megatron-LM" + }, + "miles": { + "commit": "cc693c6930be178f411f711f55c3d334a087c097", + "install": "base's editable install (pip install -e /root/miles --no-deps) kept (same path); /root is on the base's sys.path, so without PYTHONPATH import miles is a namespace package over the /root/miles checkout -- yeto runs with PYTHONPATH=$HOME/miles (regular package)", + "path": "/root/miles", + "repository": "https://github.com/michaellchung/miles", + "upstream_commit": "9e4260de047a704208535c0e90c531929879ab40" + }, + "peft": { + "install": "wheel unpacked into the base's site-packages (pip install --no-deps equivalent)", + "version": "0.20.0", + "wheel": "peft-0.20.0-py3-none-any.whl", + "wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" + }, + "pure_python_overlay": true, + "schema": 1, + "sglang": { + "commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "install": "pip install -e /sgl-workspace/sglang/python[all] --no-deps (base's editable install, same path)", + "path": "/sgl-workspace/sglang", + "repository": "https://github.com/michaellchung/sglang", + "upstream_commit": "571212b636baca45e10fa3b4da11a289123f3235", + "version": "0.5.21.dev70+g2fa8801" + }, + "tag": "cc693c6-2fa8801" +} \ No newline at end of file diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index 6408dfdb4..6de68515c 100644 --- a/tests/test_rl_miles_adapter_config.py +++ b/tests/test_rl_miles_adapter_config.py @@ -281,17 +281,14 @@ def test_optimizer_offload_with_the_full_parameter_sync_needs_the_distopt_reader # older pinned image is refused before launch. cfg = sub(make_config(), "trainable", parameter_mode="full", target_modules=()) cfg = dataclasses.replace(cfg, train_memory=rc.TrainMemoryConfig(optimizer_offload_fraction=0.5)) + monkeypatch.setattr(mc, "MILES_NEXT_COMMIT", "ed5f64fd7a183cccc8ab7ca3da90fba570108c2b") with pytest.raises(mc.MilesConfigError, match="56190a5a8 or later"): mc.translate_run_config(cfg, AlgorithmSpec()) -def test_optimizer_offload_with_the_full_parameter_sync_keeps_the_distributed_optimizer(monkeypatch): +def test_optimizer_offload_with_the_full_parameter_sync_keeps_the_distributed_optimizer(): cfg = sub(make_config(), "trainable", parameter_mode="full", target_modules=()) cfg = dataclasses.replace(cfg, train_memory=rc.TrainMemoryConfig(optimizer_offload_fraction=1.0)) - from yeto.rl.adapters.miles import pins - - monkeypatch.setattr(pins, "MILES_FORK_LINEAGE", pins.MILES_FORK_LINEAGE + (mc.MILES_DISTOPT_FULL_PARAMETER_COMMIT,)) - monkeypatch.setattr(mc, "MILES_NEXT_COMMIT", mc.MILES_DISTOPT_FULL_PARAMETER_COMMIT) launch = mc.translate_run_config(cfg, AlgorithmSpec()) assert "use_distributed_optimizer" not in launch.runtime_attrs assert ("--use-distributed-optimizer --use-precision-aware-optimizer --optimizer-cpu-offload " diff --git a/tests/test_rl_ports_image.py b/tests/test_rl_ports_image.py index fb0646ef0..f24a7b38e 100644 --- a/tests/test_rl_ports_image.py +++ b/tests/test_rl_ports_image.py @@ -70,7 +70,7 @@ def test_build_inputs_match_the_pins(): def test_commits_and_digest_are_pinned_together_with_the_build_record(): """Commit pins, tag, Dockerfile ARGs and the image digest move together: the latest build record must name exactly the pinned commits and digest.""" - record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-ed5f64f-2fa8801" + record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801" record = json.loads((record_dir / "build-record.json").read_text()) assert record["miles_commit"] == rl.MILES_NEXT_COMMIT assert record["sglang_commit"] == rl.SGLANG_NEXT_COMMIT diff --git a/yeto/rl/adapters/miles/pins.py b/yeto/rl/adapters/miles/pins.py index 494be16b7..8d379c3a8 100644 --- a/yeto/rl/adapters/miles/pins.py +++ b/yeto/rl/adapters/miles/pins.py @@ -53,7 +53,10 @@ class MilesRevisionPins(NamedTuple): # import stays light, so RayWorkerManager no longer loads megatron.bridge # (rl-trainer-cold-start-cost L2.2)). 17aeb5cb5 = 386ae3b07 + light Ray actors # stay off megatron_utils (L2.1) + the master port is held until use (S23 defect 9). -MILES_NEXT_COMMIT = "ed5f64fd7a183cccc8ab7ca3da90fba570108c2b" +# s27-distopt-fp (56190a5a8 = ed5f64fd7 + full-parameter sync reads DP=1 +# distributed-optimizer masters, incl. CPU offload; S27 G-B); cc693c693 = + the +# Adam proof leaves the hybrid optimizer's master_param mirror out (S27 IMG2). +MILES_NEXT_COMMIT = "cc693c6930be178f411f711f55c3d334a087c097" # sgl-project/sglang ``sglang-miles`` head when radixark/miles@9e4260d was # committed (upstream's Dockerfile follows that branch unpinned). SGLANG_NEXT_REPOSITORY = "https://github.com/michaellchung/sglang" @@ -112,12 +115,18 @@ class MilesRevisionPins(NamedTuple): "9462de532f747cf688e090bbaf6b24f80e378fe8", # ed5f64fd7 = + DP-invariant state export/restore carries the CPU-offload optimizer state (S26 E6 B8); no loss path. "ed5f64fd7a183cccc8ab7ca3da90fba570108c2b", + # 56190a5a8 = + full-parameter sync reads DP=1 distributed-optimizer masters, + # incl. the HybridDeviceOptimizer's offloaded FP32 copies (S27 G-B); no loss path. + "56190a5a86d14f1eb96d592038d76a9ad8651b0c", + # cc693c693 = + the optimizer-state proof skips the HybridDeviceOptimizer's + # master_param mirror, which a global apply overwrites (S27 IMG2); no loss path. + "cc693c6930be178f411f711f55c3d334a087c097", ) # First Miles commit whose full-parameter sync reads distributed-optimizer # masters (and the HybridDeviceOptimizer's offloaded FP32 copies) at DP=1 -# (S27 G-B; local fork branch s27-distopt-fp on ed5f64fd7, no image yet). -# Append it to MILES_FORK_LINEAGE together with the image re-pin. +# (S27 G-B; fork branch yeto/s27-distopt-fp on ed5f64fd7, in the image since +# 56190a5-2fa8801; the image pin is cc693c6-2fa8801). MILES_DISTOPT_FULL_PARAMETER_COMMIT = "56190a5a86d14f1eb96d592038d76a9ad8651b0c" @@ -156,7 +165,7 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # every SHA. PUBLIC on ghcr.io/michaellchung since 2026-10-07 (anonymous # pull); a private image needs SKYPILOT_DOCKER_{USERNAME,PASSWORD,SERVER} or # --rl-image-private (yeto.launcher.registry_login_for; read:packages token). -# Tag ed5f64f-2fa8801; (S27 IMG: Miles ed5f64fd7, host-Adam round cuts (B8); previous 44d41ef-2fa8801 @sha256:5b9dc2a9...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-ed5f64f-2fa8801/). Earlier: tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt +# Tag cc693c6-2fa8801; (S27 IMG2: Miles cc693c693 = 56190a5a8 (full-parameter sync with distributed optimizer + CPU offload, G-B) + Adam proof without the hybrid master mirror; previous 56190a5-2fa8801 @sha256:5254e34e... (failed on GPU: global apply changed local Adam state), ed5f64f-2fa8801 @sha256:a4769254...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/). Earlier: tag ed5f64f-2fa8801; (S27 IMG: Miles ed5f64fd7, host-Adam round cuts (B8); previous 44d41ef-2fa8801 @sha256:5b9dc2a9...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-ed5f64f-2fa8801/). Earlier: tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt # Megatron-LM patch in the layer; previous 17aeb5c-2fa8801 @sha256:b304f857...; # build record openspec/changes/rl-infra-spec/evidence/ports-image/ # 2026-10-14-44d41ef-2fa8801/). Earlier: tag 17aeb5c-2fa8801; (S24 L2.1-L2.4: Miles s24-miles-import, peft 0.20.0 and the @@ -165,7 +174,7 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # ports-image/2026-10-14-17aeb5c-2fa8801/). Earlier: tag 386ae3b-2fa8801; (S23 P0-1.11 GPU run: Miles s23-sao-streaming-roles; previous 64b591a-2fa8801 @sha256:fa2413be...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-386ae3b-2fa8801/). Earlier: (S19 #1: critic family + algo-supplement metrics baked in; previous ddce209-2fa8801 @sha256:9c252c38...) (S18 ARU-3: Miles s18-agentic-suspend = 2f7871fb2 + --agentic-suspend-between-turns; previous 2f7871f-2fa8801 @sha256:62b4f164..., efbbc63-2fa8801 @sha256:a7990076..., 8bc5223-4e4148f @sha256:4aeafd77...) ghcr.io/michaellchung (public package); linux/amd64 only. MILES_NEXT_IMAGE = ( "docker:ghcr.io/michaellchung/yeto-miles-ports@sha256:" - "a476925466dad4df51db6955006a2230f9f88e800e44d529ac5d3f8820b19d58" + "2011cbf9bf6e2f472a0819d77685d6cec56f4941957ae5652f09f3ccb7604e08" ) MILES_NEXT_IMAGE_MANIFEST = "/opt/yeto/image-manifest.json" # Nebius VM images whose /var/lib/docker already holds a docker image's From 6d574673c4ee1dbb5bf4b5ee3b5f20ce6b40ef65 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:19:58 +0000 Subject: [PATCH 149/233] Full-parameter layout: split a parameter above the fragment bound into row slices Qwen3.8-27B at TP2 holds a 248320x5120 vocabulary shard of 2.54 GB per rank, above the 2 GiB f32 fragment bound, so the owner-affine layout refused it and forced 8-GPU islands. - ParameterLayout.create(max_fragment_numel=...) splits such a parameter into balanced dimension-0 row slices. Each slice is its own fragment tensor (#rows=a:b): averaged tables stay averaged, an RDA parameter gets RDA per slice, tensor-aligned quantized pushes treat each slice as one segment. A layout that splits nothing keeps its hash. - The Miles owner plan names a flat range per wire name for fragments that hold part of a parameter (needs Miles s27-fragbound-split a477bcbe0). - --rl-sync-fragment-max-bytes (island --sweep-fragment-max-bytes) makes the bound configurable; the default stays 2 GiB. - Head memory model for 27B TP2 with the split vocabulary. Co-Authored-By: Claude Opus 5.5 --- syncer/src/head_memory.rs | 43 ++- tests/test_rl_split_parameter_layout.py | 357 ++++++++++++++++++ yeto/launch/rl_args.py | 4 + yeto/launch/rl_island_task.py | 2 + yeto/rl/adapters/miles/island_entry.py | 1 + .../miles/models/chunked_full_parameter.py | 22 +- yeto/rl/adapters/miles/run_flags.py | 7 + yeto/rl/engine/run_flags.py | 9 + yeto/rl/local_learner.py | 179 ++++++++- 9 files changed, 605 insertions(+), 19 deletions(-) create mode 100644 tests/test_rl_split_parameter_layout.py diff --git a/syncer/src/head_memory.rs b/syncer/src/head_memory.rs index 0ffc624bc..dffd56f4c 100644 --- a/syncer/src/head_memory.rs +++ b/syncer/src/head_memory.rs @@ -84,9 +84,11 @@ static METER: meter::Meter = meter::Meter; const F32: f64 = 4.0; /// Shape of one strict full-parameter session. +#[derive(Clone, Copy)] struct Session { params: f64, - /// Largest fragment (one canonical tensor per fragment in Miles). + /// Largest fragment: at most the fragment bound (a larger tensor is + /// split into dimension-0 row slices across several fragments). max_fragment: f64, islands: f64, /// Fragment rounds in flight (`--pipeline`). @@ -318,6 +320,45 @@ fn head_memory_measured_scaled_down() { std::fs::remove_dir_all(&dir).ok(); } +/// Qwen3.8-27B at TP2: each rank owns a 124160 x 5120 vocabulary shard +/// (2.54 GB f32), above the 2 GiB fragment bound, so the layout splits it +/// into two row slices. Every fragment is then at most the bound, and the +/// per-round head terms scale with the bound instead of the largest tensor. +#[test] +fn head_memory_model_for_27b_tp2_with_split_vocabulary() { + const FRAGMENT_BOUND_SCALARS: f64 = (2u64 << 30) as f64 / F32; + let vocabulary_shard = 124_160.0 * 5120.0; + assert!(vocabulary_shard > FRAGMENT_BOUND_SCALARS); + let slices = (vocabulary_shard / FRAGMENT_BOUND_SCALARS).ceil(); + assert_eq!(slices, 2.0); + println!("islands | resident | rounds fold worst | async ckpt snapshot | peak [GiB]"); + for islands in [2.0, 4.0] { + let split = Session { + params: 27e9, + max_fragment: FRAGMENT_BOUND_SCALARS, + islands, + pipeline: 2.0, + rda: true, + downlink_anchor: true, + }; + let whole = Session { max_fragment: 152_064.0 * 5120.0, ..split }; + let (e, w) = (split.estimate(), whole.estimate()); + let peak = e.resident + e.rounds_fold_worst + e.async_checkpoint_snapshot; + println!( + "{islands} | {:.1} | {:.1} | {:.1} | {:.1}", + gib(e.resident), + gib(e.rounds_fold_worst), + gib(e.async_checkpoint_snapshot), + gib(peak), + ); + // Splitting leaves the resident state alone and bounds every + // per-round term by the fragment bound. + assert_eq!(e.resident, w.resident); + assert!(e.rounds_fold_worst < w.rounds_fold_worst); + assert!(gib(e.rounds_fold_worst) <= 2.0 * (islands + 1.0) * 2.0 + 1e-9); + } +} + /// S27 P4b: per-process head bytes for 27B with the streaming weighted fold /// (any quorum, token weights) and the head split into N syncer shards. /// diff --git a/tests/test_rl_split_parameter_layout.py b/tests/test_rl_split_parameter_layout.py new file mode 100644 index 000000000..e63e45c68 --- /dev/null +++ b/tests/test_rl_split_parameter_layout.py @@ -0,0 +1,357 @@ +"""One parameter larger than the fragment byte bound spans several fragments. + +Qwen3.8-27B at TP2 holds a 248320x5120 vocabulary shard per rank (~2.54 GB +in f32), above the 2 GiB fragment bound. The layout splits such a parameter +into dimension-0 row slices; each slice is its own fragment tensor. These +tests use a small bound in place of 2 GiB: + +* layouts that split nothing keep their exact hash and Miles plan hash; +* the Miles owner plan of a split layout installs on the real Miles check; +* two learners merged by the real Rust syncer give the numpy result, for + averaged and RDA slices (RDA is per row slice), f32 and int8 pushes. +""" + +from __future__ import annotations + +# ruff: noqa: I001 -- install the optional bridge stub before Miles import. + +import shutil +import socket +import subprocess +import sys +import time +import types +from pathlib import Path + +import numpy as np +import pytest + +pytest.importorskip("miles") + +bridge_stub = types.ModuleType("yeto.rl.adapters.miles.models.deepseek_v4.bridge") +bridge_stub.ensure_deepseek_v4_bridge = lambda: None +sys.modules.setdefault("yeto.rl.adapters.miles.models.deepseek_v4.bridge", bridge_stub) + +from miles.backends.megatron_utils.full_parameter_state import ( + FullParameterShardManifest, + FullParameterShardSpec, + FullParameterTopology, + _layout_hash, + _validate_fragment_plan_against_manifest, +) + +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from yeto.fragments import MERGE_AVG, MERGE_RDA +from yeto.protocol import DTYPE_F32, DTYPE_F32_PUSH_I8, DTYPE_FLAG_TENSOR_BLOCKS, SyncerClient +from yeto.rl.adapters.miles.models.chunked_full_parameter import ( + MilesChunkedFullParameterAdapter, +) +from yeto.rl.local_learner import ( + ComponentIdentity, + ParameterLayout, + ParameterSpec, + split_parameter_specs, +) + +EMBEDDING = "embedding.word_embeddings.weight" +MATRIX = "decoder.layers.0.mlp.linear_fc1.weight" +NORM = "decoder.final_layernorm.weight" +# Shapes per TP rank: the vocabulary shard is the only parameter above the +# 16-scalar (64-byte) test bound that stands in for 2 GiB. +SHAPES = {EMBEDDING: (10, 4), MATRIX: (3, 4), NORM: (4,)} +BOUND_BYTES = 64 + + +def _topology(rank: int) -> FullParameterTopology: + return FullParameterTopology( + tp_rank=rank, tp_size=2, pp_rank=0, pp_size=1, ep_rank=0, ep_size=1, + cp_rank=0, cp_size=1, dp_rank=0, dp_size=1, + ) + + +def _manifest(rank: int, shapes=SHAPES) -> FullParameterShardManifest: + owner = _topology(rank) + specs = tuple(sorted( + FullParameterShardSpec( + role="actor", shard_id=owner.shard_id, name=name, shape=shape, + dtype="float32", numel=int(np.prod(shape)), + ) + for name, shape in shapes.items() + )) + return FullParameterShardManifest(owner, "actor", _layout_hash(owner, specs), specs) + + +def _component() -> ComponentIdentity: + return ComponentIdentity("actor", "a" * 40, "b" * 64) + + +def _adapter(*, bound: int = BOUND_BYTES, matrix_merge: str = "avg", shapes=SHAPES): + return MilesChunkedFullParameterAdapter.create( + (_manifest(0, shapes), _manifest(1, shapes)), + algorithm="grpo", + components=(_component(),), + minimum_fragments=2, + max_fragment_bytes=bound, + max_chunk_bytes=32, + matrix_merge=matrix_merge, + ) + + +def _spec(name: str, shape: tuple[int, ...]) -> ParameterSpec: + return ParameterSpec("actor", name, shape, "float32", int(np.prod(shape)), "tp0") + + +def test_split_rows_are_balanced_and_each_fits_the_bound(): + (whole,) = split_parameter_specs([_spec(MATRIX, (3, 4))], 16) + assert whole.is_whole and whole.tensor_name == "actor::tp0::" + MATRIX + slices = split_parameter_specs([_spec(EMBEDDING, (10, 4))], 16) + assert [(unit.row_start, unit.row_stop) for unit in slices] == [(0, 3), (3, 6), (6, 10)] + assert [unit.flat_range for unit in slices] == [(0, 12), (12, 24), (24, 40)] + assert all(unit.numel <= 16 for unit in slices) + assert slices[1].tensor_name == f"actor::tp0::{EMBEDDING}#rows=3:6" + assert slices[1].shape == (3, 4) + # Qwen3.8-27B at TP2: 124160 vocabulary rows of 5120 per rank, 2 GiB bound. + big = split_parameter_specs([_spec(EMBEDDING, (124160, 5120))], (2 << 30) // 4) + assert len(big) == 2 and all(unit.numel * 4 <= 2 << 30 for unit in big) + with pytest.raises(ValueError, match="one dimension-0 row"): + split_parameter_specs([_spec(EMBEDDING, (2, 20))], 16) + + +def test_layouts_that_split_nothing_keep_their_hash(): + specs = [_spec(MATRIX, (3, 4)), _spec(NORM, (4,)), _spec(EMBEDDING, (10, 4))] + for merge in ("avg", "rda"): + plain = ParameterLayout.create( + algorithm="grpo", components=(_component(),), specs=specs, num_fragments=3, + fragment_strategy="owner_affine", matrix_merge=merge, + ) + bounded = ParameterLayout.create( + algorithm="grpo", components=(_component(),), specs=specs, num_fragments=3, + fragment_strategy="owner_affine", matrix_merge=merge, max_fragment_numel=40, + ) + assert bounded.layout_hash == plain.layout_hash + assert not bounded.has_split_parameters + large = _adapter(bound=4096) + assert not large.layout.has_split_parameters + assert all(fragment.flat_ranges == () for plan in large.plans for fragment in plan.fragments) + + +def test_split_layout_changes_the_hash_and_installs_on_the_miles_check(): + adapter = _adapter() + layout = adapter.layout + assert layout.has_split_parameters + assert layout.layout_hash != _adapter(bound=4096).layout.layout_hash + assert all(fragment.numel * 4 <= BOUND_BYTES for fragment in layout.fragments.fragments) + # Every rank's vocabulary shard spans three fragments, each one averaged. + split = [ + fragment_id for fragment_id in range(layout.fragments.num_fragments) + if any(not unit.is_whole for unit in layout.fragment_slices(fragment_id)) + ] + assert len(split) == 6 + assert all(layout.fragments.fragments[f].merge_mode == MERGE_AVG for f in split) + for plan in adapter.plans: + owner = plan.topology.tp_rank + _validate_fragment_plan_against_manifest(plan, _manifest(owner)) + assert sum(fragment.numel for fragment in plan.fragments) == sum( + int(np.prod(shape)) for shape in SHAPES.values() + ) + # Plans and layout are deterministic. + assert _adapter().layout.layout_hash == layout.layout_hash + assert [plan.plan_hash for plan in _adapter().plans] == [plan.plan_hash for plan in adapter.plans] + + +def test_rda_split_matrix_keeps_rda_per_row_slice(): + shapes = {EMBEDDING: (2, 4), MATRIX: (10, 4), NORM: (4,)} + layout = _adapter(matrix_merge="rda", shapes=shapes).layout + modes = { + unit.tensor_name: layout.fragments.fragments[fragment_id].merge_mode + for fragment_id in range(layout.fragments.num_fragments) + for unit in layout.fragment_slices(fragment_id) + } + matrix_slices = [name for name in modes if MATRIX in name] + assert len(matrix_slices) == 6 and all("#rows=" in name for name in matrix_slices) + assert {modes[name] for name in matrix_slices} == {MERGE_RDA} + assert {mode for name, mode in modes.items() if EMBEDDING in name} == {MERGE_AVG} + + +# -- real Rust syncer ------------------------------------------------------- + +ROOT = Path(__file__).resolve().parent.parent +CARGO = shutil.which("cargo") or str(Path.home() / ".cargo/bin/cargo") +LEARNERS = 2 +needs_cargo = pytest.mark.skipif(not Path(CARGO).exists(), reason="needs cargo to build the syncer") + + +@pytest.fixture(scope="module") +def syncer_binary() -> Path: + subprocess.run([CARGO, "build", "-q"], cwd=ROOT / "syncer", check=True) + return ROOT / "syncer/target/debug/yeto-syncer" + + +@pytest.fixture(autouse=True) +def _hmac_key(monkeypatch): + monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "split-parameter-test-key") + + +def _start_syncer(binary: Path, tmp_path: Path, fragments: int, *extra: str): + from yeto import launcher + + args = _prepared_args(fragments, *extra) + args.quorum = LEARNERS # a strict sweep waits for both learners + command = launcher.role_syncer_commands(args, LEARNERS, binary=str(binary))["actor"] + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + port = probe.getsockname()[1] + command = (command + .replace(f"--port {launcher.SYNCER_PORT}", f"--port {port}") + .replace(launcher.RL_SYNCER_CHECKPOINT, str(tmp_path / "state.ckpt")) + .replace(launcher.RL_SYNCER_EVENT_TAPE, str(tmp_path / "tape.jsonl"))) + process = subprocess.Popen(["bash", "-c", command], stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, text=True) + return process, port, launcher.rl_syncer_profile(args, LEARNERS).sha256 + + +def _prepared_args(fragments: int, *extra: str): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", ( + "--rl-sync-preset", "full-parameter", "--rl-placement", "fixed-partition", + "--rl-rollout-gpus", "1", "--rl-parameter-layout-sha256", "e" * 64, + "--fragments", str(fragments), "--total-steps", "1", *extra), + gpu="modal:2xh100") + args.controller = "local" + _prepare_rl_args(args) + return args + + +def test_fragment_max_bytes_flag_reaches_the_island_and_is_checked(): + from yeto.launch.rl_island_task import _rl_trainable_island_flags + + default = _rl_trainable_island_flags(_prepared_args(2), 0, 1) + assert "--sweep-fragment-max-bytes" not in default + island = _rl_trainable_island_flags( + _prepared_args(2, "--rl-sync-fragment-max-bytes", "1073741824"), 0, 1) + assert "--sweep-fragment-max-bytes 1073741824" in island + for bad in ("6", "4294967296"): + with pytest.raises(ValueError, match="--rl-sync-fragment-max-bytes a multiple of 4"): + _prepared_args(2, "--rl-sync-fragment-max-bytes", bad) + + +def _wait(predicate, what: str, timeout: float = 30.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + value = predicate() + if value: + return value + time.sleep(0.01) + raise TimeoutError(what) + + +def _rda(deltas: list[np.ndarray]) -> np.ndarray: + norms = [float(np.linalg.norm(d.astype(np.float64))) for d in deltas] + direction = sum(d.astype(np.float64) / n for d, n in zip(deltas, norms)) / len(deltas) + return direction * (sum(norms) / len(norms)) / np.linalg.norm(direction) + + +def _sweep(binary, tmp_path, layout, *, dtype=DTYPE_F32, extra=(), push_i8_block=128): + """One strict sweep of two learners; returns (initial, deltas, finals).""" + + fragments = layout.fragments + process, port, profile = _start_syncer(binary, tmp_path, fragments.num_fragments, *extra) + rng = np.random.default_rng(7) + initial = {f: rng.standard_normal(frag.numel).astype(" dict[str, str]: ("rl_sync_push_dtype", "f32", "--rl-sync-push-dtype"), ("rl_sync_checkpoint_every", 1, "--rl-sync-checkpoint-every"), ("rl_sync_push_int8_block", 128, "--rl-sync-push-int8-block"), + ("rl_sync_fragment_max_bytes", 2 << 30, "--rl-sync-fragment-max-bytes"), ("rl_sync_push_int4_block", 128, "--rl-sync-push-int4-block"), ("rl_sync_push_int4_stochastic", False, "--rl-sync-push-int4-stochastic"), ("rl_sync_push_int4_error_feedback", "off", "--rl-sync-push-int4-error-feedback"), @@ -253,6 +254,9 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: problems.append("no --rl-sync-overlap (sweep only)") elif checkpoint_every != 1 and (checkpoint_every < 1 or checkpoint_every % max(args.fragments, 1)): problems.append("--rl-sync-checkpoint-every 1 or a multiple of --fragments") + fragment_max_bytes = getattr(args, "rl_sync_fragment_max_bytes", 2 << 30) + if type(fragment_max_bytes) is not int or not 4 <= fragment_max_bytes <= 2 << 30 or fragment_max_bytes % 4: + problems.append("--rl-sync-fragment-max-bytes a multiple of 4 between 4 and 2147483648") int8_block = int(getattr(args, "rl_sync_push_int8_block", 128)) if not 32 <= int8_block <= 1 << 20: problems.append("--rl-sync-push-int8-block between 32 and 1048576") diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 1f45104b4..7f3f17622 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -118,6 +118,8 @@ def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: sweep += " --sweep-overlap" if getattr(args, "rl_sync_push_no_tensor_blocks", False): sweep += " --sweep-push-no-tensor-blocks" + if getattr(args, "rl_sync_fragment_max_bytes", 2 << 30) != 2 << 30: + sweep += f" --sweep-fragment-max-bytes {int(args.rl_sync_fragment_max_bytes)}" if getattr(args, "rl_sync_push_dtype", "f32") == "int8": sweep += f" --sweep-push-int8-block {int(getattr(args, 'rl_sync_push_int8_block', 128))}" if getattr(args, "rl_sync_push_dtype", "f32") == "int4": diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index f178541e3..25a9bcc53 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -921,6 +921,7 @@ def _dense_full_sync_config( wan_stall_timeout=float(getattr(args, "wan_stall_timeout", 0.0) or 0.0), push_dtype=getattr(args, "sweep_push_dtype", "f32"), push_int8_block=int(getattr(args, "sweep_push_int8_block", 128)), + max_fragment_bytes=int(getattr(args, "sweep_fragment_max_bytes", 2 << 30)), push_int4_block=int(getattr(args, "sweep_push_int4_block", 128)), push_tensor_blocks=not getattr(args, "sweep_push_no_tensor_blocks", False), push_int4_stochastic=bool(getattr(args, "sweep_push_int4_stochastic", False)), diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index ad4c8e273..332c8e330 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -14,6 +14,7 @@ ParameterLayout, ParameterSpec, owner_affine_minimum_fragments, + split_parameter_specs, ) from yeto.rl.moe_layout import check_expert_partition, describe_topology from yeto.rl.trajectory_evidence import TrajectoryBatchEvidence @@ -325,8 +326,12 @@ def create( # its own experts); a dense model passes through unchanged. check_expert_partition(specs) start = max(minimum_fragments, owner_affine_minimum_fragments(specs, matrix_merge)) + # A parameter above the bound spans several fragments as dimension-0 + # row slices; with no such parameter the layout is unchanged. + max_fragment_numel = max_fragment_bytes // 4 + units = split_parameter_specs(specs, max_fragment_numel) layout = None - for count in range(start, len(specs) + 1): + for count in range(start, len(units) + 1): candidate = ParameterLayout.create( algorithm=algorithm, components=components, @@ -335,6 +340,7 @@ def create( fragment_strategy="owner_affine", stream_role=stream_role, matrix_merge=matrix_merge, + max_fragment_numel=max_fragment_numel, ) if all( fragment.numel * 4 <= max_fragment_bytes @@ -343,9 +349,7 @@ def create( layout = candidate break if layout is None: - raise ValueError( - "a single topology-owned parameter exceeds the fragment byte bound" - ) + raise RuntimeError("owner-affine layout did not fit the fragment byte bound") manifests_by_owner = {} for manifest in ordered: @@ -365,11 +369,19 @@ def create( for fragment_id, fragment in enumerate(layout.fragments.fragments): if layout.fragment_owner(fragment_id) != owner: continue + slices = layout.fragment_slices(fragment_id) fragments.append( FullParameterFragmentPlan( fragment_id=fragment_id, - wire_names=tuple(name for name, _ in fragment.tensors), + wire_names=tuple(unit.spec.wire_name for unit in slices), numel=fragment.numel, + # Only a fragment holding part of a parameter names + # ranges, so unsplit plans keep their plan hash. + flat_ranges=( + () + if all(unit.is_whole for unit in slices) + else tuple(unit.flat_range for unit in slices) + ), ) ) plans.append( diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index febb0d818..fe72049f4 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -342,6 +342,13 @@ default="f32", help="full-parameter sweep: wire dtype of pushed outer deltas (syncer masters stay f32)", ), + island_flag( + "--sweep-fragment-max-bytes", + type=int, + default=2 << 30, + help="full-parameter: f32 byte bound of one sync fragment; a larger " + "parameter spans several fragments as dimension-0 row slices", + ), island_flag( "--sweep-push-int8-block", type=int, diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index dad8e9859..0e18ff9d0 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -949,6 +949,15 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "by one outer step; rollouts stay on-policy, so --rl-max-policy-age stays 0. " "Needs --rl-sync-fragment-pipeline equal to --fragments", ), + cli_flag( + "--rl-sync-fragment-max-bytes", + type=int, + default=2 << 30, + help="full-parameter only: f32 byte bound of one sync fragment (multiple of 4, " + "4..2147483648). A parameter above the bound spans several fragments as " + "dimension-0 row slices, each merged as its own tensor (RDA per slice). " + "A smaller bound changes the layout hash, so the probe must use the same value", + ), cli_flag( "--rl-sync-push-int8-block", type=int, diff --git a/yeto/rl/local_learner.py b/yeto/rl/local_learner.py index bffa6248e..0e5110cf8 100644 --- a/yeto/rl/local_learner.py +++ b/yeto/rl/local_learner.py @@ -9,6 +9,7 @@ from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass, field from fractions import Fraction +from itertools import pairwise from typing import Protocol import torch @@ -124,6 +125,113 @@ def wire_name(self) -> str: return f"{self.role}::{self.shard_id}::{self.name}" +@dataclass(frozen=True, order=True) +class ParameterSlice: + """Rows ``[row_start, row_stop)`` of one parameter along dimension 0. + + A layout splits a parameter into slices only when the parameter alone is + larger than the fragment bound; every other parameter is one whole slice. + """ + + spec: ParameterSpec + row_start: int + row_stop: int + + def __post_init__(self) -> None: + if not 0 <= self.row_start < self.row_stop <= self.spec.shape[0]: + raise ValueError(f"invalid row slice of {self.spec.wire_name!r}") + + @property + def is_whole(self) -> bool: + return self.row_start == 0 and self.row_stop == self.spec.shape[0] + + @property + def row_numel(self) -> int: + return self.spec.numel // self.spec.shape[0] + + @property + def numel(self) -> int: + return (self.row_stop - self.row_start) * self.row_numel + + @property + def shape(self) -> tuple[int, ...]: + return (self.row_stop - self.row_start, *self.spec.shape[1:]) + + @property + def flat_range(self) -> tuple[int, int]: + return self.row_start * self.row_numel, self.row_stop * self.row_numel + + @property + def tensor_name(self) -> str: + """The fragment tensor name; a whole parameter keeps its wire name.""" + + if self.is_whole: + return self.spec.wire_name + return f"{self.spec.wire_name}#rows={self.row_start}:{self.row_stop}" + + # The owner-affine grouping reads these through the spec. + @property + def role(self) -> str: + return self.spec.role + + @property + def shard_id(self) -> str: + return self.spec.shard_id + + @property + def name(self) -> str: + return self.spec.name + + +def split_parameter_specs( + specs: Sequence[ParameterSpec], + max_numel: int | None, +) -> tuple[ParameterSlice, ...]: + """Split every parameter above ``max_numel`` into balanced row slices. + + The slice count is the smallest one whose largest slice fits the bound; + row boundaries are ``rows * i // count``. ``None`` splits nothing. + """ + + slices = [] + for spec in sorted(specs): + rows = spec.shape[0] + if max_numel is None or spec.numel <= max_numel: + slices.append(ParameterSlice(spec, 0, rows)) + continue + row_numel = spec.numel // rows + if row_numel > max_numel: + raise ValueError( + f"one dimension-0 row of {spec.wire_name!r} exceeds the fragment bound" + ) + count = -(-spec.numel // max_numel) + while -(-rows // count) * row_numel > max_numel: + count += 1 + bounds = [rows * index // count for index in range(count + 1)] + slices.extend( + ParameterSlice(spec, start, stop) + for start, stop in pairwise(bounds) + ) + return tuple(slices) + + +def _join_adjacent_slices(members: Sequence[ParameterSlice]) -> list[ParameterSlice]: + """Join neighbouring slices of one parameter that share a fragment.""" + + joined: list[ParameterSlice] = [] + for member in members: + previous = joined[-1] if joined else None + if ( + previous is not None + and previous.spec == member.spec + and previous.row_stop == member.row_start + ): + joined[-1] = ParameterSlice(member.spec, previous.row_start, member.row_stop) + else: + joined.append(member) + return joined + + @dataclass(frozen=True) class ParameterLayout: """Deterministic role-qualified layout shared by dense and PULSE exchange.""" @@ -139,6 +247,11 @@ class ParameterLayout: compare=False, ) stream_role: str | None = None + _slices_by_tensor_name: Mapping[str, ParameterSlice] = field( + default_factory=dict, + repr=False, + compare=False, + ) @classmethod def create( @@ -151,11 +264,19 @@ def create( fragment_strategy: str = "balanced", stream_role: str | None = None, matrix_merge: str = "avg", + max_fragment_numel: int | None = None, ) -> ParameterLayout: """Lay ``specs`` out as ``num_fragments`` fragments. ``matrix_merge`` other than "avg" requires the owner-affine strategy: its fragments never mix merge modes, so each fragment carries one. + + ``max_fragment_numel`` (owner-affine only) splits a parameter larger + than the bound into dimension-0 row slices (see + ``split_parameter_specs``). A slice is its own fragment tensor, named + ``#rows=:``: the syncer merges it as one + tensor, so an RDA parameter that is split gets RDA per row slice. + A layout that splits nothing keeps its hash. """ if algorithm not in _ROLES_BY_ALGORITHM: raise ValueError(f"unsupported local RL algorithm: {algorithm!r}") @@ -193,29 +314,45 @@ def create( ) if matrix_merge != "avg" and fragment_strategy != "owner_affine": raise ValueError("a non-average matrix merge requires owner-affine fragments") + if max_fragment_numel is not None and ( + fragment_strategy != "owner_affine" + or type(max_fragment_numel) is not int + or max_fragment_numel < 1 + ): + raise ValueError("parameter splitting needs an owner-affine positive bound") + units = split_parameter_specs(ordered_specs, max_fragment_numel) if fragment_strategy == "balanced": minimum_fragments = 1 else: minimum_fragments = owner_affine_minimum_fragments(ordered_specs, matrix_merge) - if num_fragments < minimum_fragments or num_fragments > len(ordered_specs): + if num_fragments < minimum_fragments or num_fragments > len(units): raise ValueError("parameter fragment count is outside the layout") if fragment_strategy == "balanced": - bins = _binpack_parameter_specs(ordered_specs, num_fragments) + bins = [ + [ParameterSlice(spec, 0, spec.shape[0]) for spec in members] + for members in _binpack_parameter_specs(ordered_specs, num_fragments) + ] else: - bins = _owner_affine_parameter_bins(ordered_specs, num_fragments, matrix_merge) + bins = [ + _join_adjacent_slices(members) + for members in _owner_affine_parameter_bins(units, num_fragments, matrix_merge) + ] fragments = FragmentLayout( [ Fragment( # Owner-affine bins never mix merge modes (see # _merge_groups), so the first member decides. merge_mode=parameter_merge_mode(members[0].name, matrix_merge), - tensors=[(spec.wire_name, spec.numel) for spec in members], - identity_shapes={spec.wire_name: spec.shape for spec in members}, + tensors=[(unit.tensor_name, unit.numel) for unit in members], + identity_shapes={unit.tensor_name: unit.shape for unit in members}, ) for members in bins ] ) + slices_by_tensor_name = { + unit.tensor_name: unit for members in bins for unit in members + } payload = { "schema": 1 if fragment_strategy == "balanced" else 2, "algorithm": algorithm, @@ -263,16 +400,32 @@ def create( fragment_strategy=fragment_strategy, stream_role=stream_role, _specs_by_wire_name={spec.wire_name: spec for spec in ordered_specs}, + _slices_by_tensor_name=slices_by_tensor_name, ) - def fragment_specs(self, fragment_id: int) -> tuple[ParameterSpec, ...]: + @property + def has_split_parameters(self) -> bool: + return any(not unit.is_whole for unit in self._slices_by_tensor_name.values()) + + def fragment_slices(self, fragment_id: int) -> tuple[ParameterSlice, ...]: + """The row slices of one fragment, in fragment tensor order.""" + if not 0 <= fragment_id < self.fragments.num_fragments: raise ValueError("parameter fragment ID is outside the layout") return tuple( - self._specs_by_wire_name[name] + self._slices_by_tensor_name[name] for name, _ in self.fragments.fragments[fragment_id].tensors ) + def fragment_specs(self, fragment_id: int) -> tuple[ParameterSpec, ...]: + """The whole parameters of one fragment's tensors. + + A split parameter appears in every fragment that holds a slice of it; + callers that need the element range use ``fragment_slices``. + """ + + return tuple(unit.spec for unit in self.fragment_slices(fragment_id)) + def fragment_owner(self, fragment_id: int) -> tuple[str, str]: """Return the unique role/topology owner of an owner-affine fragment.""" @@ -308,9 +461,9 @@ def _binpack_parameter_specs( def _merge_groups( - specs: Sequence[ParameterSpec], + specs: Sequence[ParameterSpec | ParameterSlice], matrix_merge: str, -) -> list[list[ParameterSpec]]: +) -> list[list[ParameterSpec | ParameterSlice]]: """Split specs into owner-affine groups that each have one merge mode. A group is a maximal run of one topology owner's canonical spec order with @@ -343,10 +496,10 @@ def owner_affine_minimum_fragments( def _owner_affine_parameter_bins( - specs: Sequence[ParameterSpec], + specs: Sequence[ParameterSpec | ParameterSlice], num_fragments: int, matrix_merge: str = "avg", -) -> list[list[ParameterSpec]]: +) -> list[list[ParameterSpec | ParameterSlice]]: """Allocate deterministic balanced bins without crossing merge groups.""" groups = _merge_groups(specs, matrix_merge) @@ -372,9 +525,9 @@ def _owner_affine_parameter_bins( def _contiguous_parameter_bins( - specs: Sequence[ParameterSpec], + specs: Sequence[ParameterSpec | ParameterSlice], num_fragments: int, -) -> list[list[ParameterSpec]]: +) -> list[list[ParameterSpec | ParameterSlice]]: """Split canonical owner specs into balanced contiguous fragments. Owner ranks stream and apply their fragments independently. Keeping the From 391769c940587ac296f5c0efd88124883afd7d39 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:20:49 +0000 Subject: [PATCH 150/233] Miles plan: pass flat_ranges only for split fragments; refuse a split on an older Miles Co-Authored-By: Claude Opus 5.5 --- .../miles/models/chunked_full_parameter.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index 332c8e330..2995159b6 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -351,6 +351,14 @@ def create( if layout is None: raise RuntimeError("owner-affine layout did not fit the fragment byte bound") + if layout.has_split_parameters and "flat_ranges" not in getattr( + FullParameterFragmentPlan, "__dataclass_fields__", {} + ): + raise RuntimeError( + "a parameter exceeds the fragment byte bound and this Miles " + "cannot split it (needs FullParameterFragmentPlan.flat_ranges, " + "Miles s27-fragbound-split)" + ) manifests_by_owner = {} for manifest in ordered: owners = {(spec.role, spec.shard_id) for spec in manifest.specs} @@ -370,18 +378,20 @@ def create( if layout.fragment_owner(fragment_id) != owner: continue slices = layout.fragment_slices(fragment_id) + # Only a fragment holding part of a parameter names flat + # ranges (Miles s27-fragbound-split); unsplit plans keep + # their plan hash and install on older Miles images. + ranges = ( + {} + if all(unit.is_whole for unit in slices) + else {"flat_ranges": tuple(unit.flat_range for unit in slices)} + ) fragments.append( FullParameterFragmentPlan( fragment_id=fragment_id, wire_names=tuple(unit.spec.wire_name for unit in slices), numel=fragment.numel, - # Only a fragment holding part of a parameter names - # ranges, so unsplit plans keep their plan hash. - flat_ranges=( - () - if all(unit.is_whole for unit in slices) - else tuple(unit.flat_range for unit in slices) - ), + **ranges, ) ) plans.append( From 70c3dcdcd405921bcf08c8e6b61045b586bd8f5e Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:30:45 +0000 Subject: [PATCH 151/233] Miles cut restore: rebase full-parameter scheduler anchor after the load G-D2 rec-c: after a peer snapshot restore (digests equal) the island's first apply failed in Miles with 'scheduler progress has no local-step receipt', because the restore loads the donor's LR scheduler progress without a local step. Both restore paths now call Miles' rebase_full_parameter_scheduler_anchor (Miles 0b64b40c1) when full-parameter tracking is installed. The test fake now refuses the apply like Miles does without the rebase, and a regression test covers it. Co-Authored-By: Claude Opus 5.5 --- tests/test_peer_recovery_miles.py | 54 +++++++++++++++++++++++++++- yeto/rl/adapters/miles/cut_plugin.py | 17 +++++++++ 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/tests/test_peer_recovery_miles.py b/tests/test_peer_recovery_miles.py index 5b8d2a0aa..bb9a92dd2 100644 --- a/tests/test_peer_recovery_miles.py +++ b/tests/test_peer_recovery_miles.py @@ -12,7 +12,9 @@ import asyncio import hashlib import json +import sys import threading +import types from pathlib import Path from types import SimpleNamespace @@ -50,7 +52,14 @@ class FakeGroup: def __init__(self, *, tp=2, dp=1, adam_step=0, missing_rank=None): self.coords = [{"tp": r % tp, "pp": 0, "dp": r // tp, "tp_size": tp, "pp_size": 1, "cp_size": 1, "ep_size": 1, "dp_size": dp} for r in range(tp * dp)] - self.states = [{"adam_step": adam_step, "m": [0.25 * r] * 3} for r in range(tp * dp)] + self.states = [{"adam_step": adam_step, "scheduler_steps": adam_step, "m": [0.25 * r] * 3} + for r in range(tp * dp)] + # Miles' full-parameter tracking on the trainer (one actor stands for every rank). + self.actor = SimpleNamespace( + opt_param_scheduler=SimpleNamespace(num_steps=adam_step), + _external_full_policy_version=0, _external_full_local_step_generation=0, + _external_full_exported_local_step_generation=0, _external_full_scheduler_anchor_steps=adam_step, + _external_full_last_local_rollout_id=None, _external_full_last_local_step_receipt=None) self.missing_rank = missing_rank self.installed = 0 self.generation = 0 @@ -78,6 +87,9 @@ async def run_plugin(self, name, kwargs): data = (Path(kwargs["directory"]) / entry["path"]).read_bytes() assert hashlib.sha256(data).hexdigest() == entry["sha256"] self.states[rank] = json.loads(data) + # The real restore loads the donor's LR scheduler, then rebases Miles' anchor. + self.actor.opt_param_scheduler.num_steps = self.states[0]["scheduler_steps"] + cut_plugin.rebase_full_parameter_tracking(self.actor) return [{"ok": True} for _ in self.coords] raise AssertionError(name) @@ -109,6 +121,9 @@ def assemble_mixed_target(self, current, *, target_policy_version, authoritative async def apply(self, group, target, *, commit_token): assert commit_token + actor = group.actor # Miles prepare_full_parameter_chunked_shard checks scheduler congruence + if actor.opt_param_scheduler.num_steps != actor._external_full_scheduler_anchor_steps: + raise RuntimeError("full-parameter scheduler progress has no local-step receipt") group.applied.append(dict(target.fragments)) group.installed = target.policy_version group.generation = 0 @@ -190,6 +205,26 @@ def config(store, *, recover, timeout_s=10.0): peer_recovery=PeerRecoveryOptions(store=str(store), recover=recover, timeout_s=timeout_s)) +def _install_miles_rebase_stub(monkeypatch): + """Without Miles, stand in for its scheduler-anchor rebase (same rule, no checks).""" + try: + import miles.backends.megatron_utils.full_parameter_state # noqa: F401 + return + except ImportError: + pass + + def rebase(actor, *, reason): + assert reason + actor._external_full_scheduler_anchor_steps = actor.opt_param_scheduler.num_steps + return actor._external_full_scheduler_anchor_steps + + for name in ("miles", "miles.backends", "miles.backends.megatron_utils"): + monkeypatch.setitem(sys.modules, name, types.ModuleType(name)) + module = types.ModuleType("miles.backends.megatron_utils.full_parameter_state") + module.rebase_full_parameter_scheduler_anchor = rebase + monkeypatch.setitem(sys.modules, module.__name__, module) + + @pytest.fixture def stubs(monkeypatch): clients = [] @@ -203,6 +238,7 @@ async def initialize(cls, group, **_kwargs): return FakeAdapter(), SimpleNamespace(policy_version=0, policy_hash="p0") monkeypatch.setattr(MilesChunkedFullParameterAdapter, "initialize", classmethod(initialize)) + _install_miles_rebase_stub(monkeypatch) monkeypatch.setattr(role_stream_module, "SyncerClient", make_client) monkeypatch.setattr(role_stream_module, "MilesRayInboundChunkSink", lambda sink: sink) monkeypatch.setattr(role_stream_module, "role_stream_session_contract_hash", lambda *_a: b"\0" * 32) @@ -294,3 +330,19 @@ def test_streams_without_peer_recovery_serve_nothing(tmp_path, stubs): stream = MilesFullParameterRoleStream(RoleStreamConfig(**{**config(tmp_path, recover=False).__dict__, "peer_recovery": None})) assert asyncio.run(stream.serve_peer_snapshots()) == 0 + + +def test_restore_without_scheduler_rebase_fails_the_first_apply(tmp_path, stubs, monkeypatch): + """G-D2 rec-c: restore OK and digests equal, but Miles refused the first apply.""" + store_root = tmp_path / "store" + donor_client = FakeClient(requests=[rp.SnapshotRequest(1, 9, 5, (3, 4))]) + asyncio.run(donor_stream(store_root, FakeGroup(adam_step=17), donor_client).serve_peer_snapshots()) + (_, versions, sha256, locator), = donor_client.sent + resent = [BcastFragment(0, 5, stored(0, 5, "g0"), stored=True), + BcastFragment(1, 6, stored(1, 6, "g1"), stored=True)] + stubs.clients.append(FakeClient(events=[rp.RecoveryClock(5, (3, 4), 1), + rp.RecoverySource(rp.SOURCE_PEER, 0, versions, sha256, locator), + rp.RecoveryDone(5, 6, 1, (5, 6))], updates=resent)) + monkeypatch.setattr(cut_plugin, "rebase_full_parameter_tracking", lambda actor: None) + with pytest.raises(Exception, match="no local-step receipt"): + asyncio.run(MilesFullParameterRoleStream(config(store_root, recover=True)).initialize(FakeGroup())) diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 70bd951f5..cd23b11a4 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -1362,6 +1362,21 @@ def rank_coords(actor: Any) -> dict[str, int]: RANK_COORDS = f"{_MODULE}.rank_coords" +def rebase_full_parameter_tracking(actor: Any) -> None: + """After a restore loads the donor's LR scheduler, move Miles' full-parameter + scheduler anchor to it. Without this the island's first global apply fails + with "scheduler progress has no local-step receipt" (G-D2 rec-c): the + restored progress has no local step behind it on this trainer. A trainer + without full-parameter tracking (LoRA cut path) is left alone.""" + if getattr(actor, "_external_full_policy_version", None) is None: + return + try: + from miles.backends.megatron_utils.full_parameter_state import rebase_full_parameter_scheduler_anchor + except ImportError as exc: + raise CutPluginError("this Miles cannot rebase full-parameter scheduler tracking after a restore") from exc + rebase_full_parameter_scheduler_anchor(actor, reason="peer snapshot restore") + + def restore_cut_shard(actor: Any, *, directory: str, files: list[Mapping[str, Any]], cut_id: str) -> dict[str, Any]: """Verify and load this rank's shard; return the post-restore summary (identical digests expected).""" install_grad_norm_recorder() # a rebuilt trainer process has no recorder yet @@ -1462,6 +1477,7 @@ def _restore(actor: Any, *, directory: str, files: list[Mapping[str, Any]], cut_ actor.opt_param_scheduler.load_state_dict(shard["scheduler"]) if int(actor.opt_param_scheduler.num_steps) != _scheduler_steps(shard["scheduler"]): raise CutPluginError("LR scheduler progress after restore differs from the cut") + rebase_full_parameter_tracking(actor) backend.set_megatron_counters(shard.get("megatron_counters") or {}) set_miles_counters(actor, shard.get("miles_counters") or {}) refresh_actor_backup(actor) # colocated update_weights may read the CPU backup @@ -1611,6 +1627,7 @@ def _restore_resharded(actor, *, directory, files, cut_id, source_dp, rng_policy actor.opt_param_scheduler.load_state_dict(first["scheduler"]) if int(actor.opt_param_scheduler.num_steps) != _scheduler_steps(first["scheduler"]): raise CutPluginError("LR scheduler progress after restore differs from the cut") + rebase_full_parameter_tracking(actor) backend.set_megatron_counters(first.get("megatron_counters") or {}) set_miles_counters(actor, first.get("miles_counters") or {}) # same counter on the resharded (E3) path refresh_actor_backup(actor) From fb76c721afe3dda8ef68ebaefb290452f4bf7b0c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:35:04 +0000 Subject: [PATCH 152/233] Ports image: Miles 0b64b40c1 (flat-range fragment split, scheduler anchor rebase) Re-pin MILES_NEXT_COMMIT to 0b64b40c1 and append a477bcbe0 and 0b64b40c1 to MILES_FORK_LINEAGE. 0b64b40c1 = cc693c693 + FullParameterFragmentPlan flat ranges (a parameter above the fragment bound spans several fragments) + rebase_full_parameter_scheduler_anchor for peer-snapshot recovery, which the cut plugin now calls. Image ghcr.io/michaellchung/yeto-miles-ports:0b64b40-2fa8801 @sha256:2b8c36f6cbdeddcbc3d867d23808b94a12a20217ce9fd2eb55ba94aa2c12f7f0 (the interim a477bcb-2fa8801 image is void and not pinned). Co-Authored-By: Claude Opus 5.5 --- docker/miles-ports/Dockerfile | 2 +- .../build-record.json | 12 +++++ .../image-manifest.json | 50 +++++++++++++++++++ tests/test_rl_ports_image.py | 2 +- yeto/rl/adapters/miles/pins.py | 16 ++++-- 5 files changed, 76 insertions(+), 6 deletions(-) create mode 100644 openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/build-record.json create mode 100644 openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/image-manifest.json diff --git a/docker/miles-ports/Dockerfile b/docker/miles-ports/Dockerfile index 2f4e7bf56..11f4c7d44 100644 --- a/docker/miles-ports/Dockerfile +++ b/docker/miles-ports/Dockerfile @@ -19,7 +19,7 @@ FROM docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c ARG MILES_REPOSITORY=https://github.com/michaellchung/miles -ARG MILES_COMMIT=cc693c6930be178f411f711f55c3d334a087c097 +ARG MILES_COMMIT=0b64b40c15fa972d20bd1393aca6067dfd2492a4 ARG SGLANG_REPOSITORY=https://github.com/michaellchung/sglang ARG SGLANG_COMMIT=2fa880182eefdb31e64d8ea70317be27f9f32121 # setuptools-scm version of SGLANG_COMMIT (describe v0.5.20-69-g4e4148f1b); diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/build-record.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/build-record.json new file mode 100644 index 000000000..cd7991919 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/build-record.json @@ -0,0 +1,12 @@ +{ + "tag": "ghcr.io/michaellchung/yeto-miles-ports:0b64b40-2fa8801", + "digest": "sha256:2b8c36f6cbdeddcbc3d867d23808b94a12a20217ce9fd2eb55ba94aa2c12f7f0", + "layer_tar_sha256": "e5ef26715d37a37b44c4816b9769fab3959c3b1a810941be6eb74f8c46b34393", + "base": "docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "base_index": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "miles_commit": "0b64b40c15fa972d20bd1393aca6067dfd2492a4", + "sglang_commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "sglang_version": "0.5.21.dev70+g2fa8801", + "peft_version": "0.20.0", + "peft_wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/image-manifest.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/image-manifest.json new file mode 100644 index 000000000..e2d7ea08e --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/image-manifest.json @@ -0,0 +1,50 @@ +{ + "base": { + "index_digest": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "manifest_digest": "sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "miles_commit": "9e4260de047a704208535c0e90c531929879ab40", + "platform": "linux/amd64", + "repository": "docker.io/radixark/miles", + "sglang_commit": "e5bba1f9c433da5352e521a8bb826bab960e7836" + }, + "builder": "scripts/build_miles_ports_image.sh (crane append); docker/miles-ports/Dockerfile equivalent", + "image": "ghcr.io/michaellchung/yeto-miles-ports", + "megatron_bridge": { + "commit": "8cd3466d14d2337c8492827b3712482c2b3e4866", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_bridge_init.py)" + ] + }, + "megatron_lm": { + "commit": "a84b105473eca9c9c6e49dbf1c8aa460c5cf796e", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_lm_lazy_modelopt.patch)" + ], + "path": "/root/Megatron-LM", + "repository": "https://github.com/radixark/Megatron-LM" + }, + "miles": { + "commit": "0b64b40c15fa972d20bd1393aca6067dfd2492a4", + "install": "base's editable install (pip install -e /root/miles --no-deps) kept (same path); /root is on the base's sys.path, so without PYTHONPATH import miles is a namespace package over the /root/miles checkout -- yeto runs with PYTHONPATH=$HOME/miles (regular package)", + "path": "/root/miles", + "repository": "https://github.com/michaellchung/miles", + "upstream_commit": "9e4260de047a704208535c0e90c531929879ab40" + }, + "peft": { + "install": "wheel unpacked into the base's site-packages (pip install --no-deps equivalent)", + "version": "0.20.0", + "wheel": "peft-0.20.0-py3-none-any.whl", + "wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" + }, + "pure_python_overlay": true, + "schema": 1, + "sglang": { + "commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "install": "pip install -e /sgl-workspace/sglang/python[all] --no-deps (base's editable install, same path)", + "path": "/sgl-workspace/sglang", + "repository": "https://github.com/michaellchung/sglang", + "upstream_commit": "571212b636baca45e10fa3b4da11a289123f3235", + "version": "0.5.21.dev70+g2fa8801" + }, + "tag": "0b64b40-2fa8801" +} \ No newline at end of file diff --git a/tests/test_rl_ports_image.py b/tests/test_rl_ports_image.py index f24a7b38e..f0e38dd77 100644 --- a/tests/test_rl_ports_image.py +++ b/tests/test_rl_ports_image.py @@ -70,7 +70,7 @@ def test_build_inputs_match_the_pins(): def test_commits_and_digest_are_pinned_together_with_the_build_record(): """Commit pins, tag, Dockerfile ARGs and the image digest move together: the latest build record must name exactly the pinned commits and digest.""" - record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801" + record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801" record = json.loads((record_dir / "build-record.json").read_text()) assert record["miles_commit"] == rl.MILES_NEXT_COMMIT assert record["sglang_commit"] == rl.SGLANG_NEXT_COMMIT diff --git a/yeto/rl/adapters/miles/pins.py b/yeto/rl/adapters/miles/pins.py index 8d379c3a8..f46a0df99 100644 --- a/yeto/rl/adapters/miles/pins.py +++ b/yeto/rl/adapters/miles/pins.py @@ -55,8 +55,10 @@ class MilesRevisionPins(NamedTuple): # stay off megatron_utils (L2.1) + the master port is held until use (S23 defect 9). # s27-distopt-fp (56190a5a8 = ed5f64fd7 + full-parameter sync reads DP=1 # distributed-optimizer masters, incl. CPU offload; S27 G-B); cc693c693 = + the -# Adam proof leaves the hybrid optimizer's master_param mirror out (S27 IMG2). -MILES_NEXT_COMMIT = "cc693c6930be178f411f711f55c3d334a087c097" +# Adam proof leaves the hybrid optimizer's master_param mirror out (S27 IMG2); +# a477bcbe0 = + one parameter may span several fragments via flat ranges (S27); +# 0b64b40c1 = + rebase_full_parameter_scheduler_anchor for peer-snapshot recovery. +MILES_NEXT_COMMIT = "0b64b40c15fa972d20bd1393aca6067dfd2492a4" # sgl-project/sglang ``sglang-miles`` head when radixark/miles@9e4260d was # committed (upstream's Dockerfile follows that branch unpinned). SGLANG_NEXT_REPOSITORY = "https://github.com/michaellchung/sglang" @@ -121,6 +123,12 @@ class MilesRevisionPins(NamedTuple): # cc693c693 = + the optimizer-state proof skips the HybridDeviceOptimizer's # master_param mirror, which a global apply overwrites (S27 IMG2); no loss path. "cc693c6930be178f411f711f55c3d334a087c097", + # a477bcbe0 = + FullParameterFragmentPlan.flat_ranges: a parameter above the + # fragment bound spans several fragments (S27); plans without ranges keep their hash. + "a477bcbe04bc6278bfed017cf8931ae90bb8e3e3", + # 0b64b40c1 = + rebase_full_parameter_scheduler_anchor after a peer snapshot + # restore (S27); no loss path. + "0b64b40c15fa972d20bd1393aca6067dfd2492a4", ) # First Miles commit whose full-parameter sync reads distributed-optimizer @@ -165,7 +173,7 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # every SHA. PUBLIC on ghcr.io/michaellchung since 2026-10-07 (anonymous # pull); a private image needs SKYPILOT_DOCKER_{USERNAME,PASSWORD,SERVER} or # --rl-image-private (yeto.launcher.registry_login_for; read:packages token). -# Tag cc693c6-2fa8801; (S27 IMG2: Miles cc693c693 = 56190a5a8 (full-parameter sync with distributed optimizer + CPU offload, G-B) + Adam proof without the hybrid master mirror; previous 56190a5-2fa8801 @sha256:5254e34e... (failed on GPU: global apply changed local Adam state), ed5f64f-2fa8801 @sha256:a4769254...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/). Earlier: tag ed5f64f-2fa8801; (S27 IMG: Miles ed5f64fd7, host-Adam round cuts (B8); previous 44d41ef-2fa8801 @sha256:5b9dc2a9...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-ed5f64f-2fa8801/). Earlier: tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt +# Tag 0b64b40-2fa8801; (S27 IMG3: Miles 0b64b40c1 = cc693c693 + flat-range fragment split + scheduler anchor rebase; previous cc693c6-2fa8801 @sha256:2011cbf9... (a477bcb-2fa8801 built but voided); build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/). Earlier: tag cc693c6-2fa8801; (S27 IMG2: Miles cc693c693 = 56190a5a8 (full-parameter sync with distributed optimizer + CPU offload, G-B) + Adam proof without the hybrid master mirror; previous 56190a5-2fa8801 @sha256:5254e34e... (failed on GPU: global apply changed local Adam state), ed5f64f-2fa8801 @sha256:a4769254...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/). Earlier: tag ed5f64f-2fa8801; (S27 IMG: Miles ed5f64fd7, host-Adam round cuts (B8); previous 44d41ef-2fa8801 @sha256:5b9dc2a9...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-ed5f64f-2fa8801/). Earlier: tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt # Megatron-LM patch in the layer; previous 17aeb5c-2fa8801 @sha256:b304f857...; # build record openspec/changes/rl-infra-spec/evidence/ports-image/ # 2026-10-14-44d41ef-2fa8801/). Earlier: tag 17aeb5c-2fa8801; (S24 L2.1-L2.4: Miles s24-miles-import, peft 0.20.0 and the @@ -174,7 +182,7 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # ports-image/2026-10-14-17aeb5c-2fa8801/). Earlier: tag 386ae3b-2fa8801; (S23 P0-1.11 GPU run: Miles s23-sao-streaming-roles; previous 64b591a-2fa8801 @sha256:fa2413be...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-386ae3b-2fa8801/). Earlier: (S19 #1: critic family + algo-supplement metrics baked in; previous ddce209-2fa8801 @sha256:9c252c38...) (S18 ARU-3: Miles s18-agentic-suspend = 2f7871fb2 + --agentic-suspend-between-turns; previous 2f7871f-2fa8801 @sha256:62b4f164..., efbbc63-2fa8801 @sha256:a7990076..., 8bc5223-4e4148f @sha256:4aeafd77...) ghcr.io/michaellchung (public package); linux/amd64 only. MILES_NEXT_IMAGE = ( "docker:ghcr.io/michaellchung/yeto-miles-ports@sha256:" - "2011cbf9bf6e2f472a0819d77685d6cec56f4941957ae5652f09f3ccb7604e08" + "2b8c36f6cbdeddcbc3d867d23808b94a12a20217ce9fd2eb55ba94aa2c12f7f0" ) MILES_NEXT_IMAGE_MANIFEST = "/opt/yeto/image-manifest.json" # Nebius VM images whose /var/lib/docker already holds a docker image's From b57103301c410e9023142003474b1f0768303c6b Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:03:48 +0000 Subject: [PATCH 153/233] Full-parameter local init: digest-checked version 0 without the initial broadcast Every island loads the same base checkpoint, so async full-parameter rounds no longer broadcast learner 0's INIT back to every island (X6: 116.5 GB per island at ~55 MB/s timed out the 900 s initial-cut wait). Islands send per-fragment SHA-256 digests of their INIT bytes after each HELLO (INIT_DIGEST); the syncer checks all islands agree and that the applied INIT matches, then sends INIT_CONFIRM instead of the version-0 broadcast. A mismatch stops the run as init_digest_mismatch with a clear error on every island. Learner 0 still uploads INIT, so syncer masters, merges, checkpoints and resume are unchanged; a resumed syncer, islands without digests and late joiners still get the broadcast. --rl-sync-init-mode auto|local|upload (auto = local for async rounds). The initial-cut wait now scales with model bytes moved at 32 MiB/s. Co-Authored-By: Claude Opus 5.5 --- docs/PROTOCOL.md | 22 ++ syncer/src/local_init.rs | 214 ++++++++++++++ syncer/src/main.rs | 1 + syncer/src/protocol.rs | 13 + syncer/src/server.rs | 110 ++++++- tests/test_rl_dense_sweep_real_syncer.py | 8 + tests/test_rl_local_init_real_syncer.py | 276 ++++++++++++++++++ yeto/launch/rl_args.py | 1 + yeto/launch/rl_island_task.py | 2 + yeto/protocol.py | 45 +++ .../miles/full_parameter/role_stream.py | 64 +++- yeto/rl/adapters/miles/island_entry.py | 9 + yeto/rl/adapters/miles/run_flags.py | 8 + yeto/rl/engine/run_flags.py | 11 + yeto/syncer_shards.py | 8 + 15 files changed, 783 insertions(+), 9 deletions(-) create mode 100644 syncer/src/local_init.rs create mode 100644 tests/test_rl_local_init_real_syncer.py diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md index 211e0b769..1d706b73a 100644 --- a/docs/PROTOCOL.md +++ b/docs/PROTOCOL.md @@ -633,6 +633,28 @@ It keeps `status.json` next to the tape (schema `yeto.syncer.elastic-status/v1`: state, syncer epoch, outer version, live islands with capacity and remaining lease). +### Local init (INIT_DIGEST / INIT_CONFIRM) + +When every island loads the same base checkpoint (`--rl-sync-init-mode +local`, the default with `--rl-sync-full-parameter-rounds async`), each island +sends INIT_DIGEST (type 28) on its control socket right after every HELLO +until its initial cut arrives: `u32 learner_id | [32] layout_hash | u32 +first_fragment | u32 count | count x [32] sha256`. Each digest is the SHA-256 +of the fragment's INIT_PARAMS tensor bytes; the list covers exactly the +fragments the receiving syncer (shard) owns. Only f32-bulk sessions accept it. + +Learner 0 still sends INIT_PARAMS, so the syncer's f32 masters, merges, +checkpoints and resume are unchanged. The syncer checks that every island's +layout hash and digests agree and that each applied INIT hashes to the agreed +digest. Any difference stops the run as the strict failure +`init_digest_mismatch`; every connected island first receives MSG_ERROR +`init digest mismatch: ...` and fails with the same metric. + +When the checks pass on a fresh start, an island that sent digests receives +INIT_CONFIRM (type 29: `[32]` SHA-256 over the concatenated digests) instead +of the version-0 broadcast and keeps its own weights as version 0. An island +without digests, a resumed syncer, and late joiners still get the broadcast. + ### Island departure (LEAVE) An island that leaves a pool sends LEAVE (type 27: `u32 learner_id | u8 diff --git a/syncer/src/local_init.rs b/syncer/src/local_init.rs new file mode 100644 index 000000000..c6c0884e9 --- /dev/null +++ b/syncer/src/local_init.rs @@ -0,0 +1,214 @@ +//! Local-init startup (S27, `MSG_INIT_DIGEST` / `MSG_INIT_CONFIRM`). +//! +//! Every island loads the same base checkpoint, so all of them already hold +//! the version-0 cut. Instead of learner 0's INIT followed by a full +//! broadcast to every island, each island reports the SHA-256 of every +//! fragment's INIT_PARAMS tensor bytes. The syncer still initializes its f32 +//! masters from learner 0's INIT (so merges, checkpoints and resume are +//! unchanged), and checks that +//! +//! * every island reports the same layout hash and fragment digests, and +//! * the INIT bytes it applied hash to those digests. +//! +//! When both hold, an island that reported digests receives one 32-byte +//! INIT_CONFIRM instead of the initial broadcast; an island that did not +//! report digests still gets the broadcast. Any disagreement is fatal: the +//! islands did not start from the same weights. + +use std::collections::BTreeSet; +use std::ops::Range; + +use anyhow::{bail, Result}; +use sha2::{Digest, Sha256}; + +use crate::protocol::Reader; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InitDigests { + pub learner_id: u32, + pub layout_hash: [u8; 32], + pub first: u32, + pub digests: Vec<[u8; 32]>, +} + +fn hex(h: &[u8; 32]) -> String { + h.iter().map(|b| format!("{b:02x}")).collect() +} + +fn read32(r: &mut Reader<'_>) -> Result<[u8; 32]> { + let mut out = [0u8; 32]; + for byte in &mut out { + *byte = r.u8()?; + } + Ok(out) +} + +impl InitDigests { + pub fn decode(payload: &[u8], connected_learner: u32) -> Result { + let mut r = Reader(payload); + let learner_id = r.u32()?; + if learner_id != connected_learner { + bail!("INIT_DIGEST learner id {learner_id} does not match connected learner {connected_learner}"); + } + let layout_hash = read32(&mut r)?; + let first = r.u32()?; + let count = r.u32()? as usize; + if count == 0 || r.remaining() != 32 * count { + bail!("INIT_DIGEST declares {count} digest(s) but carries {} byte(s)", r.remaining()); + } + let digests = (0..count).map(|_| read32(&mut r)).collect::>>()?; + Ok(Self { learner_id, layout_hash, first, digests }) + } + + pub fn range(&self) -> Range { + self.first as usize..self.first as usize + self.digests.len() + } + + /// INIT_CONFIRM payload: SHA-256 over the concatenated digests. + pub fn combined(&self) -> [u8; 32] { + let mut hasher = Sha256::new(); + for digest in &self.digests { + hasher.update(digest); + } + hasher.finalize().into() + } +} + +/// Startup bookkeeping of the local-init check (one per syncer shard). +pub struct LocalInit { + owned: Range, + agreed: Option, + holders: BTreeSet, + /// SHA-256 of each applied INIT (index = fragment id). + applied: Vec>, +} + +impl LocalInit { + pub fn new(owned: Range, fragments: usize) -> Self { + Self { owned, agreed: None, holders: BTreeSet::new(), applied: vec![None; fragments] } + } + + /// Record one island's digests; an error means the islands disagree. + pub fn on_digests(&mut self, digests: InitDigests) -> Result<()> { + if digests.range() != self.owned { + bail!( + "init digest mismatch: learner {} reported fragments {:?}, this syncer owns {:?}", + digests.learner_id, + digests.range(), + self.owned + ); + } + for (offset, digest) in digests.digests.iter().enumerate() { + let fid = self.owned.start + offset; + if let Some(applied) = &self.applied[fid] { + check(fid, digests.learner_id, digest, "the applied INIT_PARAMS", applied)?; + } + } + match &self.agreed { + None => self.agreed = Some(digests.clone()), + Some(agreed) => { + if agreed.layout_hash != digests.layout_hash { + bail!( + "init digest mismatch: learner {} layout hash {} differs from learner {} \ + layout hash {}; every island must build the same fragment layout", + digests.learner_id, + hex(&digests.layout_hash), + agreed.learner_id, + hex(&agreed.layout_hash) + ); + } + for (offset, (mine, theirs)) in digests.digests.iter().zip(&agreed.digests).enumerate() { + let other = format!("learner {}", agreed.learner_id); + check(self.owned.start + offset, digests.learner_id, mine, &other, theirs)?; + } + } + } + self.holders.insert(digests.learner_id); + Ok(()) + } + + /// Record the hash of an applied INIT; an error means it differs from + /// the digests the islands agreed on. + pub fn on_init(&mut self, fid: usize, learner_id: u32, sha256: [u8; 32]) -> Result<()> { + if let Some(agreed) = &self.agreed { + if let Some(expected) = agreed.digests.get(fid.wrapping_sub(self.owned.start)) { + let other = format!("the islands' digest (learner {})", agreed.learner_id); + check(fid, learner_id, &sha256, &other, expected)?; + } + } + if let Some(slot) = self.applied.get_mut(fid) { + slot.get_or_insert(sha256); + } + Ok(()) + } + + /// Whether `learner_id` reported digests that agree. + pub fn holds(&self, learner_id: u32) -> bool { + self.holders.contains(&learner_id) + } + + /// The INIT_CONFIRM payload, once the agreed digests cover every applied + /// INIT of this shard; `None` means the initial broadcast is needed. + pub fn confirm(&self) -> Option<[u8; 32]> { + let agreed = self.agreed.as_ref()?; + let complete = self.owned.clone().all(|fid| { + self.applied[fid].as_ref() == agreed.digests.get(fid - self.owned.start) + }); + complete.then(|| agreed.combined()) + } +} + +fn check(fid: usize, learner: u32, mine: &[u8; 32], other: &str, theirs: &[u8; 32]) -> Result<()> { + if mine != theirs { + bail!( + "init digest mismatch: fragment {fid} of learner {learner} has sha256 {} but {other} \ + has {}; every island must load the same base checkpoint (or start with \ + --rl-init-mode upload)", + hex(mine), + hex(theirs) + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn digests(learner_id: u32, values: &[u8]) -> InitDigests { + InitDigests { + learner_id, + layout_hash: [7; 32], + first: 0, + digests: values.iter().map(|v| [*v; 32]).collect(), + } + } + + #[test] + fn agreement_and_init_give_a_confirm() { + let mut init = LocalInit::new(0..2, 2); + init.on_digests(digests(0, &[1, 2])).unwrap(); + init.on_digests(digests(1, &[1, 2])).unwrap(); + assert!(init.confirm().is_none(), "no INIT applied yet"); + init.on_init(0, 0, [1; 32]).unwrap(); + init.on_init(1, 0, [2; 32]).unwrap(); + assert_eq!(init.confirm(), Some(digests(0, &[1, 2]).combined())); + assert!(init.holds(1) && !init.holds(2)); + } + + #[test] + fn mismatches_are_refused() { + let mut init = LocalInit::new(0..2, 2); + init.on_digests(digests(0, &[1, 2])).unwrap(); + let error = init.on_digests(digests(1, &[1, 3])).unwrap_err().to_string(); + assert!(error.contains("fragment 1 of learner 1"), "{error}"); + let mut other_layout = digests(1, &[1, 2]); + other_layout.layout_hash = [8; 32]; + assert!(init.on_digests(other_layout).unwrap_err().to_string().contains("layout hash")); + assert!(init.on_init(0, 0, [9; 32]).is_err()); + // An INIT applied before any digest is checked when digests arrive. + let mut late = LocalInit::new(0..2, 2); + late.on_init(0, 0, [5; 32]).unwrap(); + assert!(late.on_digests(digests(1, &[1, 2])).is_err()); + } +} diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 8dc059ac6..cdb18cb47 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -3,6 +3,7 @@ mod hmac; mod shard; #[cfg(test)] mod head_memory; +mod local_init; mod iso_worker; mod membership; mod recovery; diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index 0056a677f..c4f442152 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -29,6 +29,19 @@ pub const MSG_BUDGET_DONE: u8 = 14; /// Sent before the island closes its connection; the reason goes on the /// `pool_leave` tape record. pub const MSG_LEAVE: u8 = 27; +/// Local-init startup (S27): every island loaded the same base checkpoint and +/// reports the SHA-256 of each owned fragment's INIT_PARAMS tensor bytes: +/// `u32 learner_id | [32] layout_hash | u32 first_fragment | u32 count | +/// count x [32] digest`. Valid only on f32-bulk sessions. +pub const MSG_INIT_DIGEST: u8 = 28; +/// Syncer -> island: every island's digests and the applied INIT agree, so +/// the island adopts its local weights as version 0 and no initial BCAST +/// follows. Payload: `[32]` SHA-256 over the concatenated fragment digests. +pub const MSG_INIT_CONFIRM: u8 = 29; +/// Largest INIT_DIGEST payload for a layout of `fragments` fragments. +pub fn init_digest_payload_len(fragments: usize) -> u64 { + 4 + 32 + 4 + 4 + 32 * fragments as u64 +} pub const LEAVE_REASON_NAMES: [&str; 7] = [ "finished", "lease_expired", diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 31a093653..f8ea31afa 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -19,6 +19,7 @@ use tokio::net::{TcpListener, TcpStream}; use tokio::sync::mpsc; use tracing::{info, warn}; +use crate::local_init::{InitDigests, LocalInit}; use crate::recovery::{self, HelloExt, PeerRecovery}; use crate::membership::{Admission, AdmitError, CapacityQuorum, Member, Membership}; use crate::protocol::*; @@ -470,6 +471,13 @@ enum Event { member: Member, fragment_id: u32, values: Vec, + /// SHA-256 of the INIT_PARAMS tensor bytes as received. + sha256: [u8; 32], + }, + /// Local-init startup: the island's per-fragment INIT digests. + InitDigest { + member: Member, + digests: InitDigests, }, Push { member: Member, @@ -1929,6 +1937,7 @@ async fn read_loop( header_at = Some(Instant::now()); match msg_type { MSG_INIT_PARAMS => Ok(group.max_init_payload), + MSG_INIT_DIGEST => Ok(init_digest_payload_len(group.layout.fragments.len())), MSG_PUSH_FRAGMENT => Ok(group.max_push_payload), MSG_HEARTBEAT => Ok(12), MSG_LEAVE => Ok(5), @@ -2111,6 +2120,7 @@ async fn dispatch_inner( bytes.len() ); } + let sha256: [u8; 32] = Sha256::digest(bytes).into(); let mut values = Vec::new(); decode_tensor(bulk_dtype(group.dtype), bytes, &mut values)?; event_tx @@ -2118,6 +2128,28 @@ async fn dispatch_inner( member: group.member, fragment_id, values, + sha256, + }) + .await + .ok(); + } + MSG_INIT_DIGEST => { + if bulk_dtype(group.dtype) != DTYPE_F32 { + bail!("INIT_DIGEST needs an f32-bulk session (dtype {})", group.dtype); + } + let digests = InitDigests::decode(payload, group.member.learner_id)?; + if digests.first as usize + digests.digests.len() > group.layout.fragments.len() { + bail!( + "INIT_DIGEST covers fragments {}..{}, layout has {}", + digests.first, + digests.first as usize + digests.digests.len(), + group.layout.fragments.len() + ); + } + event_tx + .send(Event::InitDigest { + member: group.member, + digests, }) .await .ok(); @@ -2412,6 +2444,17 @@ async fn recv_init_event( event.context("event channel closed") } +/// The islands did not start from the same weights: record the strict +/// failure, tell every connected island why, and stop the syncer. +async fn refuse_local_init(cfg: &Config, registry: &Registry, error: anyhow::Error) -> Result<()> { + let message = format!("{error:#}"); + append_strict_failure(cfg.event_tape.as_deref(), "init_digest_mismatch", &message); + for group in current_groups(registry) { + let _ = group.send_small(MSG_ERROR, message.clone().into()).await; + } + bail!("RL strict failure init_digest_mismatch: {message}"); +} + /// Connected learners the scheduler waits for before the first round: every /// learner with static membership, the launch quorum with dynamic membership. fn startup_members(cfg: &Config) -> u32 { @@ -2438,6 +2481,9 @@ async fn scheduler( // LEAVE reasons of islands whose disconnect has not arrived yet. let mut leave_reasons: HashMap = HashMap::new(); let checkpoint_each_round = checkpoint_before_broadcast(&cfg); + // Local-init digests (`local_init`); created with the state. + let mut local_init: Option = None; + let mut resumed = false; // Phase 1: wait until every fragment is initialized (via INIT_PARAMS or // a resumed checkpoint) and the expected learners have connected (late @@ -2527,15 +2573,33 @@ async fn scheduler( } } info!(step = st.global_step, "resumed from checkpoint"); + resumed = true; } } + local_init = Some(LocalInit::new( + st.owned_fragments(), + st.layout.fragments.len(), + )); state = Some(st); } } + Event::InitDigest { member, digests } => { + if !is_current_member(®istry, member) { + continue; + } + let checked = local_init + .as_mut() + .context("INIT_DIGEST before HELLO")? + .on_digests(digests); + if let Err(error) = checked { + refuse_local_init(&cfg, ®istry, error).await?; + } + } Event::Init { member, fragment_id, values, + sha256, } => { if !is_current_member(®istry, member) { warn!( @@ -2546,6 +2610,16 @@ async fn scheduler( continue; } let st = state.as_mut().context("INIT before HELLO")?; + if !st.initialized[fragment_id as usize] { + let checked = local_init.as_mut().unwrap().on_init( + fragment_id as usize, + member.learner_id, + sha256, + ); + if let Err(error) = checked { + refuse_local_init(&cfg, ®istry, error).await?; + } + } st.init_fragment(fragment_id as usize, values)?; if st.all_initialized() { info!("global parameters initialized"); @@ -2640,8 +2714,32 @@ async fn scheduler( // learners start bit-identical (also serves recovery for late joiners). // A checkpoint already at the terminal cut goes straight to lossless // FINAL delivery; a redundant full BCAST would only inflate receiver RSS. + // Local init: islands whose digests agree with the applied INIT already + // hold version 0 and only get INIT_CONFIRM. A resumed cut is not the + // base checkpoint, so it is always broadcast. if st.global_step < last_step { - broadcast_all_fragments(&st, ®istry, &budget_cutoff).await; + let confirm = local_init + .as_ref() + .filter(|_| !resumed) + .and_then(|init| init.confirm().map(|digest| (init, digest))); + let groups = current_groups(®istry); + let (confirmed, rest): (Vec<_>, Vec<_>) = groups.into_iter().partition(|group| { + confirm.is_some_and(|(init, _)| init.holds(group.member.learner_id)) + }); + if let Some((_, digest)) = confirm { + for group in &confirmed { + let payload = bytes::Bytes::copy_from_slice(&digest); + if !send_small_until_cutoff(group, MSG_INIT_CONFIRM, payload, &budget_cutoff).await { + break; + } + } + info!( + confirmed = confirmed.len(), + broadcast = rest.len(), + "local init: islands kept their own version-0 weights" + ); + } + broadcast_all_fragments(&st, &rest, &budget_cutoff).await; } // Phase 2: gather, compute, and commit are separate stages. Torch SVD @@ -3431,7 +3529,7 @@ async fn scheduler( break 'outer; } } - Event::Init { .. } => {} // already initialized; ignore + Event::Init { .. } | Event::InitDigest { .. } => {} // already initialized; ignore Event::FinalAck { member, .. } => { warn!( learner_id = member.learner_id, @@ -3874,7 +3972,7 @@ async fn collect_budget_reports( generation = member.generation, "disconnected while waiting for learner-budget cutoff" ), - Event::Init { .. } | Event::FinalAck { .. } | Event::Leave { .. } => {} + Event::Init { .. } | Event::InitDigest { .. } | Event::FinalAck { .. } | Event::Leave { .. } => {} Event::Recovery { member, .. } => ignore_recovery_message(member, "the budget cutoff"), } } @@ -5158,10 +5256,9 @@ fn is_current_member(registry: &Registry, member: Member) -> bool { async fn broadcast_all_fragments( st: &GlobalState, - registry: &Registry, + groups: &[Arc], budget_cutoff: &BudgetCutoff, ) { - let groups = current_groups(registry); // Fan each fragment out to every learner before moving to the next one. // Each learner has independent socket-writer queues, so this keeps all // links busy concurrently. Iterating learner-by-learner here serializes @@ -5169,7 +5266,7 @@ async fn broadcast_all_fragments( // The current streaming encoder keeps peak host memory bounded while the // per-learner socket queues provide link concurrency. for p in st.owned_fragments() { - for group in &groups { + for group in groups { match send_fragment_until_cutoff( st, group, @@ -5444,6 +5541,7 @@ async fn finalize_learners( | Event::Heartbeat { .. } | Event::Leave { .. } | Event::Init { .. } + | Event::InitDigest { .. } | Event::Recovery { .. } | Event::BudgetDone { .. } => { // The authoritative cut is frozen; late training traffic is diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index 763510372..7072fe0fd 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -292,6 +292,14 @@ def test_bf16_push_flag_reaches_the_island_command(): assert "--sweep-push-dtype" not in _rl_trainable_island_flags(_dense_args(), 0, 1) +def test_init_mode_flag_reaches_the_island_command(): + from yeto.launch.rl_island_task import _rl_trainable_island_flags + + args = _dense_args("--rl-sync-init-mode", "upload") + assert "--init-mode upload" in _rl_trainable_island_flags(args, 0, 1) + assert "--init-mode" not in _rl_trainable_island_flags(_dense_args(), 0, 1) + + def _tail(data: bytes) -> torch.Tensor: return torch.frombuffer(bytearray(data), dtype=torch.float32) diff --git a/tests/test_rl_local_init_real_syncer.py b/tests/test_rl_local_init_real_syncer.py new file mode 100644 index 000000000..34078e73a --- /dev/null +++ b/tests/test_rl_local_init_real_syncer.py @@ -0,0 +1,276 @@ +"""S27 local init against the real Rust syncer (CPU, no Ray). + +Every island loaded the same base weights. Islands report per-fragment +INIT digests; learner 0 still uploads INIT (the syncer keeps exact f32 +masters), but an island whose digests agree gets INIT_CONFIRM instead of the +version-0 broadcast. Merged rounds must equal the old INIT path bit for bit, +a disagreeing island must stop the run with a clear error, and a resumed +syncer must broadcast its checkpointed cut instead of confirming. +""" + +from __future__ import annotations + +import hashlib +import time + +import numpy as np +import pytest + +from test_rl_outer_merge_real_syncer import ( # noqa: F401 - fixtures + _deltas, + _f32, + _fake_sky, + _frag_of, + _hmac_key, + _initial, + _no_modal_listing, + Reference, + Rig, + syncer_binary, +) +from yeto.protocol import DTYPE_F32, SessionRejectedError, SyncerClient + +LAYOUT_HASH = hashlib.sha256(b"local-init-test-layout").digest() +LR, MU = 0.35, 0.6 + + +def _digests(initial): + return [hashlib.sha256(_f32(values)).digest() for values in initial] + + +class LocalRig(Rig): + """``Rig`` whose islands may report local-init digests before start.""" + + def __init__(self, binary, tmp_path, islands, *options, **kwargs): + # The launcher-built command adds --resume when the checkpoint exists. + super().__init__(binary, tmp_path, islands, "--rl-sync-outer-lr", str(LR), + "--rl-sync-outer-momentum", str(MU), *options, **kwargs) + + def start_local(self, initial, digests_by_island): + """Islands in ``digests_by_island`` report digests; learner 0 sends INIT.""" + for island, digests in digests_by_island.items(): + self.clients[island].set_init_digests(LAYOUT_HASH, digests) + for client in self.clients: + client.start() + time.sleep(0.5) # let every island's digests reach the syncer first + for fid, values in enumerate(initial): + self.clients[0].send_init(fid, _f32(values)) + + def wait_confirmed(self, island, timeout=20.0): + client, deadline = self.clients[island], time.monotonic() + timeout + while not client.initial_cut_confirmed(): + client.check_health() + assert time.monotonic() < deadline, f"island {island} got no INIT_CONFIRM" + time.sleep(0.01) + client.mark_initial_cut_received() + + def wait_exit(self, timeout=15.0): + deadline = time.monotonic() + timeout + while self.process.poll() is None and time.monotonic() < deadline: + time.sleep(0.05) + return self.process.poll() + + +@pytest.fixture +def make_local_rig(syncer_binary, tmp_path): + rigs = [] + + def make(islands, *options, workdir=None, **kwargs): + if workdir is None: + workdir = tmp_path / f"rig{len(rigs)}" + workdir.mkdir() + rigs.append(LocalRig(syncer_binary, workdir, islands, *options, **kwargs)) + return rigs[-1] + + yield make + for rig in rigs: + rig.close() + + +def _run_rounds(rig, initial, steps, *, ref=None, seed=40): + """Push both islands' deltas for ``steps``; return every merged BCAST.""" + ref = ref or Reference(rig.layout, initial, LR, MU) + got = {} + for step in steps: + fid = _frag_of(rig, step) + deltas = _deltas(rig.layout, fid, seed + step, 2) + base = max(step - rig.layout.num_fragments, 0) + for island in (0, 1): + rig.push(island, step, deltas[island], base=base) + want = ref.step(fid, deltas, [1.0, 1.0]) + for island in (0, 1): + got[(island, step)] = rig.bcast(island, fid, step).copy() + np.testing.assert_allclose(got[(0, step)], want, rtol=2e-5, atol=2e-5) + return got + + +def test_local_init_confirms_without_broadcast_and_merges_like_init_broadcast(make_local_rig): + steps = range(1, 7) + old = make_local_rig(2, total_steps=4) + initial = _initial(old.layout) + old.start(initial) + reference = _run_rounds(old, initial, steps) + + rig = make_local_rig(2, total_steps=4) + rig.start_local(initial, {0: _digests(initial), 1: _digests(initial)}) + for island in (0, 1): + rig.wait_confirmed(island) + assert not rig.clients[island].drain_updates(), "no version-0 broadcast in local init" + local = _run_rounds(rig, initial, steps) + for key, values in reference.items(): # identical bits, not just close + assert values.tobytes() == local[key].tobytes(), key + + +def test_island_without_digests_still_gets_the_broadcast(make_local_rig): + rig = make_local_rig(2, total_steps=4) + initial = _initial(rig.layout) + rig.start_local(initial, {0: _digests(initial)}) + rig.wait_confirmed(0) + for fid in range(rig.layout.num_fragments): + assert rig.bcast(1, fid, 0).tobytes() == np.asarray(initial[fid]).tobytes() + _run_rounds(rig, initial, range(1, 3)) + + +@pytest.mark.parametrize("bad", ["island", "init"]) +def test_digest_mismatch_stops_the_run_with_a_clear_error(make_local_rig, bad): + rig = make_local_rig(2) + initial = _initial(rig.layout) + digests = _digests(initial) + other = list(digests) + other[1] = hashlib.sha256(b"another checkpoint").digest() + # "island": island 1 loaded other weights; "init": learner 0 uploads + # bytes that differ from the digests both islands agreed on. + by_island = {0: digests, 1: other} if bad == "island" else {0: other, 1: other} + rig.start_local(initial, by_island) + assert rig.wait_exit() not in (None, 0) + rig.close() + assert "init digest mismatch: fragment 1" in rig.output + assert "same base checkpoint" in rig.output + # The island that hears the refusal fails as a strict invariant. + from yeto.rl.core import StrictRlInvariantError + + errors = [client._err for client in rig.clients] + assert any(isinstance(e, SessionRejectedError) and e.metric == "init_digest_mismatch" + for e in errors), errors + with pytest.raises(StrictRlInvariantError, match="init digest mismatch"): + for client in rig.clients: + client.check_health() + + +def test_resumed_syncer_broadcasts_its_checkpoint_instead_of_confirming(make_local_rig, tmp_path): + workdir = tmp_path / "resume" + workdir.mkdir() + rig = make_local_rig(2, total_steps=4, workdir=workdir) + initial = _initial(rig.layout) + rig.start_local(initial, {0: _digests(initial), 1: _digests(initial)}) + for island in (0, 1): + rig.wait_confirmed(island) + ref = Reference(rig.layout, initial, LR, MU) + first = _run_rounds(rig, initial, range(1, 4), ref=ref) + time.sleep(0.5) # the round-3 checkpoint commits before its broadcast; be safe + rig.close() + assert (workdir / "state.ckpt").exists() + + resumed = make_local_rig(2, total_steps=4, workdir=workdir) + assert "--resume" in resumed.command + resumed.start_local(initial, {0: _digests(initial), 1: _digests(initial)}) + # Fragment 0 was last merged in round 3, fragment 1 in round 2. + for island in (0, 1): + assert resumed.bcast(island, 0, 3).tobytes() == first[(0, 3)].tobytes() + assert resumed.bcast(island, 1, 2).tobytes() == first[(0, 2)].tobytes() + assert not resumed.clients[island].initial_cut_confirmed() + _run_rounds(resumed, initial, range(4, 7), ref=ref) + + +def test_set_init_digests_validates_its_input(): + from yeto.fragments import MERGE_AVG, Fragment, FragmentLayout + + layout = FragmentLayout([Fragment(MERGE_AVG, [("w", 4)])]) + client = SyncerClient(("127.0.0.1", 9), 0, layout, dtype=DTYPE_F32) + with pytest.raises(ValueError): + client.set_init_digests(LAYOUT_HASH, []) + with pytest.raises(ValueError): + client.set_init_digests(b"short", [b"x" * 32]) + client.set_init_digests(LAYOUT_HASH, [b"x" * 32]) + assert not client.initial_cut_confirmed() + + +# ---- role stream (no syncer) ---------------------------------------------------------- + +class _Stored: + def __init__(self, fragment_id, payload=b""): + self.fragment_id, self.version, self.payload, self.released = fragment_id, 0, payload, False + + def release(self): + self.released = True + + +def _fake_adapter(payloads): + from types import SimpleNamespace + + from yeto.fragments import MERGE_AVG, Fragment, FragmentLayout + + fragments = FragmentLayout([Fragment(MERGE_AVG, [(f"w{i}", len(p) // 4)]) + for i, p in enumerate(payloads)]) + layout = SimpleNamespace(fragments=fragments, layout_hash="1" * 64) + + def fragment_parts(_initial, fragment_id, ray_module=None): + data = payloads[fragment_id] + return lambda: [data[:4], data[4:]] # parts, as Miles streams them + + return SimpleNamespace(layout=layout, fragment_parts=fragment_parts) + + +def test_role_stream_keeps_local_version_zero_after_the_confirm(monkeypatch): + from dataclasses import replace + + from test_rl_full_parameter_sync import _role_stream_config + from yeto.rl.adapters.miles.full_parameter import role_stream as module + + payloads = [_f32([1, 2, 3]), _f32([4, 5])] + adapter = _fake_adapter(payloads) + stream = module.MilesFullParameterRoleStream( + replace(_role_stream_config("actor", 30100), local_init=True)) + local, digests = stream._local_version_zero(adapter, object(), lambda fid, v, p: _Stored(fid, bytes(p))) + assert [s.payload for s in local] == payloads + assert digests == [hashlib.sha256(p).digest() for p in payloads] + + class ConfirmingClient: + marked = False + + def check_health(self): + pass + + def drain_updates(self): + return [] + + def initial_cut_confirmed(self): + return True + + def mark_initial_cut_received(self): + self.marked = True + + client = ConfirmingClient() + assert stream._wait_for_initial_fragments(client, adapter, None, local) == local + assert client.marked and not any(s.released for s in local) + + +def test_initial_cut_timeout_scales_with_the_model_and_the_init_mode(): + from yeto.rl.adapters.miles.full_parameter.role_stream import initial_cut_timeout + + x6 = 116_500_000_000 # X6: f32 Qwen3-30B-A3B + upload = initial_cut_timeout(900.0, x6, 2, local_init=False) + local = initial_cut_timeout(900.0, x6, 2, local_init=True) + assert upload == pytest.approx(3 * local) and local > 3000 # 900 s timed out in X6 + assert initial_cut_timeout(900.0, 1 << 20, 2, local_init=False) == 900.0 + + +def test_local_init_is_the_default_for_async_rounds_only(): + from types import SimpleNamespace + + from yeto.rl.adapters.miles.island_entry import _local_init + + assert _local_init(SimpleNamespace(init_mode="auto"), True) + assert not _local_init(SimpleNamespace(init_mode="auto"), False) + assert _local_init(SimpleNamespace(init_mode="local"), False) + assert not _local_init(SimpleNamespace(init_mode="upload"), True) diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 17f72fa7f..83b16c4f0 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -171,6 +171,7 @@ def sao_contract_mounts(args, learner_id: int) -> dict[str, str]: ("rl_sync_overlap", False, "--rl-sync-overlap"), ("rl_sync_full_parameter_rounds", "sweep", "--rl-sync-full-parameter-rounds"), ("rl_sync_expert_token_weights", False, "--rl-sync-expert-token-weights"), + ("rl_sync_init_mode", "auto", "--rl-sync-init-mode"), ) diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 7f3f17622..cd1917c2c 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -112,6 +112,8 @@ def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: sweep += " --full-parameter-rounds async" if getattr(args, "rl_sync_expert_token_weights", False): sweep += " --sweep-expert-token-weights" + if getattr(args, "rl_sync_init_mode", "auto") != "auto": + sweep += f" --init-mode {args.rl_sync_init_mode}" if getattr(args, "rl_sync_push_dtype", "f32") != "f32": sweep += f" --sweep-push-dtype {args.rl_sync_push_dtype}" if getattr(args, "rl_sync_overlap", False): diff --git a/yeto/protocol.py b/yeto/protocol.py index 06c068003..1002f11e6 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -69,6 +69,11 @@ MSG_BW_PROBE = 26 # An island leaves the pool: u32 learner_id | u8 reason (yeto.rl.pool_client.LeaveReason). MSG_LEAVE = 27 +# Local-init startup (S27): the island's per-fragment INIT digests, and the +# syncer's confirmation that every island holds the same version 0 (see +# syncer/src/local_init.rs). +MSG_INIT_DIGEST = 28 +MSG_INIT_CONFIRM = 29 FINALIZATION_REVISION = 1 FINALIZATION_TIMEOUT = 900.0 @@ -722,6 +727,8 @@ class ProtocolError(RuntimeError): ("HELLO authentication failed", "island_auth_failed"), # The island was launched for an older syncer incarnation. (rp.SYNCER_EPOCH_FENCE_PREFIX, "syncer_epoch_fenced"), + # Local init: the islands did not load the same base weights. + ("init digest mismatch", "init_digest_mismatch"), ) @@ -1014,6 +1021,11 @@ def __init__( # the kernel on a slow link (s26-e3b-p06-C32a). self._initial_cut_received = threading.Event() self._init_generation: int | None = None + # Local init: the INIT_DIGEST frame sent after every HELLO until the + # initial cut arrives, and the INIT_CONFIRM digest it must earn. + self._init_digest_frame: bytes | None = None + self._init_combined: bytes | None = None + self._init_confirmed = threading.Event() self.session_contract_hash = session_contract_hash self.syncer_profile_hash = syncer_profile_hash # secret-handling-hardening: HELLO carries an HMAC when the island has @@ -1167,6 +1179,8 @@ def _connect_group(self, patient: bool) -> None: self.hmac_key, ), ) + if self._init_digest_frame is not None and not self._initial_cut_received.is_set(): + control.sendall(self._init_digest_frame) socks.append(control) for idx in range(self.num_streams): s = self._connect_one() if patient else self._dial(RECONNECT_DIAL_TIMEOUT) @@ -1383,6 +1397,29 @@ def init_replay_due(self) -> bool: and generation != self._init_generation ) + def set_init_digests(self, layout_hash: bytes, digests: Sequence[bytes]) -> None: + """Local init: report the SHA-256 of every fragment's INIT_PARAMS + tensor bytes (one per layout fragment). Call before :meth:`start`; + the digests then follow every HELLO until the initial cut arrives, + and :meth:`initial_cut_confirmed` turns true when the syncer agrees.""" + if len(layout_hash) != 32 or len(digests) != self.layout.num_fragments: + raise ValueError("local init needs a 32-byte layout hash and one digest per fragment") + if bulk_dtype(self.dtype) != DTYPE_F32: + raise ValueError("local init needs an f32-bulk session") + owned = [bytes(digests[index]) for index in self.owned_fragments] + if any(len(digest) != 32 for digest in owned): + raise ValueError("fragment digests must be 32-byte SHA-256 values") + if self._connected.is_set(): + raise RuntimeError("set_init_digests must be called before start()") + payload = (struct.pack(" bool: + """True once the syncer confirmed this island's local version 0.""" + return self._init_confirmed.is_set() + def mark_initial_cut_received(self) -> None: """The initial cut arrived: INIT is no longer replayed, and the stall timeout now also bounds unacknowledged data in the kernel.""" @@ -2151,6 +2188,7 @@ def _incoming_frame_limit(self, msg_type: int) -> int: MSG_FINAL_MANIFEST: _FINAL_MANIFEST_HEAD.size + 8 * self.layout.num_fragments, MSG_SHUTDOWN: 0, + MSG_INIT_CONFIRM: 32, MSG_ERROR: MAX_ERROR_FRAME, MSG_CHUNK: _CHUNK_HEAD.size + CHUNK_SIZE, rp.MSG_RECOVERY_CLOCK: rp.MAX_RECOVERY_FRAME, @@ -2558,6 +2596,13 @@ def _dispatch_live( self._recover.clear() self._recovery_events.put(event) return + if msg_type == MSG_INIT_CONFIRM: + if self._init_combined is None: + raise ProtocolError("INIT_CONFIRM for an island that sent no INIT_DIGEST") + if bytes(payload) != self._init_combined: + raise ProtocolError("INIT_CONFIRM digest differs from this island's fragments") + self._init_confirmed.set() + return if msg_type == MSG_PULL_REQ: if len(payload) != 16: raise ValueError("PULL_REQ payload must be 16 bytes") diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 0e4072f97..703c46dff 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -59,6 +59,21 @@ # Seconds between progress lines while an island waits for its initial cut. INITIAL_CUT_REPORT_SECONDS = 30.0 +# The initial-cut wait grows with the model: long enough to move the cut at +# an assumed worst-case 32 MiB/s head link (X6: 116.5 GB of f32 timed out the +# flat 900 s default, which stays the floor). +INITIAL_CUT_MIN_BYTES_PER_SECOND = 32 << 20 + + +def initial_cut_timeout(configured: float, model_bytes: int, islands: int, *, local_init: bool) -> float: + """Seconds an island waits for its initial cut. + + Learner 0 uploads ``model_bytes`` of INIT. Without local init the syncer + then broadcasts them to every island over the head's one link, so the cut + moves ``(1 + islands) * model_bytes``; with local init only the upload. + The configured timeout is the floor.""" + moved = model_bytes * (1 if local_init else 1 + islands) + return max(configured, moved / INITIAL_CUT_MIN_BYTES_PER_SECOND) _LOWER_HEX = frozenset("0123456789abcdef") @@ -160,6 +175,12 @@ class RoleStreamConfig: # may skip fragment versions. None = fixed roster. Not part of the # session contract (membership is a syncer option). island_lease_s: float | None = None + # Local init (S27): this island loaded the run's base checkpoint and keeps + # it as version 0 once the syncer confirms every island's fragment digests + # agree (MSG_INIT_CONFIRM); no initial broadcast is needed. False waits for + # the broadcast of learner 0's INIT. Not part of the session contract: it + # changes which bytes travel at startup, never the version-0 values. + local_init: bool = False @property def tensor_blocks_active(self) -> bool: @@ -559,7 +580,13 @@ async def initialize(self, group) -> None: if self._recovering: await self._initialize_recovering(group, adapter, initial, client, sink, heartbeat) return + local: tuple[StoredAuthoritativeFragment, ...] | None = None try: + if self.config.local_init: + local, digests = await asyncio.to_thread( + self._local_version_zero, adapter, initial, sink) + client.set_init_digests( + hashlib.sha256(adapter.layout.layout_hash.encode()).digest(), digests) client.start() if heartbeat is not None: heartbeat.start() # the initial cut may take longer than the lease @@ -579,9 +606,14 @@ def init_parts(fragment_id: int): # island's other role from connecting, and two islands blocked # on different roles deadlock (s23-g3-sao-main-20261014a). authoritative = await asyncio.to_thread( - self._wait_for_initial_fragments, client, adapter, init_parts + self._wait_for_initial_fragments, client, adapter, init_parts, local ) adapter.verify_initial_fragments(initial, authoritative, allow_newer=self.may_rejoin) + if local is not None: # a rejoin's newer fragments replace version 0 + kept = {id(stored) for stored in authoritative} + for stored in local: + if id(stored) not in kept: + stored.release() except BaseException: if heartbeat is not None: heartbeat.stop() @@ -736,20 +768,46 @@ async def serve_peer_snapshots(self) -> int: self._donor_refusals_reported = len(self._donor.refused) return served + def _local_version_zero(self, adapter, initial, sink): + """Local init: store this island's own INIT bytes as its version-0 + anchors and hash them exactly as the syncer hashes INIT_PARAMS.""" + stored, digests = [], [] + for fragment_id in range(adapter.layout.fragments.num_fragments): + payload = bytearray() + for part in adapter.fragment_parts(initial, fragment_id, ray_module=self.ray_module)(): + payload += part + digests.append(hashlib.sha256(payload).digest()) + stored.append(sink(fragment_id, 0, payload)) + return tuple(stored), digests + def _wait_for_initial_fragments( self, client, adapter: MilesChunkedFullParameterAdapter, init_parts=None, + local: Sequence[StoredAuthoritativeFragment] | None = None, ) -> tuple[StoredAuthoritativeFragment, ...]: + """The version-0 cut: the syncer's broadcast, or with ``local`` this + island's own fragments once the syncer confirms the digests agree.""" started = time.monotonic() - deadline = started + self.config.wait_timeout + count = adapter.layout.fragments.num_fragments + fragments = getattr(adapter.layout.fragments, "fragments", ()) + model_bytes = sum(4 * fragment.numel for fragment in fragments) + deadline = started + initial_cut_timeout( + self.config.wait_timeout, model_bytes, len(self.config.learner_generations), + local_init=local is not None) next_report = started + INITIAL_CUT_REPORT_SECONDS observed: dict[int, StoredAuthoritativeFragment] = {} - count = adapter.layout.fragments.num_fragments self._report_initial_cut(client, observed, count, started, "waiting") while len(observed) < count: client.check_health() + if local is not None and client.initial_cut_confirmed(): + for stored in observed.values(): + if all(stored is not mine for mine in local): + stored.release() + self._report_initial_cut(client, observed, count, started, "confirmed local") + client.mark_initial_cut_received() + return tuple(local) if self.may_rejoin and _is_set(getattr(client, "finalizing", None)): # The pool reached its final cut before this island caught up; # waiting for broadcasts would only run into the timeout. diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 25a9bcc53..7c3a7bc6d 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -882,6 +882,14 @@ def _full_parameter_async(args) -> bool: return getattr(args, "full_parameter_rounds", "sweep") == "async" +def _local_init(args, async_rounds: bool) -> bool: + """Local init (S27): every island launched by one run loads the same base + checkpoint, so async rounds keep it as version 0 by default; the syncer + refuses the run if the islands' fragment digests disagree.""" + mode = getattr(args, "init_mode", "auto") + return mode == "local" or (mode == "auto" and async_rounds) + + def _dense_full_sync_config( args, *, model_config_sha256: str, training_contract_sha256: str, algorithm_spec_sha256: str ): @@ -932,6 +940,7 @@ def _dense_full_sync_config( # overlap's outer-step lag is bound into the session contract. overlap=bool(getattr(args, "sweep_overlap", False)), partial_quorum=async_rounds, + local_init=_local_init(args, async_rounds), matrix_merge=getattr(args, "rl_sync_matrix_merge", "avg"), peer_recovery=peer_recovery_options(args), # an island pool sets the epoch (and lease); else peer recovery may diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index fe72049f4..1036225ea 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -336,6 +336,14 @@ "asynchronous sharded rounds (H = --local-horizon, a quorum below the " "island count may close a round without this island)", ), + island_flag( + "--init-mode", + choices=("auto", "local", "upload"), + default="auto", + help="full-parameter: local = keep the loaded base weights as version 0 " + "after the syncer confirms every island's fragment digests agree; upload " + "= wait for the syncer's initial broadcast; auto = local for async rounds", + ), island_flag( "--sweep-push-dtype", choices=("f32", "bf16", "int8", "int4", "int4-log"), diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 0e18ff9d0..a8dff483a 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -1020,6 +1020,17 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. ), # S27 P1: the outer-merge settings each RL preset used to fix. Unset keeps # the preset's value, so the syncer command and profile stay unchanged. + cli_flag( + "--rl-sync-init-mode", + choices=("auto", "local", "upload"), + default="auto", + help="full-parameter only (S27): how islands get the version-0 weights. " + "local = every island keeps the base checkpoint it loaded and reports " + "per-fragment digests; the syncer checks they agree and sends no initial " + "broadcast (learner 0 still uploads INIT for the syncer's f32 masters). " + "upload = learner 0 uploads INIT and the syncer broadcasts it to every " + "island. auto (default) = local with --rl-sync-full-parameter-rounds async", + ), cli_flag( "--rl-sync-full-parameter-rounds", choices=FULL_PARAMETER_ROUNDS, diff --git a/yeto/syncer_shards.py b/yeto/syncer_shards.py index 7511eae26..f183840d6 100644 --- a/yeto/syncer_shards.py +++ b/yeto/syncer_shards.py @@ -161,6 +161,14 @@ def init_replay_due(self) -> bool: """True when some shard needs INIT again (see ``SyncerClient``).""" return any(client.init_replay_due() for client in self.shards) + def set_init_digests(self, layout_hash: bytes, digests) -> None: + """Local init: each shard gets the digests of the fragments it owns.""" + for client in self.shards: + client.set_init_digests(layout_hash, digests) + + def initial_cut_confirmed(self) -> bool: + return all(client.initial_cut_confirmed() for client in self.shards) + def mark_initial_cut_received(self) -> None: for client in self.shards: client.mark_initial_cut_received() From 0c2fc16278ae21b189705c6682321006e12b8461 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:43:23 +0000 Subject: [PATCH 154/233] verl: readback undoes the served model's vLLM prefix renaming (Qwen3.5 VL publication check) Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_verl_adapter.py | 16 ++++++++++++++++ yeto/rl/adapters/verl/param_names.py | 26 +++++++++++++++++++++++--- yeto/rl/adapters/verl/vllm_readback.py | 18 ++++++++++++++---- 3 files changed, 53 insertions(+), 7 deletions(-) diff --git a/tests/test_rl_verl_adapter.py b/tests/test_rl_verl_adapter.py index 1a417dcd1..cc01ec03e 100644 --- a/tests/test_rl_verl_adapter.py +++ b/tests/test_rl_verl_adapter.py @@ -471,3 +471,19 @@ def test_fully_async_runtime_env_keeps_nccl_host_buffers_on_shm(monkeypatch): assert verl_main.CUMEM_HOST_ENV not in verl_main._runtime_env()["env_vars"] monkeypatch.setenv(verl_main.CUMEM_HOST_ENV, "1") assert verl_main._runtime_env(fully_async=True)["env_vars"][verl_main.CUMEM_HOST_ENV] == "1" + + +def test_vllm_prefix_renaming_is_undone(): + """Qwen3.5 VL in vLLM serves 'model.language_model.*' as 'language_model.model.*' (E10).""" + from types import SimpleNamespace + + from yeto.rl.adapters.verl import param_names as pn + + mapper = SimpleNamespace(orig_to_new_prefix={"model.language_model.": "language_model.model.", + "model.visual.": "visual.", "mtp.": None}) + hf = pn.hf_prefixes_from_mapper(mapper) + assert hf == {"language_model.model.": "model.language_model.", "visual.": "model.visual."} + assert pn.vllm_to_canonical("language_model.model.layers.0.linear_attn.out_proj", "B", None, hf) == \ + "base_model.model.model.language_model.layers.0.linear_attn.out_proj.lora_B.weight" + assert pn.vllm_to_canonical("model.layers.0.self_attn.o_proj", "A", None, hf) == \ + "base_model.model.model.layers.0.self_attn.o_proj.lora_A.weight" diff --git a/yeto/rl/adapters/verl/param_names.py b/yeto/rl/adapters/verl/param_names.py index fc663200f..243b901b0 100644 --- a/yeto/rl/adapters/verl/param_names.py +++ b/yeto/rl/adapters/verl/param_names.py @@ -34,6 +34,9 @@ "train_rule": "drop '.' before '.weight' in '.lora_{A,B}..weight'", "adapter": ADAPTER_NAME, "infer_rule": "prefix 'base_model.model.'; unpack vLLM packed modules in listed order", + # vLLM renames HF module prefixes for some models (Qwen3.5 VL: 'model.language_model.' + # -> 'language_model.model.'); the readback undoes the model's own mapper (E10) + "infer_prefix_rule": "invert the vLLM model's hf_to_vllm_mapper orig_to_new_prefix", "vllm_packed_modules": {k: list(v) for k, v in VLLM_PACKED_MODULES.items()}, } @@ -68,11 +71,28 @@ def canonical_to_train(name: str) -> str: return f"{CANONICAL_PREFIX}{match['module']}.lora_{match['ab']}.{ADAPTER_NAME}.weight" -def vllm_to_canonical(module_name: str, ab: str, packed_index: int | None = None) -> str: - """vLLM ``LoRAModel`` entry (module, A/B, slot in a packed module) -> canonical name.""" +def hf_prefixes_from_mapper(mapper) -> dict[str, str]: + """{vLLM prefix: HF prefix} from a vLLM ``WeightsMapper`` (its ``orig_to_new_prefix`` + inverted; prefixes vLLM drops map to None there and are skipped).""" + prefixes = getattr(mapper, "orig_to_new_prefix", None) or {} + return {new: orig for orig, new in prefixes.items() if new} + + +def vllm_to_hf_module(module_name: str, hf_prefixes: dict[str, str] | None) -> str: + """Undo vLLM's prefix renaming (longest matching vLLM prefix wins).""" + for new in sorted(hf_prefixes or {}, key=len, reverse=True): + if module_name.startswith(new): + return hf_prefixes[new] + module_name[len(new):] + return module_name + + +def vllm_to_canonical(module_name: str, ab: str, packed_index: int | None = None, + hf_prefixes: dict[str, str] | None = None) -> str: + """vLLM ``LoRAModel`` entry (module, A/B, slot in a packed module) -> canonical name. + ``hf_prefixes`` (:func:`hf_prefixes_from_mapper`) maps vLLM module prefixes back to HF.""" if ab not in ("A", "B"): raise ParamNameError(f"lora side must be A or B, got {ab!r}") - module_name = module_name.removeprefix(CANONICAL_PREFIX) + module_name = vllm_to_hf_module(module_name.removeprefix(CANONICAL_PREFIX), hf_prefixes) parent, _, leaf = module_name.rpartition(".") if leaf in VLLM_PACKED_MODULES: if packed_index is None: diff --git a/yeto/rl/adapters/verl/vllm_readback.py b/yeto/rl/adapters/verl/vllm_readback.py index e3c386880..bfb1944c4 100644 --- a/yeto/rl/adapters/verl/vllm_readback.py +++ b/yeto/rl/adapters/verl/vllm_readback.py @@ -33,7 +33,16 @@ def _registered_lora_model(worker, lora_int_id: int): return model -def canonical_tensors_from_lora_model(lora_model) -> tuple[dict, dict]: +def _hf_prefixes(worker) -> dict[str, str]: + """The served model's vLLM -> HF module prefixes ({} when the model renames none).""" + from .param_names import hf_prefixes_from_mapper + + runner = worker.model_runner + model = runner.get_model() if hasattr(runner, "get_model") else getattr(runner, "model", None) + return hf_prefixes_from_mapper(getattr(model, "hf_to_vllm_mapper", None)) + + +def canonical_tensors_from_lora_model(lora_model, hf_prefixes: dict[str, str] | None = None) -> tuple[dict, dict]: """{canonical name: tensor} from a vLLM ``LoRAModel`` (+ a note dict).""" from .param_names import vllm_to_canonical @@ -45,8 +54,8 @@ def canonical_tensors_from_lora_model(lora_model) -> tuple[dict, dict]: else: pairs = [(None, a, b)] for slot, la, lb in pairs: - out[vllm_to_canonical(module_name, "A", slot)] = la - out[vllm_to_canonical(module_name, "B", slot)] = lb + out[vllm_to_canonical(module_name, "A", slot, hf_prefixes)] = la + out[vllm_to_canonical(module_name, "B", slot, hf_prefixes)] = lb return out, {"entries": len(lora_model.loras), "transposed": transposed} @@ -68,7 +77,8 @@ def after_add_lora(worker, lora_int_id: int, expected_shapes: dict | None = None try: from .publish import lora_checksum - tensors, note = canonical_tensors_from_lora_model(_registered_lora_model(worker, lora_int_id)) + tensors, note = canonical_tensors_from_lora_model(_registered_lora_model(worker, lora_int_id), + _hf_prefixes(worker)) if expected_shapes: # vLLM may keep lora_a as (in, rank) on some versions; record, do not guess silently. for name, tensor in list(tensors.items()): From 492a563f7e33bcbe0d7fa3c38928f01b14faffd7 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:30:21 +0000 Subject: [PATCH 155/233] rl: stop local training only when every fragment's last round is ready With a fragment-sharded syncer head each shard runs its steps independently, so the terminal-step PULL of one shard can arrive while another shard has not yet asked for an earlier fragment's last round. The island treated the terminal PULL alone as the end of training, stopped producing local rounds, and could no longer answer the other shard's PULL; the shard then waited for that push while the finished shard waited for final acknowledgements. RolePairing now sets terminal_seen only once every fragment's last round is either ready to submit on this island or already closed. With one syncer process PULLs arrive in step order, so behaviour there is unchanged. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_full_parameter_sync.py | 30 ++++++++++++++++++++++++++ yeto/rl/engine/full_parameter_sync.py | 31 +++++++++++++++++++++++++-- 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/tests/test_rl_full_parameter_sync.py b/tests/test_rl_full_parameter_sync.py index e7fe47168..5314c9e10 100644 --- a/tests/test_rl_full_parameter_sync.py +++ b/tests/test_rl_full_parameter_sync.py @@ -386,6 +386,36 @@ async def test_terminal_final_manifest_drains_all_retained_pipeline_pairs(): assert pairing.acknowledge_final() +def test_terminal_pull_of_one_shard_does_not_stop_training_before_other_last_rounds(): + # Sharded head: shard 1 (fragment 1) asks for the terminal step 4 while + # shard 0 has not yet asked for step 3, fragment 0's last round. The + # island must keep training so it can still answer step 3. + events = [] + actor = _PairedRole("actor", [_candidate("actor", 4, base_version=2)], events, + release_broadcasts=False) + actor.fragment_versions = (1, 2) + pairing = _pairing(actor) + + pairing.submit_ready() + assert ("actor", 4, True) in events + assert not pairing.terminal_seen and not pairing.any_finalizing + + actor.candidates.append(_candidate("actor", 3, base_version=1)) + pairing.submit_ready() + assert ("actor", 3, True) in events + assert pairing.terminal_seen + + +def test_terminal_pull_stops_training_when_other_last_rounds_already_closed(): + events = [] + actor = _PairedRole("actor", [_candidate("actor", 4, base_version=2)], events) + actor.fragment_versions = (3, 2) # fragment 0's last round (step 3) closed + pairing = _pairing(actor) + + pairing.submit_ready() + assert pairing.terminal_seen + + @pytest.mark.asyncio async def test_one_role_is_the_same_rule_without_a_peer(): events = [] diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index 29a93c3b3..f876c6fd8 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -266,6 +266,9 @@ def __init__(self, streams: Mapping[str, RoleStream], profile: FullParameterSync self.profile = profile self.inflight: dict[LogicalKey, PairedSubmission] = {} self.journal: list[dict[str, Any]] = [] + # Fragments whose last round this island can submit (or that closed + # without it); the island stops training once it covers all of them. + self._final_rounds_ready: set[int] = set() self.terminal_seen = False self.failure: BaseException | None = None self._pull_skew_since: dict[LogicalKey, float] = {} @@ -342,8 +345,7 @@ def submit_ready(self) -> dict[str, RoleProgress]: for role, stream in self.streams.items() } now = time.monotonic() - if any(key[0] == self.profile.total_fragment_steps for keys in by_role.values() for key in keys): - self.terminal_seen = True + self._note_final_rounds(by_role) for key, pair in tuple(self.inflight.items()): observed = {role: by_role[role].pop(key, None) for role in self.streams} self._push_once(pair, observed, produced) @@ -364,6 +366,31 @@ def submit_ready(self) -> dict[str, RoleProgress]: raise return {role: RoleProgress(submissions=tuple(values)) for role, values in produced.items()} + def _note_final_rounds(self, by_role: Mapping[str, Mapping[LogicalKey, Any]]) -> None: + """Set ``terminal_seen`` once every fragment's last round is ready. + + Fragment ``f`` is last merged at step ``total - P + 1 + f``. With one + syncer process PULLs arrive in step order, so a ready terminal-step + PULL implied every other last round was ready too. A sharded head + runs each shard's steps independently: the terminal PULL of one shard + can arrive while another shard still has to ask for an earlier last + round. Stopping local training then leaves no round to answer that + PULL and both sides wait on each other. So training stops only when + each fragment's last round is ready here or already closed. + """ + + total, fragments = self.profile.total_fragment_steps, self.profile.fragments + first_last = total - fragments + 1 + for keys in by_role.values(): + self._final_rounds_ready.update(key[1] for key in keys if key[0] >= first_last) + for fragment_id in range(fragments): + if fragment_id not in self._final_rounds_ready and self._round_closed( + fragment_id, first_last + fragment_id + ): + self._final_rounds_ready.add(fragment_id) + if len(self._final_rounds_ready) == fragments: + self.terminal_seen = True + def _new_pair(self, candidates: dict[str, Any]) -> PairedSubmission: first = next(iter(candidates.values())) for candidate in candidates.values(): From d6dd6bfe2b0dee931f13980f8d8ab716af193bbe Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:52:52 +0000 Subject: [PATCH 156/233] rl: a retried round attempt supersedes the older PULL permit The syncer retries a round after a quorum timeout under a new attempt for the same step and base, and refuses pushes of the old attempt. The role stream raised on the second PULL and the island exited. Keep the newest attempt instead; a retained pair already rebinds its exact operands to a newer attempt. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_full_parameter_async.py | 31 +++++++++++++++++++ .../miles/full_parameter/role_stream.py | 12 +++++++ 2 files changed, 43 insertions(+) diff --git a/tests/test_rl_full_parameter_async.py b/tests/test_rl_full_parameter_async.py index d160dea59..c83a9b587 100644 --- a/tests/test_rl_full_parameter_async.py +++ b/tests/test_rl_full_parameter_async.py @@ -291,6 +291,37 @@ def test_newer_pull_supersedes_a_closed_round_only_with_a_partial_quorum(): strict._queue_pulls() +@pytest.mark.parametrize("partial_quorum", [True, False]) +def test_retried_round_attempt_supersedes_the_older_permit(partial_quorum): + # A quorum timeout makes the syncer resend the same round under a new + # attempt; it refuses pushes of the old one. + stream = _stream(partial_quorum) + stream.client.pulls = [PullRequest(0, 1, 1)] + stream._queue_pulls() + stream.client.pulls = [PullRequest(0, 1, 2)] + stream._queue_pulls() + assert stream._pending == {0: PullRequest(0, 1, 2)} + # A late duplicate of the old attempt never wins back. + stream.client.pulls = [PullRequest(0, 1, 1), PullRequest(0, 1, 2)] + stream._queue_pulls() + assert stream._pending == {0: PullRequest(0, 1, 2)} + # Several attempts in one drain keep the newest. + stream.client.pulls = [PullRequest(1, 2, 3), PullRequest(1, 2, 5), PullRequest(1, 2, 4)] + stream._queue_pulls() + assert stream._pending == {0: PullRequest(0, 1, 2), 1: PullRequest(1, 2, 5)} + + +@pytest.mark.parametrize("partial_quorum", [True, False]) +def test_retried_attempt_reopens_a_round_answered_under_the_old_attempt(partial_quorum): + stream = _stream(partial_quorum) + # Attempt 1 was pushed (answered, permit consumed); the retry asks again. + stream._queue_pulls() # records the link generations + stream._answered = {(1, 1): stream._link_generation(0)} + stream.client.pulls = [PullRequest(0, 1, 1), PullRequest(0, 1, 2)] + stream._queue_pulls() + assert stream._pending == {0: PullRequest(0, 1, 2)} + + @pytest.mark.asyncio @pytest.mark.parametrize("partial_quorum", [True, False]) async def test_applying_a_round_drops_its_unanswered_pull(partial_quorum): diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 703c46dff..9b5164037 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -1417,6 +1417,18 @@ def _queue_pulls(self) -> None: if permit.global_step > previous.global_step: self._pending[permit.fragment_id] = permit continue + if ( + previous is not None + and previous.global_step == permit.global_step + and previous.round_attempt != permit.round_attempt + ): + # The syncer retries a round (same step and base) under a new + # attempt after a quorum timeout and refuses pushes of the old + # attempt. The newest attempt supersedes; a retained pair + # rebinds its exact operands to it (``retained_replay_candidate``). + if permit.round_attempt > previous.round_attempt: + self._pending[permit.fragment_id] = permit + continue if previous is not None and previous != permit: raise RuntimeError("role-stream received conflicting PULLs") self._pending.setdefault(permit.fragment_id, permit) From 60239fffa3051dfce37c5f674abefba35fa4716a Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:07:48 +0000 Subject: [PATCH 157/233] Island pool with peer recovery: one option for a replacement island's state source The island pool decides which island fills a slot; peer recovery decides where that island takes its trainer state. The two now run together. - Add --rl-island-state-source {peer,catch-up,cut}; it and --rl-peer-recovery resolve to one value, contradictions are refused, cut is refused for synced islands (not implemented). - Peer recovery defaults on in an island pool where it applies; the pool already passes the syncer epoch, so it is not passed twice. - Role stream: drop the pool refusal; the lease heartbeat runs during and after the restore. - Syncer: in pool mode, record pool_join (recovered, base_version) when a recovery completes and refresh the pool status. - Tests: Rust unit test for weight 0 until promotion, lease expiry and re-join; integration test of a pool slot filled by peer recovery; Miles role-stream and launcher option tests. Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 57 +++++++++++++++ tests/test_peer_recovery.py | 71 ++++++++++++++++++- tests/test_peer_recovery_launch.py | 34 +++++++++ tests/test_peer_recovery_miles.py | 27 ++++++- tests/test_rl_sync_mode.py | 10 ++- yeto/launch/peer_recovery.py | 24 ++++--- yeto/launch/rl_args.py | 45 ++++++++---- yeto/launch/syncer.py | 3 +- .../miles/full_parameter/role_stream.py | 2 - yeto/rl/engine/run_flags.py | 15 +++- 10 files changed, 255 insertions(+), 33 deletions(-) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index f8ea31afa..4a1a9c46e 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -3524,6 +3524,20 @@ async fn scheduler( info!(learner_id = member.learner_id, steps = record.steps, path = ?record.source, "peer recovery done"); } send_recovery(®istry, vec![done]).await; + if dynamic { + // The pool took the island back with recovered state; its + // pushes are fenced at the versions it now holds. + append_pool_record( + cfg.event_tape.as_deref(), + "pool_join", + cfg.syncer_epoch, + &format!( + "\"island_id\":{},\"generation\":{},\"catch_up\":true,\"recovered\":true,\"base_version\":{}", + member.learner_id, member.generation, st.global_step + ), + ); + write_pool_status(&cfg, ®istry, st.global_step, "running"); + } if !repull_rounds_for(&group, ®istry, &mut inflight, fixed_roster, &budget_cutoff).await { cutoff_established = true; break 'outer; @@ -9221,6 +9235,49 @@ mod tests { assert!(r.current.get(&0) == Some(&back)); } + #[tokio::test] + async fn pool_member_under_peer_recovery_has_weight_zero_until_promoted() { + let registry = leased_registry(1); + let recover = HelloExt { syncer_epoch: Some(0), recover: true, ..Default::default() }; + let (donor, lost) = (member(1, 20), member(0, 10)); + let t0 = Instant::now() - Duration::from_secs(5); + { + let mut r = registry.lock().unwrap(); + r.register_group(test_group(donor), Instant::now()).unwrap(); + // A pool slot filled by an island that asks for peer recovery: + // admitted to the pool, kept off the round roster and donors. + r.register_group_with(test_group(lost), t0, recover).unwrap(); + assert!(r.is_recovering(lost)); + assert!(r.membership.is_live(lost)); + } + assert!(!is_current_member(®istry, lost)); + let candidates = recovery_candidates(®istry, &[]); + assert_eq!(candidates.iter().map(|c| c.member).collect::>(), vec![donor]); + + // Its lease runs out during the restore: it leaves the pool and the + // recovery; the next generation asks again (re-JOIN). + let mut rounds = vec![test_round(vec![donor])]; + let mut rates = StepRates::default(); + assert_eq!(expire_member_leases(®istry, &mut rounds, &mut rates).await, vec![lost]); + assert_eq!(rounds[0].expected_members, vec![donor]); + let mut r = registry.lock().unwrap(); + assert!(!r.is_recovering(lost)); + r.unregister(lost); + let back = member(0, 11); + assert_eq!(r.register_group_with(test_group(back), Instant::now(), recover).map(|x| x.1), + Ok(Admission::Rejoin)); + assert!(r.is_recovering(back) && !r.current.contains_key(&0)); + + // Promoted with the versions it holds after the replay: pushes based + // on older versions are fenced. + assert!(r.promote_recovered(back)); + r.membership.record_catch_up(back, &[7, 8]); + assert_eq!(r.current.get(&0), Some(&back)); + assert!(r.membership.check_push_base(back, 0, 6).is_err()); + r.membership.check_push_base(back, 0, 7).unwrap(); + assert!(!r.promote_recovered(back), "promotion happens once per generation"); + } + #[test] fn membership_options_are_refused_with_a_fixed_roster() { let mut cfg = round_test_config(4); diff --git a/tests/test_peer_recovery.py b/tests/test_peer_recovery.py index 27dbeb8d0..a7e1ac581 100644 --- a/tests/test_peer_recovery.py +++ b/tests/test_peer_recovery.py @@ -24,6 +24,7 @@ from yeto import recovery_protocol as rp from yeto.fragments import MERGE_AVG, Fragment, FragmentLayout from yeto.protocol import DTYPE_F32, BcastFragment, SyncerClient, encode_hello, layout_fingerprint +from yeto.rl.pool_client import LeaseHeartbeat from yeto.rl.peer_recovery import DirectorySnapshotStore, DonorService, Snapshot, read_array_state, recover ROOT = Path(__file__).resolve().parent.parent @@ -53,7 +54,8 @@ class FakeIsland(threading.Thread): serves donor snapshots, and finalizes. Implements RecoverableIsland.""" def __init__(self, learner_id: int, port: int, target: np.ndarray, store: DirectorySnapshotStore, - *, recover_on_start: bool = False, syncer_epoch: int | None = None) -> None: + *, recover_on_start: bool = False, syncer_epoch: int | None = None, + lease_s: float | None = None) -> None: super().__init__(daemon=True) self.learner_id = learner_id self.layout = make_layout() @@ -72,6 +74,9 @@ def __init__(self, learner_id: int, port: int, target: np.ndarray, store: Direct self.client = SyncerClient(("127.0.0.1", port), learner_id, self.layout, DTYPE_F32, num_streams=1, syncer_epoch=syncer_epoch, recover=recover_on_start) self.donor = DonorService(self.client, self, store) + # Island pool member: heartbeats keep its lease (also during a restore). + self.heartbeat = (None if lease_s is None + else LeaseHeartbeat(self.client, lease_s, progress=lambda: self.steps_total)) self.killed = threading.Event() self.outcome = None self.imported: dict[str, np.ndarray] | None = None @@ -131,6 +136,8 @@ def run(self) -> None: except BaseException as exc: # surfaced by the test self.exc = exc finally: + if self.heartbeat is not None: + self.heartbeat.stop() self.client.close() def _adam(self) -> None: @@ -145,6 +152,8 @@ def _adam(self) -> None: def _run(self) -> None: self.client.start() + if self.heartbeat is not None: + self.heartbeat.start() if self.recover_on_start: self.outcome = recover(self.client, self, self.store, timeout=60) self.after_recovery = self.export_recovery_state() @@ -183,10 +192,11 @@ def _run(self) -> None: time.sleep(0.004) -def start_syncer(binary: Path, port: int, learners: int, quorum: int, tape: Path, *extra: str) -> subprocess.Popen: +def start_syncer(binary: Path, port: int, learners: int, quorum: int, tape: Path, *extra: str, + total_steps: int = 80) -> subprocess.Popen: return subprocess.Popen( [str(binary), "--port", str(port), "--learners", str(learners), "--quorum", str(quorum), - "--grace-ms", "50", "--total-steps", "80", "--outer-lr", "0.7", "--outer-momentum", "0.0", + "--grace-ms", "50", "--total-steps", str(total_steps), "--outer-lr", "0.7", "--outer-momentum", "0.0", "--min-round-interval-ms", "40", "--quorum-timeout-s", "30", "--event-tape", str(tape), "--recovery-bound-steps", str(BOUND_H), *extra], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, @@ -311,6 +321,61 @@ def test_killed_island_recovers_from_a_peer_within_bound(tmp_path, transfer_s, e assert ("peer recovery missed its bound" in out) == expect_missed +@pytest.mark.timeout(240) +def test_island_pool_slot_filled_by_peer_recovery(tmp_path): + """Island pool and peer recovery together: the pool admits a new + generation for a lost island; its state comes from a peer. The restore + takes longer than the lease, so only heartbeats keep the island in the + pool; it has weight 0 (no pushes) until RECOVERY_DONE.""" + lease_s = 1.5 + binary = build_syncer() + port = free_port() + tape = tmp_path / "tape.jsonl" + store = RecordingStore(tmp_path / "store") + rng = np.random.default_rng(1) + targets = [rng.standard_normal(FRAGMENTS * NUMEL).astype(np.float32) for _ in range(3)] + proc = start_syncer(binary, port, 3, 2, tape, "--island-lease-s", str(lease_s), "--syncer-epoch", "5", + "--max-base-lag", "2", total_steps=120) + pool = {"syncer_epoch": 5, "lease_s": lease_s} + try: + islands = [FakeIsland(i, port, targets[i], store, **pool) for i in range(3)] + for island in islands: + island.start() + wait_until(lambda: min(islands[0].versions) >= 8, 60, "no early rounds") + islands[2].kill() + islands[2].join(timeout=10) + kill_versions = list(islands[0].versions) + wait_until(lambda: sum(islands[0].versions) > sum(kill_versions) + 8, 60, "rounds stopped without island 2") + back = FakeIsland(2, port, targets[2], SlowStore(store.root, 2 * lease_s, store.puts), + recover_on_start=True, **pool) + back.start() + wait_until(lambda: back.outcome is not None or back.exc is not None, 60, "recovery did not finish") + if back.exc: + raise back.exc + finish(proc, [islands[0], islands[1], back]) + except BaseException: + if proc.poll() is None: + proc.kill() + print((proc.stdout.read() if proc.stdout else "")[-4000:]) + raise + outcome = back.outcome + assert outcome.path == "peer" and outcome.donor in (0, 1), outcome + assert outcome.seconds > lease_s, "the restore outlasted the lease" + snapshot = store.puts[islands[outcome.donor].donor.served[0][1]] + for key in ("adam_m", "adam_v", "adam_step"): + np.testing.assert_array_equal(back.imported[key], snapshot.state[key], err_msg=key) + records = [json.loads(line) for line in tape.read_text().splitlines()] + generation = back.client.connection_generation + left = [r for r in records if r.get("kind") == "pool_leave" and r.get("island_id") == 2 + and r.get("generation") == generation] + assert not left, "heartbeats kept the recovering island's lease" + joined = [r for r in records if r.get("kind") == "pool_join" and r.get("recovered")] + assert [(r["island_id"], r["generation"]) for r in joined] == [(2, generation)] + assert back.heartbeat.errors == [] + assert back.pushes_after_recovery > 0 + assert back.final_manifest is not None and back.final_manifest == islands[0].final_manifest + + @pytest.mark.timeout(240) def test_recovery_still_in_transfer_at_the_end_of_training_takes_the_final_cut(tmp_path): binary = build_syncer() diff --git a/tests/test_peer_recovery_launch.py b/tests/test_peer_recovery_launch.py index 5523efcc9..2c41f4914 100644 --- a/tests/test_peer_recovery_launch.py +++ b/tests/test_peer_recovery_launch.py @@ -45,6 +45,35 @@ def test_explicit_store_and_ttl(): assert _flag(plan["role_syncers"][0]["command"], "--syncer-epoch") == "0" +POOL = ("--rl-sync-quorum-theta", "0.75", "--rl-sync-max-base-lag", "2", "--rl-sync-island-lease-s", "45") + + +def test_island_pool_takes_peer_recovery_with_one_epoch_flag(): + """The pool decides who fills a slot; peer recovery is that island's state source.""" + args, plan = _plan(RECOVERY + POOL) + assert args.rl_island_state_source == "peer" + command = plan["role_syncers"][0]["command"].split() + assert command.count("--syncer-epoch") == 1, "pool and recovery share one syncer epoch" + assert command[command.index("--syncer-epoch") + 1] == "4" + assert command[command.index("--island-lease-s") + 1] == "45.0" + assert command[command.index("--recovery-bound-steps") + 1] == "8" + for island in plan["island_requests"]: + words = island["learner_command"].split() + assert words.count("--rl-syncer-epoch") == 1 + assert words[words.index("--rl-syncer-epoch") + 1] == "4" + assert "--rl-island-pool" in words + assert _flag(island["learner_command"], "--rl-peer-recovery-store") == "/shared/run-store" + assert relaunch_envs(args) == {PEER_RECOVER_ENV: "1"} + + +def test_island_pool_with_catch_up_source_has_no_recovery(): + args, plan = _plan(ASYNC + POOL + ("--rl-island-state-source", "catch-up")) + assert args.rl_peer_recovery is False + assert "--recovery-bound-steps" not in plan["role_syncers"][0]["command"] + assert "--rl-peer-recovery-store" not in plan["island_requests"][0]["learner_command"] + assert relaunch_envs(args) == {} + + def test_without_peer_recovery_commands_are_unchanged(): _, plan = _plan(ASYNC) assert "--syncer-epoch" not in plan["role_syncers"][0]["command"] @@ -61,6 +90,11 @@ def test_without_peer_recovery_commands_are_unchanged(): (RECOVERY + ("--rl-recovery-max-donor-tries", "0"), "must be >= 1"), (ASYNC + ("--rl-syncer-epoch", "2"), "needs island pool options"), (RECOVERY + ("--syncer-shards", "2"), "one syncer process"), + (RECOVERY + ("--rl-island-state-source", "catch-up"), "contradicts --rl-peer-recovery"), + (ASYNC + ("--no-rl-peer-recovery", "--rl-island-state-source", "peer"), + "contradicts --no-rl-peer-recovery"), + (ASYNC + ("--rl-island-state-source", "cut"), "own round cuts"), + (ASYNC + ("--rl-island-state-source", "peer"), "--rl-peer-recovery-store or --rl-checkpoint-store"), ], ) def test_refusals(extra, message): diff --git a/tests/test_peer_recovery_miles.py b/tests/test_peer_recovery_miles.py index bb9a92dd2..415a406c2 100644 --- a/tests/test_peer_recovery_miles.py +++ b/tests/test_peer_recovery_miles.py @@ -193,8 +193,9 @@ def close(self): self.closed = True -def config(store, *, recover, timeout_s=10.0): +def config(store, *, recover, timeout_s=10.0, island_lease_s=None): return RoleStreamConfig( + island_lease_s=island_lease_s, algorithm="grpo", role="actor", component=ComponentIdentity(role="actor", model_revision="a" * 40, config_hash="c" * 64), syncer_addr=("127.0.0.1", 30100), learner_id=0, learner_generation=0, learner_generations=(0, 0), @@ -285,6 +286,30 @@ def test_donor_snapshot_and_recovering_restore_through_the_role_stream(tmp_path, assert DirectorySnapshotStore(store_root).locators() == [], "a restored snapshot is removed" +def test_island_pool_member_recovers_from_a_peer_while_it_heartbeats(tmp_path, stubs): + """An island pool fills a slot; the island's state still comes from a peer.""" + store_root = tmp_path / "store" + donor_client = FakeClient(requests=[rp.SnapshotRequest(1, 9, 5, (3, 4))]) + donor = donor_stream(store_root, FakeGroup(adam_step=17), donor_client) + assert asyncio.run(donor.serve_peer_snapshots()) == 1 + (_, versions, sha256, locator), = donor_client.sent + client = FakeClient(events=[rp.RecoveryClock(5, (3, 4), 1), + rp.RecoverySource(rp.SOURCE_PEER, 0, versions, sha256, locator), + rp.RecoveryDone(5, 6, 1, (5, 6))], + updates=[BcastFragment(0, 5, stored(0, 5, "g0"), stored=True), + BcastFragment(1, 6, stored(1, 6, "g1"), stored=True)]) + stubs.clients.append(client) + stream = MilesFullParameterRoleStream(config(store_root, recover=True, island_lease_s=30.0)) + asyncio.run(stream.initialize(FakeGroup(adam_step=0))) + try: + assert stubs.kwargs[0]["recover"] is True + assert stream.heartbeat is not None, "the lease is renewed during and after the restore" + assert stream.recovery_outcome.path == "peer" + assert stream.fragment_versions == (5, 6) + finally: + stream.heartbeat.stop() + + def test_other_tensor_parallel_layout_is_refused(tmp_path, stubs): store_root = tmp_path / "store" donor_client = FakeClient(requests=[rp.SnapshotRequest(1, 9, 5, (3, 4))]) diff --git a/tests/test_rl_sync_mode.py b/tests/test_rl_sync_mode.py index 4b0d89ed3..b7ee5ba41 100644 --- a/tests/test_rl_sync_mode.py +++ b/tests/test_rl_sync_mode.py @@ -220,8 +220,16 @@ def test_peer_recovery_defaults_on_where_it_applies(): args, _ = _plan(FULL) # no snapshot store: off assert args.rl_peer_recovery is False # Off: the store alone is then refused, which shows recovery stayed off. - for extra in (("--no-rl-peer-recovery",), ("--rl-sync-mode", "strict"), POOL): + for extra in (("--no-rl-peer-recovery",), ("--rl-island-state-source", "catch-up"), + ("--rl-sync-mode", "strict")): args = _launcher_args("ports", FULL + store + extra, gpu=TWO_ISLANDS) with pytest.raises(ValueError, match="or --rl-peer-recovery"): _prepare_rl_args(args) assert args.rl_peer_recovery is False + + +def test_peer_recovery_defaults_on_in_an_island_pool(): + args, _ = _plan(FULL + ("--rl-checkpoint-store", "/shared/run-store") + POOL) + assert args.rl_peer_recovery is True and args.rl_island_state_source == "peer" + args, _ = _plan(FULL + POOL) # no snapshot store: the syncer catch-up + assert args.rl_peer_recovery is False and args.rl_island_state_source == "catch-up" diff --git a/yeto/launch/peer_recovery.py b/yeto/launch/peer_recovery.py index 19f97e1e4..9314c9eff 100644 --- a/yeto/launch/peer_recovery.py +++ b/yeto/launch/peer_recovery.py @@ -1,6 +1,10 @@ """Launcher side of peer recovery (S27 E-3, ``--rl-peer-recovery``). -Full-parameter asynchronous rounds only. With the option on: +Full-parameter asynchronous rounds only, with a fixed roster or an island +pool. The pool decides which island fills a slot; this option decides the +state source of that island (``--rl-island-state-source peer``; off is +``catch-up``: every fragment from the syncer, fresh optimizer state). With the +option on: * every role syncer gets ``--syncer-epoch`` and the recovery options (:func:`yeto.launch.syncer.peer_recovery_syncer_flags`); @@ -28,8 +32,11 @@ def peer_recovery_enabled(args) -> bool: def peer_recovery_store(args) -> str | None: """The snapshot store path on the island, or None when recovery is off.""" - if not peer_recovery_enabled(args): - return None + return snapshot_store_path(args) if peer_recovery_enabled(args) else None + + +def snapshot_store_path(args) -> str | None: + """The snapshot store path the island would use, whether or not recovery is on.""" explicit = getattr(args, "rl_peer_recovery_store", None) if explicit: return str(explicit) @@ -42,7 +49,7 @@ def peer_recovery_store(args) -> str | None: def check_peer_recovery(args) -> None: """Refuse ``--rl-peer-recovery`` options that do not apply. Every recovery option without ``--rl-peer-recovery`` is refused too.""" - from .syncer import island_pool, rl_full_parameter_async, syncer_shard_count + from .syncer import rl_full_parameter_async, syncer_shard_count options = [flag for dest, flag in ( ("rl_recovery_bound_steps", "--rl-recovery-bound-steps"), @@ -58,9 +65,6 @@ def check_peer_recovery(args) -> None: if not rl_full_parameter_async(args): raise ValueError("--rl-peer-recovery needs --rl-sync-preset full-parameter with " "--rl-sync-full-parameter-rounds async") - if island_pool(args): - # Both set the syncer epoch; recovery is only wired for a fixed roster. - raise ValueError("--rl-peer-recovery does not take island pool options yet") if syncer_shard_count(args) > 1: # Recovery is one exchange with one syncer; shards would each need it. raise ValueError("--rl-peer-recovery needs one syncer process (--syncer-shards 1)") @@ -87,9 +91,11 @@ def island_flags(args) -> str: store = peer_recovery_store(args) if store is None: return "" - from .syncer import island_syncer_epoch + from .syncer import island_pool, island_syncer_epoch - flags = f" --rl-syncer-epoch {island_syncer_epoch(args)} --rl-peer-recovery-store {shlex.quote(store)}" + # An island pool already sends the epoch with its membership flags. + flags = "" if island_pool(args) else f" --rl-syncer-epoch {island_syncer_epoch(args)}" + flags += f" --rl-peer-recovery-store {shlex.quote(store)}" ttl = getattr(args, "rl_peer_recovery_ttl_s", None) if ttl is not None: flags += f" --rl-peer-recovery-ttl-s {float(ttl)!r}" diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 83b16c4f0..4ffb64915 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -405,24 +405,43 @@ def _apply_rl_sync_mode(args) -> None: ) +ISLAND_STATE_SOURCES = ("peer", "catch-up", "cut") + + def _default_peer_recovery(args) -> None: - """``--rl-peer-recovery`` left unset: on in the async mode where it applies - (full-parameter async rounds, a fixed roster, one syncer process, no - broadcast deltas, a snapshot store), off otherwise; the strict mode and - ``--no-rl-peer-recovery`` keep it off.""" - if getattr(args, "rl_peer_recovery", None) is not None: - return - args.rl_peer_recovery = False + """Resolve the state source of a replacement island (``--rl-island-state-source``) + and ``--rl-peer-recovery`` together; the two always agree afterwards. + + Who fills a slot (a relaunch, an island pool member) is not decided here; this + only picks where that island takes its trainer state. ``--rl-peer-recovery`` + left unset: on in the async mode where it applies (full-parameter async rounds, + one syncer process, no broadcast deltas, a snapshot store; a fixed roster or an + island pool), off otherwise; the strict mode and ``--no-rl-peer-recovery`` keep + it off. Peer recovery on is source ``peer``, off is ``catch-up``.""" + source = getattr(args, "rl_island_state_source", None) + explicit = getattr(args, "rl_peer_recovery", None) + if source == "cut": + raise ValueError("--rl-island-state-source cut: a synced island does not resume from its " + "own round cuts yet (the syncer is authoritative); use peer or catch-up") + if source is not None and explicit is not None and explicit != (source == "peer"): + flag = "--rl-peer-recovery" if explicit else "--no-rl-peer-recovery" + raise ValueError(f"--rl-island-state-source {source} contradicts {flag}") + if source is not None: + args.rl_peer_recovery = source == "peer" + elif explicit is None: + args.rl_peer_recovery = _peer_recovery_applies(args) + args.rl_island_state_source = "peer" if args.rl_peer_recovery else "catch-up" + + +def _peer_recovery_applies(args) -> bool: if (getattr(args, "rl_sync_mode_effective", "strict") != "async" - or not rl_full_parameter_async(args) or island_pool(args) + or not rl_full_parameter_async(args) or syncer_shard_count(args) > 1 or getattr(args, "rl_sync_bcast_delta_full_every", 0)): - return - from .peer_recovery import peer_recovery_store + return False + from .peer_recovery import snapshot_store_path - args.rl_peer_recovery = True - if peer_recovery_store(args) is None: - args.rl_peer_recovery = False + return snapshot_store_path(args) is not None def _apply_rl_outer_merge_options(args, *, num_islands: int) -> None: diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index ae4ffb43a..fe9d3f6d5 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -506,7 +506,8 @@ def peer_recovery_syncer_flags(args) -> str: options (checked by ``launch.peer_recovery``); "" when recovery is off.""" if not getattr(args, "rl_peer_recovery", False): return "" - flags = f" --syncer-epoch {island_syncer_epoch(args)}" + # An island pool already passes the epoch with its membership flags. + flags = "" if island_pool(args) else f" --syncer-epoch {island_syncer_epoch(args)}" for dest, flag in (("rl_recovery_bound_steps", "--recovery-bound-steps"), ("rl_recovery_max_donor_tries", "--recovery-max-donor-tries"), ("rl_recovery_donor_timeout_s", "--recovery-donor-timeout-s")): diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 703c46dff..00f2797ad 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -298,8 +298,6 @@ def __post_init__(self) -> None: raise ValueError("peer recovery needs asynchronous rounds without broadcast deltas") if self.peer_recovery is not None and self.syncer_shards: raise ValueError("peer recovery needs one syncer process, not a sharded head") - if self.peer_recovery is not None and self.island_lease_s is not None: - raise ValueError("peer recovery is wired for a fixed roster, not an island pool") if self.wait_timeout <= 0 or not 0 < self.poll_seconds <= 1: raise ValueError("role-stream wait bounds are invalid") if ( diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index a8dff483a..49e9001f7 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -700,9 +700,18 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. help="full-parameter async: an island relaunched after it lost its trainer " "state restores the inner optimizer state from a peer island's snapshot, " "then takes every fragment from the syncer. Default: on in --rl-sync-mode " - "async where it applies (full-parameter async rounds, a fixed island roster, " - "one syncer process, no broadcast deltas, a snapshot store from " - "--rl-peer-recovery-store or --rl-checkpoint-store), else off"), + "async where it applies (full-parameter async rounds, one syncer process, " + "no broadcast deltas, a snapshot store from --rl-peer-recovery-store or " + "--rl-checkpoint-store; a fixed roster or an island pool), else off. Same as " + "--rl-island-state-source peer (on) or catch-up (off)"), + cli_flag("--rl-island-state-source", choices=("peer", "catch-up", "cut"), default=None, + help="full-parameter async: where an island that replaces a lost island (a " + "relaunch, or a new member of an island pool) takes its trainer state. peer: " + "parameters and inner optimizer state from a healthy peer, then the newer " + "fragments from the syncer (--rl-peer-recovery); catch-up: every fragment " + "from the syncer, fresh optimizer state; cut: the island's own latest round " + "cut (not implemented for synced islands yet, refused). Default: peer where " + "--rl-peer-recovery applies, else catch-up"), cli_flag("--rl-recovery-bound-steps", type=int, default=None, metavar="H", help="--rl-peer-recovery: bound on one recovery in syncer steps; the syncer " "records a recovery past it as missed (default: measured, not checked)"), From 7e3e8d1b926240f246fa39766825811fc9eac3e5 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:17:33 +0000 Subject: [PATCH 158/233] syncer: send rejoin catch-ups one fragment per scheduler pass A rejoining learner received every owned fragment inside the scheduler's event handler. For a 27B/30B model that is ~110 GB of f32 and took 11-13 minutes, during which no event was read and no round deadline fired: a 300 s grace window committed after 23 min (s27-gh-a), and queued pushes were dropped by a quorum-timeout restart once the catch-up ended (s27-x6-gi-a). The HELLO and the replacement-disconnect paths now queue a catch-up. The scheduler sends one fragment per pass after serving round actions and polls events without sleeping while a catch-up is pending, so a deadline waits for at most one fragment transfer. The tail (record_catch_up, pool records, repull or rebind) runs once the last fragment is out; durability is still settled before every send. A learner joins no new round before its catch-up finishes, and a superseded generation's catch-up is dropped. Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 394 +++++++++++++++++++++++++++++++------------ 1 file changed, 284 insertions(+), 110 deletions(-) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index f8ea31afa..b34747aeb 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -2763,6 +2763,7 @@ async fn scheduler( let mut cutoff_established = false; let mut first_budget_report: Option<(Member, u64)> = None; let mut recovery = PeerRecovery::new(cfg.recovery.clone()); + let mut catch_ups: CatchUps> = CatchUps::new(); if dynamic { write_pool_status(&cfg, ®istry, st.global_step, "running"); } @@ -2877,7 +2878,8 @@ async fn scheduler( if busy_fragments.contains(&p) { break; } - let groups = current_groups(®istry); + let mut groups = current_groups(®istry); + groups.retain(|group| !catch_ups.holds(group.member)); if groups.is_empty() || (fixed_roster && groups.len() < cfg.quorum as usize) { next_launch = Instant::now() + Duration::from_millis(100); break; @@ -3075,6 +3077,100 @@ async fn scheduler( } i += 1; } + // Advance one rejoin catch-up by one fragment (or run its tail). + // The round deadlines above were served first, so a long catch-up + // never holds a grace window open. + for member in catch_ups.drop_superseded(|member| is_current_member(®istry, member)) { + warn!( + learner_id = member.learner_id, + generation = member.generation, + "superseded learner catch-up dropped" + ); + } + match catch_ups.next_step() { + None => {} + Some(CatchUpStep::Send(member, group, p)) => { + // Catch-up must never expose a round whose checkpoint is not + // yet durable. + if let Some(pending) = publishing.take() { + settle_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } + let dtype = bulk_dtype(st.wire_dtype); + match send_fragment_until_cutoff(&st, &group, p, MSG_BCAST_FRAGMENT, dtype, &budget_cutoff) + .await + { + Ok(true) => {} + Ok(false) => { + cutoff_established = true; + break 'outer; + } + Err(error) => { + warn!( + learner_id = member.learner_id, + generation = member.generation, + fragment = p, + %error, + "recovery fragment encode failed" + ); + catch_ups.abandon(member); + } + } + } + Some(CatchUpStep::Done(done)) => { + let group = done.group; + if done.tail.rejoin { + registry + .lock() + .unwrap() + .membership + .record_catch_up(group.member, &st.versions); + if dynamic { + append_pool_record( + cfg.event_tape.as_deref(), + "pool_join", + cfg.syncer_epoch, + &format!( + "\"island_id\":{},\"generation\":{},\"catch_up\":true,\"base_version\":{}", + group.member.learner_id, group.member.generation, st.global_step + ), + ); + write_pool_status(&cfg, ®istry, st.global_step, "running"); + } + if !repull_rounds_for(&group, ®istry, &mut inflight, fixed_roster, &budget_cutoff).await { + cutoff_established = true; + break 'outer; + } + } + if done.tail.rebind { + let repulls = rebind_current_unanswered_rounds(®istry, &mut inflight, group.member); + for repull in repulls { + if !send_small_until_cutoff(&group, MSG_PULL_REQ, repull.pull, &budget_cutoff).await { + cutoff_established = true; + break 'outer; + } + info!( + learner_id = group.member.learner_id, + old_generation = repull.old_generation, + generation = group.member.generation, + step = repull.t, + fragment = repull.p, + attempt = repull.attempt, + base_version = repull.base_version, + "rebound outstanding pull after captured generation disconnected" + ); + } + } + } + } if submitted_any { continue; // refill the pipeline before waiting again } @@ -3112,6 +3208,10 @@ async fn scheduler( if next_fragment_available { earliest = Some(earliest.map_or(next_launch, |d| d.min(next_launch))); } + if !catch_ups.is_empty() { + // Poll events without sleeping, then send the next fragment. + earliest = Some(deadline_now); + } let wake = if let Some(pending) = publishing.as_mut() { // Same event-first bias; a durable checkpoint outranks compute // so the oldest committed round is exposed first. @@ -3401,49 +3501,14 @@ async fn scheduler( continue; } - // Rejoining learner: first catch it up to the current + // Rejoining learner: catch it up to the current // parameters, then atomically rebind only its unanswered - // slots in existing rounds. The helper revalidates the - // generation because catch-up awaits socket writers and - // a newer HELLO may supersede this one meanwhile. - // Catch-up must never expose a round whose checkpoint - // is not yet durable. - if let Some(pending) = publishing.take() { - settle_pending_publish( - &cfg, - &st, - ®istry, - &mut last_sync_secs, - pending, - &budget_cutoff, - &mut busy_fragments, - ) - .await?; - } - if st.global_step < last_step { - send_all_fragments(&st, &group, &budget_cutoff).await; - } - registry - .lock() - .unwrap() - .membership - .record_catch_up(group.member, &st.versions); - if dynamic { - append_pool_record( - cfg.event_tape.as_deref(), - "pool_join", - cfg.syncer_epoch, - &format!( - "\"island_id\":{},\"generation\":{},\"catch_up\":true,\"base_version\":{}", - group.member.learner_id, group.member.generation, st.global_step - ), - ); - write_pool_status(&cfg, ®istry, st.global_step, "running"); - } - if !repull_rounds_for(&group, ®istry, &mut inflight, fixed_roster, &budget_cutoff).await { - cutoff_established = true; - break 'outer; - } + // slots in existing rounds. The scheduler sends the + // catch-up one fragment per pass (see `CatchUps`) and + // revalidates the generation before each send, because a + // newer HELLO may supersede this one meanwhile. + let fragments = if st.global_step < last_step { st.owned_fragments() } else { 0..0 }; + catch_ups.enqueue(group.member, group, fragments, CatchUpTail::REJOIN); } Event::Recovery { member, msg } => { let restored = match msg { @@ -3604,50 +3669,11 @@ async fn scheduler( if let Some(group) = group { // Catch-up must precede a replay even when the // disconnect event wins the race with the queued - // replacement HELLO. - if let Some(pending) = publishing.take() { - settle_pending_publish( - &cfg, - &st, - ®istry, - &mut last_sync_secs, - pending, - &budget_cutoff, - &mut busy_fragments, - ) - .await?; - } - if st.global_step < last_step { - send_all_fragments(&st, &group, &budget_cutoff).await; - } - let repulls = rebind_current_unanswered_rounds( - ®istry, - &mut inflight, - replacement, - ); - for repull in repulls { - if !send_small_until_cutoff( - &group, - MSG_PULL_REQ, - repull.pull, - &budget_cutoff, - ) - .await - { - cutoff_established = true; - break 'outer; - } - info!( - learner_id = replacement.learner_id, - old_generation = repull.old_generation, - generation = replacement.generation, - step = repull.t, - fragment = repull.p, - attempt = repull.attempt, - base_version = repull.base_version, - "rebound outstanding pull after captured generation disconnected" - ); - } + // replacement HELLO: the rebind runs only once + // the catch-up has finished. + let fragments = + if st.global_step < last_step { st.owned_fragments() } else { 0..0 }; + catch_ups.enqueue(replacement, group, fragments, CatchUpTail::REBIND); } } } @@ -5656,30 +5682,109 @@ impl FinalAbsence { } } -async fn send_all_fragments(st: &GlobalState, group: &Arc, budget_cutoff: &BudgetCutoff) { - for p in st.owned_fragments() { - match send_fragment_until_cutoff( - st, - group, - p, - MSG_BCAST_FRAGMENT, - bulk_dtype(st.wire_dtype), - budget_cutoff, - ) - .await - { - Ok(true) => {} - Ok(false) => return, - Err(error) => { - warn!( - learner_id = group.member.learner_id, - generation = group.member.generation, - fragment = p, - %error, - "recovery fragment encode failed" - ); - return; +/// What a finished catch-up still owes its learner. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +struct CatchUpTail { + /// HELLO path: record the catch-up and repull the learner's open rounds. + rejoin: bool, + /// Disconnect path: move the captured generation's unanswered slots to + /// the already-current replacement. + rebind: bool, +} + +impl CatchUpTail { + const REJOIN: Self = Self { rejoin: true, rebind: false }; + const REBIND: Self = Self { rejoin: false, rebind: true }; +} + +/// One rejoined learner's catch-up: the current value of every owned fragment. +struct CatchUp { + member: Member, + group: G, + remaining: std::collections::VecDeque, + tail: CatchUpTail, +} + +enum CatchUpStep { + /// Send this fragment's current value to the learner. + Send(Member, G, usize), + /// Every fragment went out (or the stream failed); run the tail. + Done(CatchUp), +} + +/// Rejoin catch-ups owned by the scheduler. +/// +/// A catch-up is the whole model (about 110 GB of f32 for a 27B model). The +/// scheduler sends one fragment per pass and then serves events and round +/// deadlines again, so a catch-up delays a grace or quorum deadline by at +/// most one fragment transfer instead of by the whole model (s27-gh-a: a +/// 300 s grace window committed after 14 min). A learner does not join new +/// rounds before its catch-up is done: its pushes would be against stale +/// bases. +struct CatchUps { + queue: std::collections::VecDeque>, +} + +impl CatchUps { + fn new() -> Self { + Self { queue: std::collections::VecDeque::new() } + } + + fn is_empty(&self) -> bool { + self.queue.is_empty() + } + + fn holds(&self, member: Member) -> bool { + self.queue.iter().any(|catch_up| catch_up.member == member) + } + + /// Queue a catch-up; a second request for the same generation only adds + /// its tail, because the fragments already sent are current. + fn enqueue( + &mut self, + member: Member, + group: G, + fragments: impl IntoIterator, + tail: CatchUpTail, + ) { + if let Some(existing) = self.queue.iter_mut().find(|catch_up| catch_up.member == member) { + existing.tail.rejoin |= tail.rejoin; + existing.tail.rebind |= tail.rebind; + return; + } + self.queue.push_back(CatchUp { member, group, remaining: fragments.into_iter().collect(), tail }); + } + + /// Remove catch-ups whose generation is no longer current. + fn drop_superseded(&mut self, is_current: impl Fn(Member) -> bool) -> Vec { + let mut dropped = Vec::new(); + self.queue.retain(|catch_up| { + let keep = is_current(catch_up.member); + if !keep { + dropped.push(catch_up.member); + } + keep + }); + dropped + } + + /// Stop sending to `member` (its stream failed); the tail still runs. + fn abandon(&mut self, member: Member) { + if let Some(catch_up) = self.queue.iter_mut().find(|catch_up| catch_up.member == member) { + catch_up.remaining.clear(); + } + } + + /// The next step, round-robin over the queued learners. + fn next_step(&mut self) -> Option> { + let mut catch_up = self.queue.pop_front()?; + match catch_up.remaining.pop_front() { + Some(p) => { + let step = CatchUpStep::Send(catch_up.member, catch_up.group.clone(), p); + self.queue.push_back(catch_up); + Some(step) } + None => Some(CatchUpStep::Done(catch_up)), } } } @@ -9169,6 +9274,75 @@ mod tests { Arc::new(Mutex::new(RegistryState::new(Membership::new(cfg, 2)))) } + fn catch_up_sends(catch_ups: &mut CatchUps<()>) -> Vec<(u32, usize)> { + let mut sent = Vec::new(); + while let Some(step) = catch_ups.next_step() { + match step { + CatchUpStep::Send(member, (), p) => sent.push((member.learner_id, p)), + CatchUpStep::Done(_) => {} + } + } + sent + } + + #[test] + fn catch_up_yields_one_fragment_per_step_then_its_tail() { + let rejoined = member(1, 7); + let mut catch_ups = CatchUps::new(); + catch_ups.enqueue(rejoined, (), 0..3, CatchUpTail::REJOIN); + assert!(catch_ups.holds(rejoined)); + for expected in 0..3 { + match catch_ups.next_step() { + Some(CatchUpStep::Send(who, (), p)) => assert_eq!((who, p), (rejoined, expected)), + _ => panic!("fragment {expected} was not the next step"), + } + } + match catch_ups.next_step() { + Some(CatchUpStep::Done(done)) => { + assert_eq!(done.member, rejoined); + assert_eq!(done.tail, CatchUpTail::REJOIN); + } + _ => panic!("finished catch-up did not yield its tail"), + } + assert!(catch_ups.is_empty() && !catch_ups.holds(rejoined)); + } + + #[test] + fn catch_up_interleaves_learners_and_merges_duplicate_requests() { + let (a, b) = (member(1, 7), member(2, 9)); + let mut catch_ups = CatchUps::new(); + catch_ups.enqueue(a, (), 0..2, CatchUpTail::REJOIN); + catch_ups.enqueue(b, (), 0..2, CatchUpTail::REJOIN); + // The disconnect of a's old generation asks again: no second copy. + catch_ups.enqueue(a, (), 0..2, CatchUpTail::REBIND); + assert_eq!(catch_up_sends(&mut catch_ups), vec![(1, 0), (2, 0), (1, 1), (2, 1)]); + + catch_ups.enqueue(a, (), 0..1, CatchUpTail::REJOIN); + catch_ups.enqueue(a, (), 0..1, CatchUpTail::REBIND); + let _ = catch_ups.next_step(); + match catch_ups.next_step() { + Some(CatchUpStep::Done(done)) => { + assert_eq!(done.tail, CatchUpTail { rejoin: true, rebind: true }) + } + _ => panic!("merged catch-up lost its tail"), + } + } + + #[test] + fn catch_up_drops_superseded_generations_and_abandons_failed_streams() { + let (old, current) = (member(1, 7), member(2, 9)); + let mut catch_ups = CatchUps::new(); + catch_ups.enqueue(old, (), 0..4, CatchUpTail::REJOIN); + catch_ups.enqueue(current, (), 0..4, CatchUpTail::REJOIN); + assert_eq!(catch_ups.drop_superseded(|m| m == current), vec![old]); + assert!(!catch_ups.holds(old)); + + // A failed stream stops the sends; the tail still runs. + catch_ups.abandon(current); + assert!(matches!(catch_ups.next_step(), Some(CatchUpStep::Done(done)) if done.member == current)); + assert!(catch_ups.next_step().is_none()); + } + #[tokio::test] async fn lease_drop_rejoin_and_stale_generation_fencing() { let registry = leased_registry(1); From 968aaa395bc2c047fd24577e02e57f243fb45ebe Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:17:33 +0000 Subject: [PATCH 159/233] wan_streams: start an idle stream's stall clock at its enqueue StreamMeter kept the time of the last byte a stream ever sent. When an idle stream got a chunk, the stall watchdog could read it before the sender thread called begin() and saw minutes of "no progress". The first push after training dropped the whole connection group three times (s27-gh-a: 2206 s; s27-x6-gi-a: 717 s and 3143 s, timeout 300 s), and each reconnect cost a full f32 catch-up. Co-Authored-By: Claude Opus 5.5 --- tests/test_wan_streams.py | 15 +++++++++++++++ yeto/wan_streams.py | 5 +++++ 2 files changed, 20 insertions(+) diff --git a/tests/test_wan_streams.py b/tests/test_wan_streams.py index 99f3d4dc1..7ed096b70 100644 --- a/tests/test_wan_streams.py +++ b/tests/test_wan_streams.py @@ -200,6 +200,21 @@ def test_watchdog_notices_an_inbound_message_that_stops_growing(): assert client._stalled_stream() is None +def test_meter_stall_clock_starts_at_the_enqueue_of_an_idle_stream(): + # s27-gh-a: a stream idle for 2206 s got a chunk; the watchdog read it + # before the sender thread picked it up and dropped the group as stalled. + clock = _Clock() + meter = StreamMeter(clock) + clock.now += 2206 + meter.enqueued(10) + assert meter.stalled_for() == 0.0 + clock.now += 4 + assert meter.stalled_for() == pytest.approx(4) + # Bytes already waiting keep their clock: more bytes are no progress. + meter.enqueued(10) + assert meter.stalled_for() == pytest.approx(4) + + def test_stall_timeout_sets_tcp_user_timeout(): if not hasattr(socket, "TCP_USER_TIMEOUT"): pytest.skip("Linux-only socket option") diff --git a/yeto/wan_streams.py b/yeto/wan_streams.py index 8d2cb6abb..cb65757fe 100644 --- a/yeto/wan_streams.py +++ b/yeto/wan_streams.py @@ -71,6 +71,11 @@ def __init__(self, clock=time.monotonic) -> None: def enqueued(self, nbytes: int) -> None: with self._lock: + if nbytes > 0 and self._item_started is None and self.queued_bytes <= 0: + # An idle stream starts its stall clock at the enqueue, not + # at the last byte it sent minutes ago (s27-gh-a: a stream + # idle for 2206 s was dropped as stalled on its next chunk). + self._last_progress = self._clock() self.queued_bytes += nbytes def begin(self) -> None: From 6157577dbb90df3a352a3ccb86c763a4a0652e68 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:17:33 +0000 Subject: [PATCH 160/233] rl islands: size the Modal disk and stop duplicating local-init anchors Full-parameter islands keep whole-model f32 cuts in the Ray object store and spill the rest to the container disk. On Modal's default disk the 27B and 30B-A3B smokes spilled 318-350 GB and failed with ENOSPC. - Local init: the version-0 anchors adopt the exported version-0 cut's Ray refs (AuthoritativeFragmentSink.adopt) instead of storing a second copy of the same bytes, and the digest is streamed instead of built in a fragment-sized bytearray. - --modal-ephemeral-disk-gib sets the island container disk. Unset, full-parameter RL sizes it from the model: 5 f32 copies + the bf16 snapshot + 64 GiB, clamped to 512..3072 GiB (634 GiB for 27B). - --rl-island-ray-object-store-gib and --rl-island-ray-spill-dir reach ray start (store size on every node; spill config on the head). Co-Authored-By: Claude Opus 5.5 --- tests/test_modal_fn_volume.py | 36 +++++++++++++ tests/test_modal_runner.py | 8 +++ tests/test_rl_launcher.py | 36 +++++++++++++ tests/test_rl_local_init_real_syncer.py | 8 ++- tests/test_rl_miles_chunked_full_parameter.py | 44 ++++++++++++++++ tests/test_run_flags.py | 5 +- yeto/cli.py | 9 ++++ yeto/launch/modal_islands.py | 40 +++++++++++++- yeto/launch/rl_island_task.py | 7 ++- yeto/launch/topology.py | 26 ++++++++++ yeto/modal_runner.py | 7 +++ .../miles/full_parameter/role_stream.py | 13 ++--- .../miles/models/chunked_full_parameter.py | 52 +++++++++++++++++++ yeto/rl/engine/run_flags.py | 9 ++++ 14 files changed, 288 insertions(+), 12 deletions(-) diff --git a/tests/test_modal_fn_volume.py b/tests/test_modal_fn_volume.py index d523d87ff..56f9eefb4 100644 --- a/tests/test_modal_fn_volume.py +++ b/tests/test_modal_fn_volume.py @@ -86,3 +86,39 @@ def test_hostmem_cgroup_v2_and_v1_and_rss(tmp_path): assert mr._proc_rss_sum(str(proc)) == 2048 * 1024 rec = mr._host_mem_used_bytes() # this machine: at least one memory source is non-null assert rec["proc_rss_sum"] and (rec["cgroup_current"] is not None or rec["meminfo_used"]) + + +def test_full_parameter_island_disk_scales_with_the_model(monkeypatch): + import argparse + + import pytest + + from yeto.launch import modal_islands + + def args(**kw): + base = dict(training_mode="rl", rl_sync_preset="full-parameter", + model="qwen36-27b", modal_ephemeral_disk_gib=None) + base.update(kw) + return argparse.Namespace(**base) + + # 27B: 54 GB bf16 -> 5 f32 copies + the snapshot + margin. + assert modal_islands.modal_ephemeral_disk_gib(args()) == 618 + assert modal_islands.modal_ephemeral_disk_gib(args(model="qwen3-8b")) == 512 # floor + assert modal_islands.modal_ephemeral_disk_gib(args(model="deepseek4flash")) == 3072 # cap + assert modal_islands.modal_ephemeral_disk_gib(args(modal_ephemeral_disk_gib=900)) == 900 + assert modal_islands.modal_ephemeral_disk_gib(args(rl_sync_preset="lora")) is None + with pytest.raises(ValueError): + modal_islands.modal_ephemeral_disk_gib(args(modal_ephemeral_disk_gib=0)) + + def unknown(*_a, **_k): + raise ValueError("offline") + + monkeypatch.setattr("yeto.shape.memory.model_weights_gb", unknown) + assert modal_islands.modal_ephemeral_disk_gib(args()) is None # Modal's default + + +def test_ephemeral_disk_validates(): + import pytest + + with pytest.raises(ValueError): + _cfg(ephemeral_disk_gib=0).validate() diff --git a/tests/test_modal_runner.py b/tests/test_modal_runner.py index 610f48835..5fb577be3 100644 --- a/tests/test_modal_runner.py +++ b/tests/test_modal_runner.py @@ -235,6 +235,7 @@ def test_define_builds_the_function_from_the_config(monkeypatch): # step after it (this order failed the first live Modal deploy). assert names == ["debian_slim", "pip_install", "env", "add_local_dir"] assert img.calls[names.index("add_local_dir")][1][1] == mr.CONTAINER_WORKDIR + assert "ephemeral_disk" not in kwargs # Modal's default disk # Unpinned single-container island: no region kwarg, no clustered. state2 = fake_modal(monkeypatch) mr.ModalOps("yeto-run").define(_cfg(learner_id=2)) @@ -645,3 +646,10 @@ def test_resolve_syncer_for_modal_sharded_head_list(): import pytest with pytest.raises(ValueError): resolve_syncer_for_modal("3.231.26.103:29400,10.0.0.1:29401", None) + + +def test_define_sizes_the_container_disk_in_mib(monkeypatch): + state = fake_modal(monkeypatch) + mr.ModalOps("yeto-run").define(_cfg(ephemeral_disk_gib=634)) + _fn, kwargs = state["functions"]["island-1"] + assert kwargs["ephemeral_disk"] == 634 * 1024 diff --git a/tests/test_rl_launcher.py b/tests/test_rl_launcher.py index 8add8173f..6e319bfc3 100644 --- a/tests/test_rl_launcher.py +++ b/tests/test_rl_launcher.py @@ -970,6 +970,42 @@ def test_miles_tasks_mount_the_same_attested_initial_adapter( ) +def test_island_ray_object_store_and_spill_dir_reach_ray_start(): + import json + import shlex + + from yeto.gpu_spec import parse_gpu_spec + + args = _args(("--gpu", "aws:2x4xa100@us-east-1", + "--rl-island-ray-object-store-gib", "200", + "--rl-island-ray-spill-dir", "/mnt/big disk/spill")) + args.model_revision = "a" * 40 + args.data_revision = "b" * 40 + args.source_sha256 = "c" * 64 + args.reward_sha256 = "d" * 64 + _prepare_rl_args(args) + run = make_miles_island_task(args, parse_gpu_spec(args.gpu)[0], 0, 1, "127.0.0.1:29400").run + _assert_island_ray_is_private(run) + store = f"--object-store-memory={200 << 30}" + head = run[run.index("ray start --head"):].split("\n", 1)[0] + worker = run[run.index('ray start --address="$MASTER_ADDR:6379"'):].split("&&", 1)[0] + assert store in head and store in worker + # The spill dir is a cluster setting: head only, created before Ray starts. + (config,) = [word.split("=", 1)[1] for word in shlex.split(head) + if word.startswith("--system-config=")] + assert "--system-config" not in worker + spilling = json.loads(json.loads(config)["object_spilling_config"]) + assert spilling == {"type": "filesystem", "params": {"directory_path": "/mnt/big disk/spill"}} + assert run.index("mkdir -p '/mnt/big disk/spill'") < run.index("ray start --head") + + plain = _args(("--gpu", "aws:2x4xa100@us-east-1")) + plain.model_revision, plain.data_revision = "a" * 40, "b" * 40 + plain.source_sha256, plain.reward_sha256 = "c" * 64, "d" * 64 + _prepare_rl_args(plain) + run = make_miles_island_task(plain, parse_gpu_spec(plain.gpu)[0], 0, 1, "127.0.0.1:29400").run + assert "--object-store-memory" not in run and "--system-config" not in run + + def _assert_island_ray_is_private(run: str) -> None: """The island's Ray must coexist with SkyPilot's runtime Ray on the node. diff --git a/tests/test_rl_local_init_real_syncer.py b/tests/test_rl_local_init_real_syncer.py index 34078e73a..b7f14dc81 100644 --- a/tests/test_rl_local_init_real_syncer.py +++ b/tests/test_rl_local_init_real_syncer.py @@ -231,7 +231,13 @@ def test_role_stream_keeps_local_version_zero_after_the_confirm(monkeypatch): adapter = _fake_adapter(payloads) stream = module.MilesFullParameterRoleStream( replace(_role_stream_config("actor", 30100), local_init=True)) - local, digests = stream._local_version_zero(adapter, object(), lambda fid, v, p: _Stored(fid, bytes(p))) + class AdoptingSink: + def adopt(self, cut, fragment_id, version, wire_payload_hash): + assert version == 0 + assert wire_payload_hash == hashlib.sha256(payloads[fragment_id]).hexdigest() + return _Stored(fragment_id, cut[fragment_id]) + + local, digests = stream._local_version_zero(adapter, payloads, AdoptingSink()) assert [s.payload for s in local] == payloads assert digests == [hashlib.sha256(p).digest() for p in payloads] diff --git a/tests/test_rl_miles_chunked_full_parameter.py b/tests/test_rl_miles_chunked_full_parameter.py index 5cd7885f7..325b4bc9f 100644 --- a/tests/test_rl_miles_chunked_full_parameter.py +++ b/tests/test_rl_miles_chunked_full_parameter.py @@ -598,6 +598,50 @@ async def test_grpo_h1_wrapper_preserves_versioned_policy_binding(): assert receipt.optimizer_steps == 1 +def test_local_init_anchor_shares_the_version_zero_refs_and_equals_a_wire_copy(): + import hashlib + + ray = FakeRay() + adapter = MilesChunkedFullParameterAdapter.create( + (manifest(0, (32,)), manifest(1, (32,))), + algorithm="grpo", + components=(component(),), + minimum_fragments=2, + max_fragment_bytes=256, + max_chunk_bytes=16, + ) + initial = adapter.validate_cut( + _cut(adapter, {0: 3, 1: 5}, policy_version=0, generation=0, ray=ray), 0, 0 + ) + sink = adapter.fragment_sink(initial, ray_module=ray) + payload = b"".join(adapter.fragment_parts(initial, 1, ray_module=ray)()) + wire_hash = hashlib.sha256(payload).hexdigest() + objects_before = len(ray.objects) + + adopted = sink.adopt(initial, 1, 0, wire_hash) + assert len(ray.objects) == objects_before # no second model copy + assert sink.adopt(initial, 1, 0, wire_hash) is adopted + with pytest.raises(RuntimeError, match="changed bytes"): + sink.adopt(initial, 1, 0, "0" * 64) + + # The same bytes stored from the wire give an identical anchor. + copied = adapter.fragment_sink(initial, ray_module=ray)(1, 0, payload) + assert copied.descriptor == adopted.descriptor + assert copied.wire_payload_hash == adopted.wire_payload_hash + assert [ray.get(r).tolist() for r in copied.refs] == [ray.get(r).tolist() for r in adopted.refs] + + # Releasing the cut leaves the anchor's refs; releasing the anchor + # leaves the cut's. + other = sink.adopt(initial, 0, 0, "f" * 64) + other.release() + assert all(r is not None for r in adopted.refs) + assert all(r is not None for refs in initial.transport_cut.shards[0].chunk_refs for r in refs) + adapter.release(initial) + assert all(r is not None for r in adopted.refs) + with pytest.raises(ValueError, match="released"): + adapter.fragment_sink(initial, ray_module=ray).adopt(initial, 1, 0, wire_hash) + + def test_inbound_chunks_commit_directly_to_exact_typed_miles_refs(): ray = FakeRay() adapter = MilesChunkedFullParameterAdapter.create( diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index 8068db135..62863dbf9 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -261,13 +261,14 @@ def test_forward_tables_name_flags_of_both_layers(): "--manifest-path", "--mark-final-checkpoint", "--master", "--matrix-merge", "--max-base-lag", "--shard-count", "--shard-index", "--weighted-fold", "--name", "--net", "--nnodes", "--no-build-isolation", "--no-checkout", "--no-deps", "--no-follow", "--node", - "--node-ip-address", "--nproc", "--num-gpus", "--num-inference-steps", "--output-dir", + "--node-ip-address", "--nproc", "--num-gpus", "--num-inference-steps", "--object-store-memory", + "--output-dir", "--password-stdin", "--policy-sweep-fragments", "--porcelain", "--port", "--profile", "--prompt", "--q-min", "--query-gpu", "--quiet", "--quorum-theta", "--quorum-timeout-s", "--raid-devices", "--recovery-bound-steps", "--recovery-donor-timeout-s", "--recovery-max-donor-tries", "--release", "--require-profile-binding", "--resume", "--rl-eval", "--rl-test", "--rm", "--run", - "--runtime", "--seed-column", "--shm-size", "--soft-deadline-s", "--syncer-epoch", + "--runtime", "--seed-column", "--shm-size", "--soft-deadline-s", "--syncer-epoch", "--system-config", "--temp-dir", "--ulimit", "--untracked-files", "--wandb-group", "--yes", # torchrun, tail and the W4A8 re-quantizer CLI in the SSH trainer island (yeto/launch/ssh_trainer_island.py) "--awq-scales", "--loop", "--master-addr", "--master-port", "--no-awq", "--node-rank", "--nproc-per-node", diff --git a/yeto/cli.py b/yeto/cli.py index 1aaf34ce0..4ba41308e 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -564,6 +564,15 @@ def int_or_auto(value: str): metavar="GIB", help="Modal islands: container memory request = limit in GiB (default 32 GiB per GPU)", ) + infra.add_argument( + "--modal-ephemeral-disk-gib", + type=int, + default=None, + metavar="GIB", + help="Modal islands: container disk in GiB (Ray spills its object store here). " + "Default: full-parameter RL sizes it from the model (5 f32 model copies + the " + "bf16 snapshot + 64 GiB, 512..3072 GiB); other runs keep Modal's default", + ) infra.add_argument( "--modal-cpu", type=int, diff --git a/yeto/launch/modal_islands.py b/yeto/launch/modal_islands.py index 9d71765e5..28068ccaa 100644 --- a/yeto/launch/modal_islands.py +++ b/yeto/launch/modal_islands.py @@ -5,6 +5,7 @@ from __future__ import annotations import json +import math import os import sys import time @@ -14,7 +15,7 @@ from . import REPO_ROOT from .code_pack import code_workdir from .remote_setup import HF_TOKEN_PATH -from .syncer import critic_syncer_address, island_hmac_secret +from .syncer import critic_syncer_address, island_hmac_secret, rl_dense_sweep from .credentials import SECRET_ENV_NAMES, _secret_value, registry_login_for from .storage import ( MODAL_CHECKPOINT_STORE_MOUNT, MODEL_STORE_MOUNT, _rl_checkpoint_mount, @@ -124,6 +125,8 @@ def build_modal_island_config(args, spec: ClusterSpec, learner_id: int, task, sy model_volume_mount=MODEL_STORE_MOUNT if rl and getattr(args, "modal_model_volume", None) else None, cpu_override=getattr(args, "modal_cpu", None), memory_gib_override=getattr(args, "modal_memory_gib", None), + ephemeral_disk_gib=modal_ephemeral_disk_gib(args) if rl else getattr( + args, "modal_ephemeral_disk_gib", None), tape_subdir=(modal_tape_subdir(args.cluster_prefix, learner_id) if getattr(args, "modal_tape_volume", None) else None), workdir=code_workdir(args), @@ -143,6 +146,41 @@ def build_modal_island_config(args, spec: ClusterSpec, learner_id: int, task, sy ) +# A full-parameter island holds whole-model f32 cuts as Ray objects: its +# anchors, the current cut, the overlap round base, staged broadcasts and a +# reconnect catch-up. What does not fit in the object store spills to the +# container disk; s27-gh-a (27B, Modal's default disk) spilled 318 GB and +# failed with ENOSPC. Size the disk for this many f32 copies. +FULL_PARAMETER_SPILL_COPIES = 5 +MODAL_EPHEMERAL_DISK_MARGIN_GIB = 64 +MODAL_EPHEMERAL_DISK_MIN_GIB = 512 +MODAL_EPHEMERAL_DISK_MAX_GIB = 3072 + + +def modal_ephemeral_disk_gib(args) -> int | None: + """Container disk of a Modal RL island: ``--modal-ephemeral-disk-gib``, or + for full-parameter sync a size scaled to the model; None keeps Modal's + default (other modes, or a model whose size is unknown).""" + explicit = getattr(args, "modal_ephemeral_disk_gib", None) + if explicit is not None: + if int(explicit) < 1: + raise ValueError("--modal-ephemeral-disk-gib must be >= 1") + return int(explicit) + if not rl_dense_sweep(args): + return None + from ..shape.memory import model_weights_gb + + try: + bf16_gb = model_weights_gb(args.model, getattr(args, "weights_gb", None)) + except ValueError as exc: + print(f"[yeto] Modal island disk: model size unknown ({exc}); keeping Modal's " + "default disk, pass --modal-ephemeral-disk-gib to size it", file=sys.stderr) + return None + bf16_gib = bf16_gb * 1e9 / (1 << 30) + need = (FULL_PARAMETER_SPILL_COPIES * 2 + 1) * bf16_gib + MODAL_EPHEMERAL_DISK_MARGIN_GIB + return min(max(math.ceil(need), MODAL_EPHEMERAL_DISK_MIN_GIB), MODAL_EPHEMERAL_DISK_MAX_GIB) + + def modal_tape_subdir(cluster_prefix: str, learner_id: int) -> str: """Per-run, per-island directory of the Modal tape Volume.""" from ..modal_runner import modal_app_name diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index cd1917c2c..97c76089d 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -19,7 +19,8 @@ ) from .topology import ( multinode_env_prelude, multinode_network_tier, rl_cross_node_flags, rl_island_bundle_map_flag, - rl_island_cpus, rl_island_gpu_env, rl_island_ray_gpus, rl_island_spec, rl_trainer_shape, + rl_island_cpus, rl_island_gpu_env, rl_island_ray_gpus, rl_island_ray_spill_dir_prelude, + rl_island_ray_storage, rl_island_spec, rl_trainer_shape, ) from .syncer import ( critic_syncer_address, elastic_debug_delay_env, elastic_debug_pause_env, island_hmac_secret, @@ -528,6 +529,7 @@ def make_miles_island_task( # command line, so pkill by that path never touches sky's. f"{rl_island_gpu_env(args, spec)}" 'MILES_RAY_DIR="$HOME/miles-ray"\n' + f"{rl_island_ray_spill_dir_prelude(args)}" # Same escalation as `ray stop --force` (TERM, short grace, KILL) # but scoped to this dir: gcs_server, raylet and the autoscaler # monitor ignore a bare SIGTERM for 30 s or more. @@ -551,6 +553,7 @@ def make_miles_island_task( # 2026-10-03: ServerUnavailable in compute_ray_pin_head_options). '--port=6379 --include-dashboard=true --temp-dir="$MILES_RAY_DIR"' + rl_island_ray_gpus(args, spec) + + rl_island_ray_storage(args, head=True) + (" --dashboard-host=0.0.0.0" if spec.num_nodes > 1 else "") + "\n" f" {shell_mark('ray_ready')}\n" " trap stop_miles_ray EXIT\n" @@ -579,7 +582,7 @@ def make_miles_island_task( " for _ in $(seq 1 150); do " 'ray start --address="$MASTER_ADDR:6379" --temp-dir="$MILES_RAY_DIR"' f"{_worker_node_ip_flag(spec.cloud)}" - f"{rl_island_ray_gpus(args, spec)} && break; " + f"{rl_island_ray_gpus(args, spec)}{rl_island_ray_storage(args, head=False)} && break; " "sleep 2; done\n" " ray status --address=\"$MASTER_ADDR:6379\" >/dev/null 2>&1 " "|| { echo '[yeto-island] worker could not join the Ray head' >&2; exit 1; }\n" diff --git a/yeto/launch/topology.py b/yeto/launch/topology.py index a5894bc82..356335789 100644 --- a/yeto/launch/topology.py +++ b/yeto/launch/topology.py @@ -3,6 +3,8 @@ from __future__ import annotations import dataclasses +import json +import shlex from ..gpu_spec import DEVICE_FAMILIES, parse_gpu_spec from .rl_backend import _rl_backend_module @@ -194,6 +196,30 @@ def rl_island_gpu_env(args, spec) -> str: f'(CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES)"\n') +def rl_island_ray_storage(args, *, head: bool) -> str: + """``ray start`` flags for the island's object store size and spill dir. + + The spill directory is a cluster setting, so only the head gets it (its + raylets read it from the GCS); every node gets the store size.""" + flags = "" + store_gib = getattr(args, "rl_island_ray_object_store_gib", None) + if store_gib is not None: + if int(store_gib) < 1: + raise ValueError("--rl-island-ray-object-store-gib must be >= 1") + flags += f" --object-store-memory={int(store_gib) << 30}" + spill_dir = getattr(args, "rl_island_ray_spill_dir", None) + if head and spill_dir: + spilling = json.dumps({"type": "filesystem", "params": {"directory_path": spill_dir}}) + flags += " --system-config=" + shlex.quote(json.dumps({"object_spilling_config": spilling})) + return flags + + +def rl_island_ray_spill_dir_prelude(args) -> str: + """Create the spill dir before ``ray start`` ("" when unset).""" + spill_dir = getattr(args, "rl_island_ray_spill_dir", None) + return f"mkdir -p {shlex.quote(spill_dir)}\n" if spill_dir else "" + + def rl_island_ray_gpus(args, spec) -> str: island = rl_island_spec(args, spec) return "" if island is spec else f" --num-gpus={island.gpus_per_node}" diff --git a/yeto/modal_runner.py b/yeto/modal_runner.py index 161b35fef..d03e729c9 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -249,6 +249,9 @@ class ModalIslandConfig: # Host resource overrides (None = per-GPU defaults below). cpu_override: int | None = None memory_gib_override: int | None = None + # Container disk in GiB (None = Modal's default). Ray spills its object + # store here (see yeto.launch.modal_islands.modal_ephemeral_disk_gib). + ephemeral_disk_gib: int | None = None @property def function_name(self) -> str: @@ -273,6 +276,8 @@ def memory_request_mib(self) -> int: def validate(self) -> None: validate_modal_shape(self.gpu, self.gpus_per_node, self.num_nodes) + if self.ephemeral_disk_gib is not None and self.ephemeral_disk_gib < 1: + raise ValueError("ephemeral_disk_gib must be >= 1") if self.training_mode == "rl": if not self.image_ref or not (re.fullmatch(r"[^\s@]+@sha256:[0-9a-fA-F]{64}", self.image_ref) or ENGINE_BUILT_IMAGE_RE.fullmatch(self.image_ref)): @@ -1083,6 +1088,8 @@ def define(self, cfg: ModalIslandConfig): ) if cfg.region: kwargs["region"] = cfg.region + if cfg.ephemeral_disk_gib is not None: + kwargs["ephemeral_disk"] = cfg.ephemeral_disk_gib * 1024 # MiB volumes = {} if cfg.volume_name and cfg.volume_mount: volumes[cfg.volume_mount] = modal.Volume.from_name(cfg.volume_name, create_if_missing=True) diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 703c46dff..87b9f0c1c 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -769,15 +769,16 @@ async def serve_peer_snapshots(self) -> int: return served def _local_version_zero(self, adapter, initial, sink): - """Local init: store this island's own INIT bytes as its version-0 - anchors and hash them exactly as the syncer hashes INIT_PARAMS.""" + """Local init: this island's own version-0 cut is its version-0 + anchors (shared, not copied), hashed exactly as the syncer hashes + INIT_PARAMS.""" stored, digests = [], [] for fragment_id in range(adapter.layout.fragments.num_fragments): - payload = bytearray() + digest = hashlib.sha256() for part in adapter.fragment_parts(initial, fragment_id, ray_module=self.ray_module)(): - payload += part - digests.append(hashlib.sha256(payload).digest()) - stored.append(sink(fragment_id, 0, payload)) + digest.update(part) + digests.append(digest.digest()) + stored.append(sink.adopt(initial, fragment_id, 0, digest.hexdigest())) return tuple(stored), digests def _wait_for_initial_fragments( diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index 2995159b6..c1f1a56df 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -139,6 +139,58 @@ def __call__( self._stored[key] = stored return stored + def adopt( + self, + cut: ReferencedPolicyCut, + fragment_id: int, + version: int, + wire_payload_hash: str, + ) -> StoredAuthoritativeFragment: + """Register ``cut``'s own chunks of one fragment without copying them. + + Local init: the island's version-0 anchors are exactly its exported + version-0 cut. Sharing the cut's Ray refs keeps one model-sized copy + in the object store instead of two (s27-gh-a: 27B spilled ~165 GB + before the first round and later filled the container disk). The + anchor holds its own ref list, so releasing either side leaves the + other intact.""" + + if cut.layout_hash != self.reference.layout_hash: + raise ValueError("adopted cut layout changed") + key = (fragment_id, version) + previous = self._stored.get(key) + if previous is not None: + if previous.wire_payload_hash != wire_payload_hash: + raise RuntimeError("authoritative fragment replay changed bytes") + return previous + shard = _fragment_owner_shard(cut.transport_cut, fragment_id) + ((descriptor, refs),) = [ + (descriptor, refs) + for descriptor, refs in zip(shard.fragments, shard.chunk_refs, strict=True) + if descriptor.fragment_id == fragment_id + ] + expected = next( + descriptor + for descriptor in _fragment_owner_shard(self.reference.transport_cut, fragment_id).fragments + if descriptor.fragment_id == fragment_id + ) + if not _same_fragment_geometry(expected, descriptor): + raise ValueError("adopted fragment geometry changed") + if any(reference is None for reference in refs): + raise ValueError("adopted fragment chunks are released") + stored = StoredAuthoritativeFragment( + fragment_id=fragment_id, + version=version, + parameter_layout_hash=shard.parameter_layout_hash, + topology=shard.topology, + plan_hash=shard.plan_hash, + descriptor=descriptor, + refs=list(refs), + wire_payload_hash=wire_payload_hash, + ) + self._stored[key] = stored + return stored + def release_all(self) -> int: return sum(fragment.release() for fragment in self._stored.values()) diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index a8dff483a..1fcac6eab 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -418,6 +418,15 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "give the island only local GPUs 0..M-1 on every node (M <= N): ray start " "--num-gpus=M, CUDA_VISIBLE_DEVICES=0..M-1, and every layout/placement/gpu_pool " "rule uses M (e.g. m5 TP8 over 2x4 on 2x8xH100 machines); unset = all N"), + cli_flag("--rl-island-ray-object-store-gib", type=int, default=None, metavar="GIB", + help="RL island: Ray object store size per node in GiB (ray start " + "--object-store-memory); unset = Ray's default (30%% of memory, capped by " + "/dev/shm). Full-parameter islands keep whole-model f32 cuts there; what does " + "not fit spills to disk (see --rl-island-ray-spill-dir)"), + cli_flag("--rl-island-ray-spill-dir", default=None, metavar="PATH", + help="RL island: directory for Ray object spilling (default: the island's Ray " + "temp dir on the container/VM disk). Point it at the largest local disk; on " + "Modal size the disk with --modal-ephemeral-disk-gib"), cli_flag("--rl-node-inventory", default=None, metavar="PATH", help="rl-ssh-node-pool: YAML inventory of the nodes yeto reaches by SSH " "(per node: ssh target, optional jump host, device family, card, card count, " From ff3044e219ecc59b781276d66198fd0b9216770b Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:44:24 +0000 Subject: [PATCH 161/233] Miles cut restore: digest the param_layout like the save does A full-parameter round cut records param_layout in the digested state; the restore re-snapshot left it out, so every full-parameter resume was refused with a state_digest mismatch. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 15477af4af51551c91fed3cc34f46d438c1687f8) --- tests/test_rl_miles_cut_plugin.py | 21 +++++++++++++++++++++ yeto/rl/adapters/miles/cut_plugin.py | 21 +++++++++++++++------ 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/tests/test_rl_miles_cut_plugin.py b/tests/test_rl_miles_cut_plugin.py index 17cd4ee06..12f645c91 100644 --- a/tests/test_rl_miles_cut_plugin.py +++ b/tests/test_rl_miles_cut_plugin.py @@ -357,3 +357,24 @@ def test_rewound_trained_trainer_restores_the_cut_like_a_fresh_one(tmp_path): for p_r, p_l in zip(reference.optimizer.param_groups[0]["params"], live.optimizer.param_groups[0]["params"]): for key in ("exp_avg", "exp_avg_sq", "step"): assert torch.equal(reference.optimizer.state[p_r][key], live.optimizer.state[p_l][key]) + + +class _LayoutBackend(type(make_rank(0, full=True)._yeto_cut_backend)): + """A backend that records a ``param_layout`` like MilesCutBackend (R-5).""" + + def param_layout(self, actor, named): + return {"schema": "test", "params": {n: {"shape": list(p.shape)} for n, p in named}} + + +def test_full_parameter_restore_with_param_layout_matches_the_saved_digest(tmp_path): + # GPU s27-glf-l9-a-pb: the save added param_layout to the digested state, the + # restore re-snapshot did not, so every full-parameter resume was refused. + run = make_rank(seed=0, full=True) + run._yeto_cut_backend = _LayoutBackend(None) + train_step(run, _batches(1)[0]) + summary = save_cut_shard(run, directory=str(tmp_path), cut_id="c1") + fresh = make_rank(seed=5, full=True) + fresh._yeto_cut_backend = _LayoutBackend(None) + restored = restore_cut_shard(fresh, directory=str(tmp_path), files=[summary], cut_id="c1") + assert restored["state_digest"] == summary["state_digest"] + assert restored["components"] == summary["components"] diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index cd23b11a4..2fc2bde35 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -677,6 +677,19 @@ def resume_state(actor: Any) -> dict[str, Any]: RESUME_STATE = f"{_MODULE}.resume_state" +def _add_param_layout(snap: dict[str, Any], actor: Any, backend: Any, named: list) -> None: + """R-5: a full-parameter shard records its ``param_layout`` for the offline TP/PP/EP + merge. Save and restore both call this, so the restored state digest covers the same + keys as the saved one.""" + layout = getattr(backend, "param_layout", None) + if _tuning(backend, named) != "full" or not callable(layout): + return + try: + snap["param_layout"] = layout(actor, named) + except Exception as error: # the cut stays valid for resume; only the offline merge refuses it + snap["param_layout"] = {"error": f"{type(error).__name__}: {error}"} + + def _snapshot(actor: Any, backend: Any, named: list, *, on_device: bool = False) -> dict[str, Any]: """The shard state at this step boundary. ``on_device`` (E1): the adapter and optimizer tensors are device copies (device-to-device, no wait for the host); @@ -1123,12 +1136,7 @@ def _capture(actor: Any, *, cut_id: str, device_copy: bool = False) -> dict[str, torch.cuda.empty_cache() if snap is None: snap = _snapshot(actor, backend, named) - layout = getattr(backend, "param_layout", None) - if _tuning(backend, named) == "full" and callable(layout): # R-5: offline TP/PP/EP merge - try: - snap["param_layout"] = layout(actor, named) - except Exception as error: # the cut stays valid for resume; only the offline merge refuses it - snap["param_layout"] = {"error": f"{type(error).__name__}: {error}"} + _add_param_layout(snap, actor, backend, named) ready = None if on_device and _sleeps_after_cut(actor): snap = _stage(snap) # host views; the device copies are dropped here @@ -1487,6 +1495,7 @@ def _restore(actor: Any, *, directory: str, files: list[Mapping[str, Any]], cut_ backend.restore_rng(shard["rng"]) # last: nothing below may consume randomness with torch.no_grad(): snap = _snapshot(actor, backend, named) + _add_param_layout(snap, actor, backend, named) # the save digests include it rng_now = backend.capture_rng() summary = { "path": name, From 92dcdb1e4186109c3f48d3c0ae78bb47df5d515b Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:57:32 +0000 Subject: [PATCH 162/233] rl: commit write-only island cuts inline when the trainer reads the sample buffer A deferred cut holds the Modal Volume writer guard until the next cut, so the sample buffer reader skipped every reload in between and the synced trainer failed the round with no groups. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit ac30c8c5617b728de13a00572aee4576cfe96a14) --- tests/test_rl_island_round_cuts.py | 15 +++++++++++++++ yeto/rl/adapters/miles/entry.py | 16 +++++++++++++++- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/tests/test_rl_island_round_cuts.py b/tests/test_rl_island_round_cuts.py index 9d98bed73..86f357d96 100644 --- a/tests/test_rl_island_round_cuts.py +++ b/tests/test_rl_island_round_cuts.py @@ -259,3 +259,18 @@ def test_verl_refuses_the_timed_interval(): with pytest.raises(ValueError, match="not wired for verl"): resume_config(SimpleNamespace(rl_cut_auto_mtbf_s=60.0)) + + +def test_write_only_cuts_commit_inline_for_a_buffer_reader(): + from yeto.rl.adapters.miles.buffer_rollout import READER_FUNCTION_PATHS + from yeto.rl.adapters.miles.entry import _commit_cuts_inline_for_buffer_reader + + reader = SimpleNamespace(background_commit=True) + _commit_cuts_inline_for_buffer_reader( + reader, miles_args=SimpleNamespace(rollout_function_path=READER_FUNCTION_PATHS[0])) + assert reader.background_commit is False + other = SimpleNamespace(background_commit=True) + _commit_cuts_inline_for_buffer_reader( + other, miles_args=SimpleNamespace(rollout_function_path="miles.rollout.sglang_rollout.generate_rollout")) + assert other.background_commit is True + _commit_cuts_inline_for_buffer_reader(None, miles_args=SimpleNamespace()) diff --git a/yeto/rl/adapters/miles/entry.py b/yeto/rl/adapters/miles/entry.py index a0a62a7d6..5dbca1571 100644 --- a/yeto/rl/adapters/miles/entry.py +++ b/yeto/rl/adapters/miles/entry.py @@ -1061,10 +1061,24 @@ def _wire_round_cuts(driver, *, elastic, miles_args, algorithm, base_model_revis or config.get("allow_config_change", False)), next_lr=next_lr_for(miles_args), interval=_cut_interval_of(elastic, miles_args)) if config.get("write_only"): - return # S27 R-2: a synced trainer's buffer lane keeps its own position (S25 C3) + # S27 R-2: a synced trainer's buffer lane keeps its own position (S25 C3) + _commit_cuts_inline_for_buffer_reader(checkpoint, miles_args=miles_args) + return _wire_sample_buffer(checkpoint, driver=driver, miles_args=miles_args) +def _commit_cuts_inline_for_buffer_reader(checkpoint: Any, *, miles_args: Any) -> None: + """A write-only cut store on a trainer that reads the sample buffer: commit + each cut inline. A deferred cut holds the store writer guard until the next + cut, so the buffer reader skips every Volume reload in between, never sees + the sampler's new groups and fails the round (S27 L11 s27-gle9-b: "0/32 + groups after 900 s").""" + from .buffer_rollout import READER_FUNCTION_PATHS + + if checkpoint is not None and getattr(miles_args, "rollout_function_path", None) in READER_FUNCTION_PATHS: + checkpoint.background_commit = False + + def _cut_interval_of(wiring: Any, miles_args: Any) -> Any: """The resume wiring's interval, else a timed one (``--rl-cut-auto-mtbf-s`` on an ``--rl-elastic`` run), else None (every N rounds).""" From 1ae040fd5ed3ce666ec17900442b57eb2bb53d3a Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 04:17:13 +0000 Subject: [PATCH 163/233] rl: resume a relaunched synced trainer from its buffer lane and the store's releases A synced trainer reads its sample buffer lane, whose cursor log holds its data position; the restart no longer asks the data source for a cursor that the write-only cut does not record. A restarted island also restores the policy versions below its start version (within the policy-age window) from the store's release records, so groups generated with versions released by an earlier publisher process pass the policy-age check. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_publisher_resume.py | 94 ++++++++++++++++++++++++ yeto/rl/adapters/miles/sampler_island.py | 7 ++ yeto/rl/engine/driver.py | 40 +++++++++- yeto/rl/engine/policy_release.py | 25 +++++++ 4 files changed, 165 insertions(+), 1 deletion(-) create mode 100644 tests/test_rl_publisher_resume.py diff --git a/tests/test_rl_publisher_resume.py b/tests/test_rl_publisher_resume.py new file mode 100644 index 000000000..598d681e3 --- /dev/null +++ b/tests/test_rl_publisher_resume.py @@ -0,0 +1,94 @@ +"""S27 L11: a relaunched synced trainer (sample buffer lane) resumes. + +* its buffer lane cursor log, not the data source, holds its data position: + the restart restores no data cursor; +* it learns the versions an earlier publisher process released (the store's + policy releases), so groups generated with them pass the policy-age check. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +from yeto.rl.adapters.miles import buffer_rollout as br +from yeto.rl.adapters.miles import sampler_island as si +from yeto.rl.engine import sample_roles +from yeto.rl.engine.driver import IslandDriver +from yeto.rl.engine.merged_release import BufferLane +from yeto.rl.engine.policy_age import check_batch_ages +from yeto.rl.engine.policy_release import PolicyReleaseStore +from yeto.rl.engine.policy_versions import PublishedPolicies +from yeto.rl.engine.resume import CheckpointStore + + +def _release(store, version, policy_hash): + store.release(version, policy_hash, lambda d: (d / "w.bin").write_text(policy_hash)) + + +def _driver(*, limit, history=None, owner=None, rollout=None): + events = [] + fake = SimpleNamespace( + published_policies=PublishedPolicies(), published_policy_history=history, + data_position_owner=owner, rollout=rollout or SimpleNamespace(), + emit=lambda event, **fields: events.append((event, fields)), + _max_policy_age=lambda: limit) + return fake, events + + +def test_records_lists_the_kept_releases_in_seq_order(tmp_path): + store = PolicyReleaseStore(CheckpointStore(tmp_path)) + assert store.records() == [] + for version in range(4): + _release(store, version, f"h{version}") + assert [(r.seq, r.version, r.policy_hash) for r in store.records()] == [(3, 2, "h2"), (4, 3, "h3")] + + +def test_restart_restores_the_published_versions_in_the_window(): + shared = [] + rollout = SimpleNamespace(set_published_policies=lambda p: shared.append(p.to_dict())) + history = lambda: [(1, "h1"), (2, "h2"), (3, "h3"), (4, "old")] # noqa: E731 + fake, events = _driver(limit=2, history=history, rollout=rollout) + IslandDriver._restore_published_policies(fake, 4) + # v1 is outside the window; v4 is the start version, published again + assert fake.published_policies.hashes() == {2: "h2", 3: "h3"} + assert shared == [{"2": "h2", "3": "h3"}] + assert events == [("rl_published_policies_restored", {"rollout_id": 4, "versions": [2, 3]})] + group = SimpleNamespace(group_id="g128", policy_token="yeto:2:h2", policy_versions=None) + assert check_batch_ages([group], 4, 2, {**fake.published_policies.hashes(), 4: "h4"}) == [] + + +def test_no_history_or_first_rollout_restores_nothing(): + fake, events = _driver(limit=2) + IslandDriver._restore_published_policies(fake, 3) + fake2, events2 = _driver(limit=2, history=lambda: [(0, "h0")]) + IslandDriver._restore_published_policies(fake2, 0) + assert fake.published_policies.hashes() == fake2.published_policies.hashes() == {} + assert events == events2 == [] + + +def test_a_buffer_lane_restart_restores_no_data_cursor(): + seeks = [] + rollout = SimpleNamespace(seek_data_cursor=lambda c: seeks.append(c)) + fake, events = _driver(limit=2, owner="buffer_lane", rollout=rollout) + IslandDriver._restore_data_cursor(fake, 2) + assert seeks == [] + assert events == [("rl_data_cursor_restored", {"rollout_id": 2, "source": "buffer_lane"})] + + +def test_wiring_a_synced_trainer_sets_the_lane_owner_and_the_release_history(tmp_path): + store = PolicyReleaseStore(CheckpointStore(tmp_path)) + _release(store, 1, "h1") + _release(store, 2, "h2") + events = [] + for index in (0, 1): # the publisher and the other synced trainer + driver = SimpleNamespace(publisher=SimpleNamespace(publish=lambda s: None), + emit=lambda event, **fields: events.append(event), + rollout=SimpleNamespace(_in_executor=lambda fn, *a: None)) + args = SimpleNamespace(lora_rank=0, yeto_rl_sample_lane=BufferLane(index, 2), + yeto_rl_sample_buffer=sample_roles.SampleBufferConfig( + role="trainer", store=str(tmp_path))) + si.wire_policy_releases(driver, miles_args=args, actor_model=SimpleNamespace(), + runner=SimpleNamespace(run=lambda value: value)) + br._ACTIVE.pop(br.TRAINER, None) + assert driver.data_position_owner == "buffer_lane" + assert driver.published_policy_history() == [(1, "h1"), (2, "h2")] diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py index 504f168e3..ecd8b3178 100644 --- a/yeto/rl/adapters/miles/sampler_island.py +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -532,6 +532,13 @@ def origin(state: Any) -> dict[str, Any]: def on_gap(previous: int, version: int) -> None: driver.emit("rl_policy_release_gap", previous_version=previous, policy_version=version) + # S27 L11: the lane cursor log resumes a relaunched synced trainer, not + # the data source; and it accepts groups made with the versions the + # store's releases name (a publisher of an earlier process released them). + driver.data_position_owner = "buffer_lane" + releases = release_store(config, environ) + driver.published_policy_history = lambda: [ + (r.version, r.policy_hash) for r in releases.records() if r.derived_from is None] driver.publisher = GateVersionPublisher( driver.publisher, lambda v: in_executor(executor_set_gate_version, int(v)), version_of) publisher_island = is_publisher(lane.index, lane.count) diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index 8ef75437d..4525d3176 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -63,7 +63,7 @@ import math import sys import time -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass from pathlib import Path from typing import Any, Protocol @@ -517,6 +517,13 @@ def __init__( self.published_state: TrainableState | None = None # P0-S: published versions within the policy-age window (version -> hash) self.published_policies = PublishedPolicies() + # S27 L11: where a restarted island's data position comes from when it is not + # the rollout data source (e.g. "buffer_lane": a synced trainer's lane cursor + # log); None: the data source cursor is restored (_restore_data_cursor). + self.data_position_owner: str | None = None + # S27 L11: () -> ((version, policy hash), ...) published before this process + # started (e.g. the store's policy releases); None: none to restore. + self.published_policy_history: Callable[[], Iterable[tuple[int, str]]] | None = None self.at_safe_point = False self.last_applied_lrs: list[float] | None = None self.eval_overlap: EvalOverlap | None = None @@ -834,6 +841,30 @@ def _record_published_policy(self, version: int, policy_hash: str) -> None: if callable(setter): setter(self.published_policies) + def _restore_published_policies(self, start_rollout_id: int) -> None: + """S27 L11: a restarted island learns the versions published before it + started, so groups the samplers generated with them pass the policy-age + check. Only versions below the start version within the policy-age window + are restored; the start version and later ones are published again.""" + history = getattr(self, "published_policy_history", None) + if start_rollout_id <= 0 or not callable(history): + return + oldest = start_rollout_id - self._max_policy_age() + restored = {} + for version, policy_hash in history(): + version = int(version) + if oldest <= version < start_rollout_id and policy_hash: + restored[version] = str(policy_hash) + if not restored: + return + for version, policy_hash in sorted(restored.items()): + self.published_policies.record(version, policy_hash) + setter = getattr(self.rollout, "set_published_policies", None) + if callable(setter): + setter(self.published_policies) + self.emit("rl_published_policies_restored", rollout_id=start_rollout_id, + versions=sorted(restored)) + def _with_policy_labels(self, batch: RolloutBatchHandle) -> RolloutBatchHandle: """P0-S: name every group's generating policy by its label. @@ -1878,6 +1909,12 @@ def _restore_data_cursor(self, start_rollout_id: int) -> None: pool that reports cursors -> DriverError (no silent re-training).""" if start_rollout_id <= 0: return + owner = getattr(self, "data_position_owner", None) + if owner is not None: + # S27 L11: the island reads no data source (a synced trainer reads its + # buffer lane, whose cursor log resumes it); no data cursor to restore. + self.emit("rl_data_cursor_restored", rollout_id=start_rollout_id, source=owner) + return seek = getattr(self.rollout, "seek_data_cursor", None) from_cut = getattr(self.sync, "resume_data_cursor", None) cut_cursor = from_cut() if callable(from_cut) else None @@ -1978,6 +2015,7 @@ def _run(self) -> TrainableState: start = None try: start = self.sync.start(self) + self._restore_published_policies(start.rollout_id) if self.ledger is not None: self.ledger.rebase(start.rollout_id) self._restore_data_cursor(start.rollout_id) diff --git a/yeto/rl/engine/policy_release.py b/yeto/rl/engine/policy_release.py index 79757a5a1..68467d7de 100644 --- a/yeto/rl/engine/policy_release.py +++ b/yeto/rl/engine/policy_release.py @@ -279,6 +279,31 @@ def latest(self) -> PolicyRelease | None: return None return PolicyRelease.from_record(json.loads(path.read_text(encoding="utf-8"))) + def records(self) -> list[PolicyRelease]: + """The release records still in the store (pruning keeps the newest + ``keep``) and LATEST, in ``seq`` order. A restarted trainer restores its + published versions from them (S27 L11).""" + found: dict[int, PolicyRelease] = {} + raws: list[bytes | str] = [] + objects = self.store.objects + if objects is not None: + keys = [k for k in objects.list(f"{self.release_dir}/") + if k.endswith(f"/{RECORD}") or k == f"{self.release_dir}/{LATEST}"] + raws = [raw for raw in (objects.read(k) for k in keys) if raw is not None] + else: + self.store.reload() + if self.root.is_dir(): + paths = [p / RECORD for p in self.root.iterdir() + if p.is_dir() and _RELEASE_NAME.fullmatch(p.name)] + [self.root / LATEST] + raws = [p.read_bytes() for p in paths if p.is_file()] + for raw in raws: + try: + record = PolicyRelease.from_record(json.loads(raw)) + except (ValueError, TypeError, PolicyReleaseError): + continue # a torn or foreign record: not a release this trainer made + found[record.seq] = record + return [found[seq] for seq in sorted(found)] + def release(self, version: int, policy_hash: str, write: Callable[[Path], None], *, buffer_selected: int | None = None, weights: str = FULL, engine_checksum: str | None = None, From 3eb7a967cc4a741e3753bb492555bf7488b36e9a Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:18:03 +0000 Subject: [PATCH 164/233] Nemotron-H parameter rules and the verl Megatron GRPO plan Frozen vision parameters (C-RADIO encoder, projector) are excluded by name from the syncer layout and the trainable set; Mamba A_log/D/dt_bias and the router bias must stay float32; routed-expert coverage is checked for HF fused and Megatron per-expert names (512 experts, latent projections are dense). The verl Megatron plan follows the fork's Nemotron Nano v3 GRPO example and scales TP/EP down; VL models are refused because verl does not read freeze_vision_tower. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_hybrid_params.py | 152 ++++++++++++++++++ yeto/rl/adapters/verl/megatron_recipe.py | 192 +++++++++++++++++++++++ yeto/rl/hybrid_params.py | 174 ++++++++++++++++++++ 3 files changed, 518 insertions(+) create mode 100644 tests/test_rl_hybrid_params.py create mode 100644 yeto/rl/adapters/verl/megatron_recipe.py create mode 100644 yeto/rl/hybrid_params.py diff --git a/tests/test_rl_hybrid_params.py b/tests/test_rl_hybrid_params.py new file mode 100644 index 000000000..a5805e0b3 --- /dev/null +++ b/tests/test_rl_hybrid_params.py @@ -0,0 +1,152 @@ +"""Nemotron-H parameter rules: frozen vision, f32 Mamba vectors, expert coverage, +and the verl Megatron plan built from the official Nano v3 example.""" + +from __future__ import annotations + +import pytest + +from yeto.fragments import build_layout +from yeto.rl import hybrid_params as hp +from yeto.rl.adapters.verl import megatron_recipe as mr + +# Names as SGLang's nano_nemotron_vl loader and HF nemotron_h define them. +VL_NAMES = { + "vision_model.radio_model.model.blocks.0.attn.qkv.weight": 1280 * 3840, + "mlp1.1.weight": 20480 * 5120, + "language_model.backbone.embeddings.weight": 131072 * 64, + "language_model.backbone.layers.0.mixer.A_log": 128, + "language_model.backbone.layers.0.mixer.in_proj.weight": 64 * 64, + "language_model.backbone.layers.1.mixer.experts.up_proj": 512 * 8 * 4, + "language_model.backbone.layers.1.mixer.fc1_latent_proj.weight": 64 * 16, + "language_model.lm_head.weight": 131072 * 64, +} + + +def test_frozen_vision_names(): + assert hp.is_frozen_vision("vision_model.radio_model.x.weight") + assert hp.is_frozen_vision("module.module.vision_projection.encoder.linear_fc1.weight") + assert hp.is_frozen_vision("mlp1.0.weight") + assert not hp.is_frozen_vision("language_model.backbone.layers.1.mixer.experts.up_proj") + assert not hp.is_frozen_vision("decoder.layers.3.mlp.linear_fc1.weight") + + +def test_layout_excludes_frozen_vision(): + named = list(VL_NAMES.items()) + kept = hp.layout_named_numels(named, freeze_vision=True) + assert {n for n, _ in kept} == {n for n in VL_NAMES if not n.startswith(("vision_model.", "mlp1."))} + layout = build_layout(kept, 4, avg_name_regex=hp.HYBRID_AVG_NAME_REGEX) + names = set(layout.tensor_names()) + assert not any(hp.is_frozen_vision(n) for n in names) + assert len(names) == len(kept) + # A_log is vector-like: avg fragment + avg = layout.fragments[0] + assert "language_model.backbone.layers.0.mixer.A_log" in dict(avg.tensors) + assert hp.layout_named_numels(named, freeze_vision=False) == named + with pytest.raises(ValueError): + hp.layout_named_numels([("vision_model.a", 1)], freeze_vision=True) + hp.check_frozen_excluded([n for n, _ in kept]) + with pytest.raises(hp.HybridParamError): + hp.check_frozen_excluded(["mlp1.1.weight"]) + + +@pytest.mark.parametrize("name", [ + "backbone.layers.0.mixer.A_log", "backbone.layers.0.mixer.D", "backbone.layers.0.mixer.dt_bias", + "module.decoder.layers.2.mixer.A_log", "backbone.layers.1.mixer.gate.e_score_correction_bias", + "decoder.layers.1.mlp.router.expert_bias", +]) +def test_requires_f32(name): + assert hp.requires_f32(name) + + +def test_check_f32(): + assert not hp.requires_f32("backbone.layers.0.mixer.in_proj.weight") + assert not hp.requires_f32("backbone.layers.0.mixer.Dx") + hp.check_f32({"backbone.layers.0.mixer.A_log": "torch.float32", "x.weight": "torch.bfloat16"}) + with pytest.raises(hp.HybridParamError, match="A_log"): + hp.check_f32({"backbone.layers.0.mixer.A_log": "torch.bfloat16"}) + + +def test_expert_coverage_hf_and_mcore(): + cov = hp.check_expert_coverage({"backbone.layers.1.mixer.experts.up_proj": (512, 8, 4), + "backbone.layers.1.mixer.experts.down_proj": (512, 4, 8)}, num_experts=512) + assert cov.moe_layers == (1,) + with pytest.raises(hp.HybridParamError, match="511"): + hp.check_expert_coverage({"backbone.layers.1.mixer.experts.up_proj": (511, 8, 4)}, num_experts=512) + mcore = {f"decoder.layers.3.mlp.experts.{p}.weight{e}": (4, 4) for p in ("linear_fc1", "linear_fc2") + for e in range(8)} + assert hp.check_expert_coverage(mcore, num_experts=8).moe_layers == (3,) + mcore.pop("decoder.layers.3.mlp.experts.linear_fc2.weight5") + with pytest.raises(hp.HybridParamError, match=r"missing \[5\]"): + hp.check_expert_coverage(mcore, num_experts=8) + # latent projections and shared experts are dense, not routed experts + with pytest.raises(hp.HybridParamError, match="no routed"): + hp.check_expert_coverage({"backbone.layers.1.mixer.fc1_latent_proj.weight": (16, 64), + "backbone.layers.1.mixer.shared_experts.up_proj.weight": (8, 64)}, + num_experts=8) + + +def test_layer_counts(): + nano = "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME" # Nemotron-3 Nano config.json + assert hp.layer_counts(nano) == {"mamba": 23, "moe": 23, "attention": 6, "mlp": 0} + with pytest.raises(ValueError): + hp.layer_counts("MEX") + + +def test_tiny_hf_nemotron_h_real_names(): + torch = pytest.importorskip("torch") + transformers = pytest.importorskip("transformers") + if not hasattr(transformers, "NemotronHConfig"): + pytest.skip("transformers without nemotron_h") + cfg = transformers.NemotronHConfig( + vocab_size=128, hidden_size=64, num_hidden_layers=3, hybrid_override_pattern="ME*", + num_attention_heads=4, num_key_value_heads=2, head_dim=16, mamba_num_heads=4, mamba_head_dim=16, + ssm_state_size=8, n_groups=2, n_routed_experts=8, num_experts_per_tok=2, moe_intermediate_size=32, + moe_shared_expert_intermediate_size=32, moe_latent_size=16, intermediate_size=32, use_mamba_kernels=False) + model = transformers.NemotronHForCausalLM(cfg).to(torch.bfloat16) + named = dict(model.named_parameters()) + f32 = [n for n in named if hp.requires_f32(n)] + assert {n.rsplit(".", 1)[1] for n in f32} >= {"A_log", "D", "dt_bias"} + # a plain .to(bf16) breaks the rule; restoring the f32 vectors passes it + with pytest.raises(hp.HybridParamError): + hp.check_f32({n: str(p.dtype) for n, p in named.items()}) + for n in f32: + named[n].data = named[n].data.float() + hp.check_f32({n: str(p.dtype) for n, p in named.items()}) + cov = hp.check_expert_coverage({n: tuple(p.shape) for n, p in named.items()}, num_experts=8) + assert cov.moe_layers == (1,) + assert any("fc1_latent_proj" in n for n in named) + + +def _plan(**kw): + base = dict(model_path="/m", train_file="/t.parquet", val_file="/v.parquet", out_dir="/o") + return mr.NemotronMegatronPlan(**{**base, **kw}) + + +def test_megatron_plan_official_and_scaled(): + o = mr.build_overrides(_plan()) + assert "actor_rollout_ref.actor.megatron.tensor_model_parallel_size=8" in o + assert "actor_rollout_ref.actor.megatron.expert_model_parallel_size=8" in o + assert "+actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32" in o + assert "+actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True" in o + small = mr.build_overrides(_plan(n_gpus=2, tp=2, ep=2, rollout_tp=2, groups_per_round=8, samples_per_group=4)) + assert "trainer.n_gpus_per_node=2" in small + assert "actor_rollout_ref.rollout.n=4" in small + no_off = mr.build_overrides(_plan(offload=False)) + assert not any("optimizer_cpu_offload" in x for x in no_off) + + +@pytest.mark.parametrize("kw,match", [ + (dict(architecture="NemotronH_Omni_Reasoning_V3"), "视觉"), + (dict(n_gpus=2, tp=4, ep=2, rollout_tp=2), "TP"), + (dict(num_experts=128, ep=3, tp=1, n_gpus=3, rollout_tp=1), "EP"), + (dict(n_gpus=2, tp=2, ep=2, rollout_tp=4), "rollout"), +]) +def test_megatron_plan_refusals(kw, match): + with pytest.raises(mr.NemotronPlanError, match=match): + mr.validate(_plan(**kw)) + + +def test_image_commands_pin_mamba_kernels(): + cmds = mr.image_commands() + assert any("--extra megatron" in c for c in cmds) + assert any("causal-conv1d==" in c for c in cmds) and any("mamba-ssm==" in c for c in cmds) diff --git a/yeto/rl/adapters/verl/megatron_recipe.py b/yeto/rl/adapters/verl/megatron_recipe.py new file mode 100644 index 000000000..27c485ff3 --- /dev/null +++ b/yeto/rl/adapters/verl/megatron_recipe.py @@ -0,0 +1,192 @@ +"""verl Megatron full-parameter GRPO overrides for Nemotron-H models. + +Source of the plan: the fork's official example +``examples/grpo_trainer/run_nemotron_nano_v3_30b_a3b_megatron.sh`` (fork +acad9875): one node of 8 GPUs, train TP8 / PP1 / EP8 / ETP1, rollout vLLM +TP8, bf16, full recompute (uniform, 1 layer), router in fp32, parameters and +optimizer offloaded to the host (``optimizer_cpu_offload``, precision-aware +optimizer), Megatron-Bridge model loading. :class:`NemotronMegatronPlan` +keeps those choices and lets the card count shrink (TP = EP = rollout TP = +cards by default, as in the example). + +This is the plan for the bare verl trainer (``verl.trainer.main_ppo``). yeto's +own verl island (``config.py``) is LoRA on one card; full-parameter verl +islands in the multi-island syncer are not wired (GOAL B10). + +Frozen vision: verl's ``actor.freeze_vision_tower`` exists in the fork's +config (``verl/workers/config/actor.py``) but no code reads it, so a VL model +would train its vision part. :func:`validate` therefore refuses a VL model; +the name rule for the freeze is :mod:`yeto.rl.hybrid_params`. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass + +VL_ARCHITECTURES = ("NemotronH_Omni_Reasoning_V3",) + + +class NemotronPlanError(ValueError): + pass + + +@dataclass(frozen=True) +class NemotronMegatronPlan: + model_path: str + train_file: str + val_file: str + out_dir: str + architecture: str = "NemotronHForCausalLM" + num_experts: int = 128 + n_gpus: int = 8 + nnodes: int = 1 + tp: int = 8 + pp: int = 1 + ep: int = 8 + etp: int = 1 + cp: int = 1 + rollout_tp: int = 8 + offload: bool = True + groups_per_round: int = 32 + samples_per_group: int = 16 + max_prompt_len: int = 2048 + max_response_len: int = 2048 + lr: float = 1e-6 + total_steps: int = 3 + gpu_memory_utilization: float = 0.7 + reward_path: str = "" + reward_name: str = "compute_score" + experiment_name: str = "yeto-nemotron-verl" + + def to_dict(self) -> dict: + return asdict(self) + + +def validate(p: NemotronMegatronPlan) -> None: + if p.architecture in VL_ARCHITECTURES: + raise NemotronPlanError( + "verl 不冻结视觉部分（freeze_vision_tower 无代码读取）；VL 模型需先接入冻结规则") + world = p.n_gpus * p.nnodes + model_parallel = p.tp * p.pp * p.cp + if world % model_parallel: + raise NemotronPlanError(f"TP*PP*CP={model_parallel} 不整除卡数 {world}") + if p.num_experts % p.ep: + raise NemotronPlanError(f"EP={p.ep} 不整除专家数 {p.num_experts}") + # Megatron: EP x ETP must divide the expert-side world (TP*CP*DP folded per PP stage). + if (world // p.pp) % (p.ep * p.etp): + raise NemotronPlanError(f"EP*ETP={p.ep * p.etp} 不整除每个流水级的卡数 {world // p.pp}") + if p.rollout_tp > p.n_gpus or p.n_gpus % p.rollout_tp: + raise NemotronPlanError(f"rollout TP={p.rollout_tp} 必须整除每机卡数 {p.n_gpus}") + if p.groups_per_round <= 0 or p.samples_per_group <= 0 or p.total_steps <= 0: + raise NemotronPlanError("groups_per_round / samples_per_group / total_steps 必须为正") + + +def build_overrides(p: NemotronMegatronPlan) -> list[str]: + """Hydra overrides for ``verl.trainer.main_ppo`` (same keys as the example).""" + validate(p) + seq = p.max_prompt_len + p.max_response_len + meg = "actor_rollout_ref.actor.megatron" + tcfg = f"+{meg}.override_transformer_config" + ocfg = "+actor_rollout_ref.actor.optim.override_optimizer_config" + o = [ + "model_engine=megatron", + "algorithm.adv_estimator=grpo", + "algorithm.use_kl_in_reward=False", + "algorithm.kl_ctrl.kl_coef=0.0", + f"data.train_files={p.train_file}", + f"data.val_files={p.val_file}", + "data.prompt_key=prompt", + "data.return_raw_chat=True", + "data.truncation=left", + f"data.max_prompt_length={p.max_prompt_len}", + f"data.max_response_length={p.max_response_len}", + f"data.train_batch_size={p.groups_per_round}", + f"actor_rollout_ref.model.path={p.model_path}", + "actor_rollout_ref.model.trust_remote_code=True", + "actor_rollout_ref.model.use_fused_kernels=False", + "actor_rollout_ref.actor.use_kl_loss=False", + "actor_rollout_ref.actor.clip_ratio_low=0.2", + "actor_rollout_ref.actor.clip_ratio_high=0.28", + "actor_rollout_ref.actor.clip_ratio_c=10.0", + "actor_rollout_ref.actor.use_dynamic_bsz=True", + f"actor_rollout_ref.actor.ppo_max_token_len_per_gpu={seq}", + "actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2", + f"actor_rollout_ref.actor.ppo_mini_batch_size={p.groups_per_round}", + f"actor_rollout_ref.actor.optim.lr={p.lr!r}", + "actor_rollout_ref.actor.optim.weight_decay=0.1", + "actor_rollout_ref.actor.optim.clip_grad=1.0", + "actor_rollout_ref.actor.entropy_coeff=0", + "actor_rollout_ref.actor.loss_agg_mode=token-mean", + f"{meg}.param_offload={p.offload}", + f"{meg}.optimizer_offload={p.offload}", + f"{meg}.tensor_model_parallel_size={p.tp}", + f"{meg}.pipeline_model_parallel_size={p.pp}", + f"{meg}.context_parallel_size={p.cp}", + f"{meg}.expert_model_parallel_size={p.ep}", + f"{meg}.expert_tensor_parallel_size={p.etp}", + f"{meg}.use_mbridge=True", + f"{meg}.dtype=bfloat16", + f"{tcfg}.apply_rope_fusion=True", + f"{tcfg}.moe_router_dtype=fp32", + f"{tcfg}.recompute_method=uniform", + f"{tcfg}.recompute_granularity=full", + f"{tcfg}.recompute_num_layers=1", + "actor_rollout_ref.rollout.name=vllm", + "actor_rollout_ref.rollout.mode=async", + "actor_rollout_ref.rollout.dtype=bfloat16", + f"actor_rollout_ref.rollout.tensor_model_parallel_size={p.rollout_tp}", + f"actor_rollout_ref.rollout.gpu_memory_utilization={p.gpu_memory_utilization}", + "actor_rollout_ref.rollout.enable_chunked_prefill=True", + f"actor_rollout_ref.rollout.max_num_batched_tokens={seq}", + "actor_rollout_ref.rollout.temperature=1.0", + "actor_rollout_ref.rollout.top_p=1.0", + "actor_rollout_ref.rollout.top_k=-1", + "actor_rollout_ref.rollout.calculate_log_probs=True", + f"actor_rollout_ref.rollout.n={p.samples_per_group}", + "actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True", + f"actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu={seq}", + "actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True", + f"actor_rollout_ref.ref.log_prob_max_token_len_per_gpu={seq}", + "trainer.logger=[console,file]", + "trainer.project_name=yeto-nemotron", + f"trainer.experiment_name={p.experiment_name}", + f"trainer.n_gpus_per_node={p.n_gpus}", + f"trainer.nnodes={p.nnodes}", + "trainer.critic_warmup=0", + "trainer.val_before_train=False", + "trainer.test_freq=-1", + "trainer.save_freq=-1", + "trainer.resume_mode=disable", + f"trainer.total_training_steps={p.total_steps}", + "trainer.total_epochs=1000", + f"trainer.default_local_dir={p.out_dir}/ckpt", + "ray_kwargs.ray_init.runtime_env.py_executable=null", + ] + if p.offload: + o += [f"{ocfg}.optimizer_offload_fraction=1", + f"{ocfg}.overlap_cpu_optimizer_d2h_h2d=True", + f"{ocfg}.use_precision_aware_optimizer=True", + f"{ocfg}.optimizer_cpu_offload=True"] + if p.reward_path: + o += [f"reward.custom_reward_function.path={p.reward_path}", + f"reward.custom_reward_function.name={p.reward_name}"] + return o + + +def image_commands(max_jobs: int = 16) -> list[str]: + """Extra image layers on top of the verl uv venv: the megatron extra plus the + Mamba training kernels the extra does not list (the example installs them by + hand: causal-conv1d and mamba-ssm). Hopper only; mamba-ssm's own CUDA + selective-scan is skipped (Mamba-2 training uses its Triton kernels).""" + from .image import VENV_BIN + from .pins import VERL_PYTHON, VERL_ROOT + env = f"MAX_JOBS={max_jobs} TORCH_CUDA_ARCH_LIST=9.0" + return [ + f"cd {VERL_ROOT} && UV_PYTHON=/usr/local/bin/python{VERL_PYTHON} uv sync --frozen --all-packages" + " --extra vllm --extra megatron", + f"uv pip install --python {VENV_BIN}/python ninja packaging wheel setuptools", + f"{env} CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --python {VENV_BIN}/python --no-build-isolation" + " 'causal-conv1d==1.6.2.post1'", + f"{env} MAMBA_SKIP_CUDA_BUILD=TRUE MAMBA_FORCE_BUILD=TRUE uv pip install --python {VENV_BIN}/python" + " --no-build-isolation --no-deps 'mamba-ssm==2.3.2.post1'", + ] diff --git a/yeto/rl/hybrid_params.py b/yeto/rl/hybrid_params.py new file mode 100644 index 000000000..1a17f997e --- /dev/null +++ b/yeto/rl/hybrid_params.py @@ -0,0 +1,174 @@ +"""Parameter rules for hybrid Mamba-2 + MoE models (NVIDIA Nemotron-H family). + +Covers the text models (``NemotronHForCausalLM``: Nemotron-3 Nano 30B-A3B, +Super 120B-A12B) and the VL model (``NemotronH_Omni_Reasoning_V3``: +Nemotron-3.5 Super VL 120B-A12B, a C-RADIO vision encoder plus an ``mlp1`` +projector in front of a Nemotron-H language model). + +Three rules, each a pure function on parameter names so that the trainer +side, the cut writer and the syncer layout apply the same rule: + +* **Frozen vision.** With the vision part frozen, the vision encoder and the + projector take no optimizer state and no syncer fragment. The rule is by + name, independent of ``requires_grad``, so that a layout built before the + trainer freezes anything is already correct. Names (verified in SGLang + 0.5.20 ``models/nano_nemotron_vl.py`` load_weights and Megatron-LM + ``LLaVAModel``): HF ``vision_model.*``, ``mlp1.*``, ``sound*``; Megatron + ``vision_model.*``, ``vision_projection.*``. +* **f32 parameters.** Mamba-2 ``A_log``, ``D`` and ``dt_bias`` are kept in + f32 by Megatron-Core (``ssm/mamba_mixer.py``: "Keep A_log in fp32") and the + router score-correction bias is f32 in HF transformers + (``_keep_in_fp32_modules_strict``). A trainer that holds them in bf16 + changes the SSM decay and the routing, so :func:`check_f32` refuses it. +* **Expert layout.** Routed experts must be complete per MoE layer: HF keeps + them as one fused 3D tensor per projection (``mixer.experts.up_proj``, + leading dim = number of experts); Megatron grouped GEMM names one tensor per + expert (``decoder.layers.L.mlp.experts.linear_fc1.weightE``, the form + :mod:`yeto.rl.moe_layout` reads). The MoE latent projections + (``fc1_latent_proj`` / ``fc2_latent_proj``) and the shared expert are dense + parameters, not routed experts. + +Import-light: no torch import at module level. +""" + +from __future__ import annotations + +import re +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass + +from .moe_layout import is_expert_parameter + +# Wrapper prefixes in front of a model's own names (DDP / Float16Module / +# FSDP / HF composite models). +_WRAPPERS = re.compile(r"^(?:(?:module|_orig_mod|_fsdp_wrapped_module|model)\.)+") + +VISION_PREFIXES = ("vision_model.", "vision_projection.", "mlp1.", "sound_encoder.", "sound_projection.") + +# Mamba-2 mixer vectors and the router bias that must stay f32. +_F32 = re.compile( + r"(?:\.mixer\.(?:A_log|D|dt_bias)" # Mamba-2 (HF and Megatron) + r"|\.gate\.e_score_correction_bias" # HF router bias + r"|\.mlp\.router\.expert_bias)\Z" # Megatron router bias +) + +# HF fused experts: ``...layers.L.mixer.experts.{up_proj,down_proj}``. +_HF_EXPERTS = re.compile(r"(?:^|\.)layers\.(\d+)\.mixer\.experts\.(up_proj|down_proj)\Z") +_MCORE_EXPERT = re.compile(r"decoder\.layers\.(\d+)\.mlp\.experts\.(linear_fc1|linear_fc2)\.weight(\d+)\Z") + +# Vector-like hybrid tensors a caller may route to the avg fragment +# (``yeto.fragments.build_layout(avg_name_regex=...)``): Mamba vectors, conv +# kernels, norms and router biases. Whether RDA merging of these 1D tensors is +# worse than avg is not measured; this is the conservative choice. +HYBRID_AVG_NAME_REGEX = ( + r"(?:\.mixer\.(?:A_log|D|dt_bias|conv1d\.(?:weight|bias)|conv1d_weight|conv1d_bias|norm\.weight)" + r"|e_score_correction_bias|router\.expert_bias|norm\.weight|layernorm\.weight)\Z" +) + + +def strip_wrappers(name: str) -> str: + return _WRAPPERS.sub("", name) + + +def is_frozen_vision(name: str) -> bool: + """A vision-encoder / projector / audio parameter of a Nemotron VL model.""" + return strip_wrappers(name).startswith(VISION_PREFIXES) + + +def requires_f32(name: str) -> bool: + return _F32.search(strip_wrappers(name)) is not None + + +def trainable_names(names: Iterable[str], *, freeze_vision: bool) -> list[str]: + """Names that take optimizer state and syncer fragments, input order kept.""" + return [n for n in names if not (freeze_vision and is_frozen_vision(n))] + + +def layout_named_numels(named_numels: Sequence[tuple[str, int]], *, + freeze_vision: bool) -> list[tuple[str, int]]: + """``(name, numel)`` pairs for :func:`yeto.fragments.build_layout`, frozen + vision parameters removed. Refuses an input with nothing left to train.""" + kept = [(n, s) for n, s in named_numels if not (freeze_vision and is_frozen_vision(n))] + if not kept: + raise ValueError("every parameter is frozen: no language-model parameters in the layout") + return kept + + +class HybridParamError(ValueError): + """The trainer's parameters break a hybrid-model rule.""" + + +def check_f32(named_dtypes: Mapping[str, str]) -> None: + """Refuse f32-required parameters held in another dtype. + + ``named_dtypes`` maps names to dtype strings (``"torch.float32"`` or + ``"float32"``).""" + bad = sorted(n for n, d in named_dtypes.items() + if requires_f32(n) and str(d).replace("torch.", "") != "float32") + if bad: + raise HybridParamError(f"must stay float32 (Mamba A_log/D/dt_bias, router bias): {bad[:8]}" + + (f" ... and {len(bad) - 8} more" if len(bad) > 8 else "")) + + +def check_frozen_excluded(trainable: Iterable[str]) -> None: + """Refuse a trainable / optimizer parameter list that still holds vision parameters.""" + bad = sorted(n for n in trainable if is_frozen_vision(n)) + if bad: + raise HybridParamError(f"frozen vision parameters in the trainable set: {bad[:8]}") + + +@dataclass(frozen=True) +class ExpertCoverage: + num_experts: int + moe_layers: tuple[int, ...] + + +def check_expert_coverage(named_shapes: Mapping[str, Sequence[int]], *, num_experts: int) -> ExpertCoverage: + """Every MoE layer holds all ``num_experts`` routed experts, each once. + + Accepts HF fused names (shape ``[E, ...]``) and Megatron per-expert names + (global expert index in the name, e.g. after ``moe_layout`` rewriting). + A mix of both conventions is refused.""" + hf: dict[tuple[int, str], int] = {} + mcore: dict[tuple[int, str], set[int]] = {} + for raw, shape in named_shapes.items(): + name = strip_wrappers(raw) + if (m := _HF_EXPERTS.search(name)) is not None: + hf[(int(m[1]), m[2])] = int(shape[0]) + elif (m := _MCORE_EXPERT.search(name)) is not None: + assert is_expert_parameter(name) + key = (int(m[1]), m[2]) + expert = int(m[3]) + if expert in mcore.setdefault(key, set()): + raise HybridParamError(f"expert {expert} of {key} appears twice") + mcore[key].add(expert) + if hf and mcore: + raise HybridParamError("HF fused and Megatron per-expert names mixed in one layout") + if not hf and not mcore: + raise HybridParamError("no routed-expert parameters found") + layers: set[int] = set() + for (layer, proj), count in hf.items(): + if count != num_experts: + raise HybridParamError(f"layer {layer} {proj}: {count} experts, expected {num_experts}") + layers.add(layer) + for (layer, proj), experts in mcore.items(): + if experts != set(range(num_experts)): + missing = sorted(set(range(num_experts)) - experts)[:8] + raise HybridParamError(f"layer {layer} {proj}: {len(experts)} of {num_experts} experts, missing {missing}") + layers.add(layer) + projections = {p for _, p in (hf or mcore)} + for layer in layers: + for proj in projections: + if (layer, proj) not in (hf or mcore): + raise HybridParamError(f"layer {layer} has no {proj} experts") + return ExpertCoverage(num_experts=num_experts, moe_layers=tuple(sorted(layers))) + + +def layer_counts(hybrid_override_pattern: str) -> dict[str, int]: + """Layer kinds of a Nemotron-H ``hybrid_override_pattern`` (M mamba, E MoE, + * attention, - dense MLP).""" + kinds = {"M": "mamba", "E": "moe", "*": "attention", "-": "mlp"} + unknown = set(hybrid_override_pattern) - set(kinds) + if unknown: + raise ValueError(f"unknown layer symbols {sorted(unknown)} in hybrid_override_pattern") + return {v: hybrid_override_pattern.count(k) for k, v in kinds.items()} From b1dfaa4760740882dc80fbc0f9e8d134fcbd4827 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:20:46 +0000 Subject: [PATCH 165/233] Nemotron-H model FLOPs for MFU and a Modal verl GRPO smoke tool verl's flops counter has no nemotron_h entry; hybrid_flops counts Mamba-2, MoE (routed top-k, shared expert, latent projections) and attention layers. The smoke tool builds the verl image with the megatron extra and the Mamba kernels, stages weights and GSM8K on a volume, and runs the bare verl trainer on pinned H200 cards. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_hybrid_params.py | 23 +++++ tools/nemotron_verl_smoke.py | 164 +++++++++++++++++++++++++++++++++ yeto/rl/hybrid_flops.py | 84 +++++++++++++++++ 3 files changed, 271 insertions(+) create mode 100644 tools/nemotron_verl_smoke.py create mode 100644 yeto/rl/hybrid_flops.py diff --git a/tests/test_rl_hybrid_params.py b/tests/test_rl_hybrid_params.py index a5805e0b3..75849454b 100644 --- a/tests/test_rl_hybrid_params.py +++ b/tests/test_rl_hybrid_params.py @@ -150,3 +150,26 @@ def test_image_commands_pin_mamba_kernels(): cmds = mr.image_commands() assert any("--extra megatron" in c for c in cmds) assert any("causal-conv1d==" in c for c in cmds) and any("mamba-ssm==" in c for c in cmds) + + +# Nemotron-3 Nano 30B-A3B config.json (HF, fields the FLOPs formula reads). +NANO = {"hybrid_override_pattern": "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME", "hidden_size": 2688, + "vocab_size": 131072, "head_dim": 128, "num_attention_heads": 32, "num_key_value_heads": 2, + "mamba_num_heads": 64, "mamba_head_dim": 64, "n_groups": 8, "ssm_state_size": 128, "conv_kernel": 4, + "n_routed_experts": 128, "num_experts_per_tok": 6, "moe_intermediate_size": 1856, + "moe_shared_expert_intermediate_size": 3712, "n_shared_experts": 1, "moe_latent_size": None, + "intermediate_size": 1856} + + +def test_hybrid_flops_matches_active_params(): + from yeto.rl import hybrid_flops as hf + per_token = hf.forward_flops(NANO, [1]) + # 2 x active parameters (model card: ~3.5B active incl. embeddings, ~0.35B of them) + assert 3.0e9 < per_token / 2 < 3.5e9 + assert hf.train_flops(NANO, [10, 20]) == 3 * hf.forward_flops(NANO, [10, 20]) + # attention adds an s^2 term only for the 6 attention layers + extra = hf.forward_flops(NANO, [4096]) - 4096 * per_token + assert extra == pytest.approx(6 * 2 * 32 * 128 * (4096 ** 2 - 4096)) + assert hf.mfu(NANO, [1000], 1.0, 1) == pytest.approx(hf.train_flops(NANO, [1000]) / 989.4e12) + latent = {**NANO, "moe_latent_size": 1024} + assert hf.forward_flops(latent, [1]) < per_token # experts run at the latent width diff --git a/tools/nemotron_verl_smoke.py b/tools/nemotron_verl_smoke.py new file mode 100644 index 000000000..cc1fabb00 --- /dev/null +++ b/tools/nemotron_verl_smoke.py @@ -0,0 +1,164 @@ +"""GPU smoke: verl Megatron full-parameter GRPO on Nemotron-3 Nano 30B-A3B (H200), MFU. + +Usage (local): + python tools/nemotron_verl_smoke.py prep RUN_ID # CPU: image build, weights, data, imports + python tools/nemotron_verl_smoke.py smoke RUN_ID [N_GPUS] # GPU: a few GRPO steps, spawned detached + python tools/nemotron_verl_smoke.py pull RUN_ID # fetch small result files, stop the app + +The run uses yeto's verl image plus :func:`megatron_recipe.image_commands` and +the overrides from :func:`megatron_recipe.build_overrides`; the trainer is the +bare verl ``main_ppo`` (not yeto's island loop or syncer). Weights stay on the +Modal volume; ``pull`` never fetches files over 50 MB. +""" + +from __future__ import annotations + +import json +import re +import subprocess +import sys +import time +from pathlib import Path + +REPO = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO)) + +APP = "yeto-nemotron-smoke-" +VOLUME = "yeto-nemotron" +MNT = "/mnt/nem" +MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" +VERL_PY = "/workspace/verl/.venv/bin/python" +MAX_PULL = 50 << 20 + + +def _sh(cmd: str, log: Path, env: dict | None = None, timeout: int = 7200) -> int: + import os + with log.open("a") as fh: + fh.write(f"\n$ {cmd}\n") + fh.flush() + p = subprocess.run(["bash", "-lc", cmd], stdout=fh, stderr=subprocess.STDOUT, + env={**os.environ, **(env or {})}, timeout=timeout) + return p.returncode + + +def prep(run: str) -> dict: + import modal + out = Path(MNT) / "runs" / run + out.mkdir(parents=True, exist_ok=True) + log = out / "prep.log" + model_dir = Path(MNT) / "models" / MODEL_ID.split("/")[1] + rc = {} + if not (model_dir / "config.json").exists(): + from huggingface_hub import snapshot_download + snapshot_download(MODEL_ID, local_dir=str(model_dir), max_workers=16) + rc["model_files"] = len(list(model_dir.glob("*.safetensors"))) + data = Path(MNT) / "data" / "gsm8k" + if not (data / "train.parquet").exists(): + rc["gsm8k"] = _sh(f"cd /workspace/verl && {VERL_PY} examples/data_preprocess/gsm8k.py " + f"--local_save_dir {data}", log) + rc["imports"] = _sh( + f"{VERL_PY} -c 'import mamba_ssm, causal_conv1d, megatron.core, megatron.bridge, vllm, torch;" + "from megatron.bridge import AutoBridge;" + f"print(torch.__version__, vllm.__version__, mamba_ssm.__version__, causal_conv1d.__version__);" + f"print(AutoBridge.can_handle(\"{model_dir}\", trust_remote_code=True))'", log) + modal.Volume.from_name(VOLUME).commit() + return rc + + +def _gpu_names() -> list[str]: + out = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output=True, text=True, timeout=60).stdout + return [n.strip() for n in out.splitlines() if n.strip()] + + +def smoke(run: str, plan: dict) -> dict: + import modal + from yeto.rl.adapters.verl.megatron_recipe import NemotronMegatronPlan, build_overrides + out = Path(MNT) / "runs" / run + out.mkdir(parents=True, exist_ok=True) + meta = {"began": time.time(), "plan": plan} + names = _gpu_names() + meta["gpus"] = names + if len(names) != plan["n_gpus"] or not all(re.search(r"H200(?![\w-])", n) for n in names): + meta["error"] = f"asked for {plan['n_gpus']}x H200, got {names}" + (out / "meta.json").write_text(json.dumps(meta, indent=1)) + modal.Volume.from_name(VOLUME).commit() + raise RuntimeError(meta["error"]) + nvml = subprocess.Popen(["bash", "-c", "nvidia-smi --query-gpu=timestamp,index,utilization.gpu,memory.used," + f"power.draw --format=csv -l 5 >> {out}/nvml.csv"]) + overrides = build_overrides(NemotronMegatronPlan(**plan)) + (out / "overrides.txt").write_text("\n".join(overrides) + "\n") + cmd = "cd /workspace/verl && " + VERL_PY + " -m verl.trainer.main_ppo " + " ".join(f"'{o}'" for o in overrides) + env = {"VERL_FILE_LOGGER_PATH": f"{out}/metrics.jsonl", "VLLM_USE_V1": "1", "RAY_TMPDIR": "/tmp/ray", + "HYDRA_FULL_ERROR": "1"} + try: + meta["rc"] = _sh(cmd, out / "train.log", env, timeout=int(plan.get("_timeout", 5400))) + finally: + nvml.terminate() + meta["seconds"] = round(time.time() - meta["began"], 1) + (out / "meta.json").write_text(json.dumps(meta, indent=1)) + modal.Volume.from_name(VOLUME).commit() + return meta + + +def build_app(run: str, n_gpus: int): + import modal + from modal._vendor import cloudpickle + from yeto.rl.adapters.verl.image import modal_image + from yeto.rl.adapters.verl.megatron_recipe import image_commands + cloudpickle.register_pickle_by_value(sys.modules[__name__]) + ignore = [".git", ".venv", ".claude", "**/__pycache__", "syncer/target", "evidence*", "openspec", "docs", + "tutorial", "data", "s1-runs"] + image = (modal_image(modal, patched=False).run_commands(*image_commands()) + .pip_install("huggingface_hub[hf_transfer]==0.36.0") + .env({"PYTHONUNBUFFERED": "1", "HF_HUB_ENABLE_HF_TRANSFER": "1", "PYTHONPATH": "/root/yeto"}) + .add_local_dir(str(REPO / "yeto"), "/root/yeto/yeto", copy=False)) + vol = modal.Volume.from_name(VOLUME, create_if_missing=True, version=2) + app = modal.App(APP + run) + fns = { + "prep": app.function(image=image, cpu=16.0, memory=65536, timeout=7200, volumes={MNT: vol}, + ephemeral_disk=600 * 1024, name="prep", serialized=True)(prep), + "smoke": app.function(image=image, gpu=f"H200!:{n_gpus}", cpu=32.0, memory=786432, timeout=7200, + retries=0, volumes={MNT: vol}, ephemeral_disk=1024 * 1024, name="smoke", + serialized=True)(smoke), + } + return app, fns + + +def plan_for(n_gpus: int) -> dict: + root = f"{MNT}" + return dict(model_path=f"{root}/models/{MODEL_ID.split('/')[1]}", train_file=f"{root}/data/gsm8k/train.parquet", + val_file=f"{root}/data/gsm8k/test.parquet", out_dir="/tmp/out", n_gpus=n_gpus, tp=n_gpus, + ep=n_gpus, rollout_tp=n_gpus, groups_per_round=32, samples_per_group=8, max_prompt_len=512, + max_response_len=1024, total_steps=4, num_experts=128) + + +def main() -> None: + import modal + action, run = sys.argv[1], sys.argv[2] + n_gpus = int(sys.argv[3]) if len(sys.argv) > 3 else 2 + local = REPO / "s1-runs" / "nemotron" / run + local.mkdir(parents=True, exist_ok=True) + if action == "pull": + listing = subprocess.run([sys.executable, "-m", "modal", "volume", "ls", "--json", VOLUME, f"runs/{run}"], + capture_output=True, text=True, cwd="/tmp") + for entry in json.loads(listing.stdout or "[]"): + size = int(str(entry.get("Size", "0")).split()[0]) if str(entry.get("Size", "0")).split()[0].isdigit() else 0 + name = entry.get("Filename") or entry.get("path") + if entry.get("Type") == "file" and size < MAX_PULL: + subprocess.run([sys.executable, "-m", "modal", "volume", "get", "--force", VOLUME, name, + str(local)], cwd="/tmp") + app_id = (local / "app_id.txt").read_text().strip() if (local / "app_id.txt").exists() else None + if app_id: + subprocess.run([sys.executable, "-m", "modal", "app", "stop", "-y", app_id], cwd="/tmp") + return + app, fns = build_app(run, n_gpus) + with modal.enable_output(), app.run(detach=True): + (local / "app_id.txt").write_text(app.app_id + "\n") + call = fns[action].spawn(run) if action == "prep" else fns["smoke"].spawn(run, plan_for(n_gpus)) + (local / "call_id.txt").write_text(call.object_id + "\n") + print("spawned", action, call.object_id, app.app_id, flush=True) + + +if __name__ == "__main__": + main() diff --git a/yeto/rl/hybrid_flops.py b/yeto/rl/hybrid_flops.py new file mode 100644 index 000000000..bae55e2b0 --- /dev/null +++ b/yeto/rl/hybrid_flops.py @@ -0,0 +1,84 @@ +"""Model FLOPs of a Nemotron-H hybrid (Mamba-2 + MoE + attention) for MFU. + +verl's ``flops_counter`` has no ``nemotron_h`` entry (fork acad9875), so its +``perf/mfu`` is not usable for this model; this module is the formula yeto +reports. Forward FLOPs per sequence of length ``s`` (multiply-add = 2 FLOPs): + +* attention layer: ``2*h*(q + 2*kv)`` + ``2*q*h`` per token for the + projections (``q = heads*head_dim``, ``kv = kv_heads*head_dim``), plus the + causal score and value products ``2 * 2*q * s^2/2`` per sequence; +* Mamba-2 layer: ``in_proj`` ``2*h*(2*d_in + 2*g*N + H)``, ``out_proj`` + ``2*d_in*h``, causal conv ``2*k*(d_in + 2*g*N)``, and the SSD scan + ``6*d_in*N`` per token (state update, state readout and the intra-chunk + term counted as one more ``2*d_in*N``; an approximation, small next to the + projections); +* MoE layer: router ``2*h*E``; latent projections ``2*2*h*L`` when + ``moe_latent_size`` is set (experts then run at width ``L``); each of the + ``top_k`` routed experts ``2*2*w*I`` (relu2, not gated: up and down only); + the shared expert ``2*2*h*I_s`` on the full hidden size; +* LM head ``2*h*V`` per token. + +Training FLOPs = 3 x forward (forward + backward). Activation recompute is +not counted: this is model FLOPs (MFU), not hardware FLOPs. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from typing import Any + +# Dense BF16 peak, no sparsity (NVIDIA H100/H200 SXM datasheets: 989.4 TFLOPS). +PEAK_BF16_TFLOPS = {"H100": 989.4, "H200": 989.4} + + +def _get(cfg: Mapping[str, Any], key: str, default: Any = None) -> Any: + value = cfg.get(key, default) + return default if value is None else value + + +def forward_flops(cfg: Mapping[str, Any], seqlens: Iterable[int]) -> float: + """Forward FLOPs of a batch of sequences for an HF ``nemotron_h`` config dict.""" + seqlens = [int(s) for s in seqlens] + tokens = sum(seqlens) + sq = sum(s * s for s in seqlens) + pattern = cfg["hybrid_override_pattern"] + h = int(cfg["hidden_size"]) + v = int(cfg["vocab_size"]) + head_dim = int(_get(cfg, "head_dim", h // int(cfg["num_attention_heads"]))) + q = int(cfg["num_attention_heads"]) * head_dim + kv = int(cfg["num_key_value_heads"]) * head_dim + m_heads = int(cfg["mamba_num_heads"]) + d_in = m_heads * int(cfg["mamba_head_dim"]) + g, n = int(cfg["n_groups"]), int(cfg["ssm_state_size"]) + k = int(_get(cfg, "conv_kernel", 4)) + e = int(_get(cfg, "n_routed_experts", 0)) + top_k = int(_get(cfg, "num_experts_per_tok", 0)) + inter = int(_get(cfg, "moe_intermediate_size", 0)) + shared = int(_get(cfg, "moe_shared_expert_intermediate_size", 0)) * int(_get(cfg, "n_shared_experts", 1)) + latent = _get(cfg, "moe_latent_size") + width = int(latent) if latent else h + dense_inter = int(_get(cfg, "intermediate_size", 0)) + + per_token = { + "*": 2 * h * (q + 2 * kv) + 2 * q * h, + "M": 2 * h * (2 * d_in + 2 * g * n + m_heads) + 2 * d_in * h + 2 * k * (d_in + 2 * g * n) + 6 * d_in * n, + "E": 2 * h * e + (4 * h * width if latent else 0) + top_k * 4 * width * inter + 4 * h * shared, + "-": 4 * h * dense_inter, + } + total = 0.0 + for symbol in pattern: + if symbol not in per_token: + raise ValueError(f"unknown layer symbol {symbol!r}") + total += per_token[symbol] * tokens + if symbol == "*": + total += 2 * q * sq # 2 products x 2 FLOPs x q x s^2/2 (causal) + return total + 2 * h * v * tokens + + +def train_flops(cfg: Mapping[str, Any], seqlens: Iterable[int]) -> float: + return 3.0 * forward_flops(cfg, seqlens) + + +def mfu(cfg: Mapping[str, Any], seqlens: Iterable[int], seconds: float, n_gpus: int, gpu: str = "H200") -> float: + """Model FLOPs utilization of a training phase that processed ``seqlens`` in ``seconds``.""" + return train_flops(cfg, seqlens) / (seconds * n_gpus * PEAK_BF16_TFLOPS[gpu] * 1e12) From 76c18185806e7d3a4d7523a51f7e99e43584539a Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 04:06:30 +0000 Subject: [PATCH 166/233] verl Nemotron image: build mamba-ssm CUDA, drop flash_attn.cute, no remote code Co-Authored-By: Claude Opus 5.5 --- tools/nemotron_verl_smoke.py | 11 +++++------ yeto/rl/adapters/verl/megatron_recipe.py | 14 +++++++++----- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/tools/nemotron_verl_smoke.py b/tools/nemotron_verl_smoke.py index cc1fabb00..1817d7805 100644 --- a/tools/nemotron_verl_smoke.py +++ b/tools/nemotron_verl_smoke.py @@ -49,8 +49,8 @@ def prep(run: str) -> dict: model_dir = Path(MNT) / "models" / MODEL_ID.split("/")[1] rc = {} if not (model_dir / "config.json").exists(): - from huggingface_hub import snapshot_download - snapshot_download(MODEL_ID, local_dir=str(model_dir), max_workers=16) + rc["download"] = _sh(f"{VERL_PY} -c 'from huggingface_hub import snapshot_download as d;" + f"d(\"{MODEL_ID}\", local_dir=\"{model_dir}\", max_workers=16)'", log) rc["model_files"] = len(list(model_dir.glob("*.safetensors"))) data = Path(MNT) / "data" / "gsm8k" if not (data / "train.parquet").exists(): @@ -60,7 +60,7 @@ def prep(run: str) -> dict: f"{VERL_PY} -c 'import mamba_ssm, causal_conv1d, megatron.core, megatron.bridge, vllm, torch;" "from megatron.bridge import AutoBridge;" f"print(torch.__version__, vllm.__version__, mamba_ssm.__version__, causal_conv1d.__version__);" - f"print(AutoBridge.can_handle(\"{model_dir}\", trust_remote_code=True))'", log) + f"print(AutoBridge.can_handle(\"{model_dir}\", trust_remote_code=False))'", log) modal.Volume.from_name(VOLUME).commit() return rc @@ -110,15 +110,14 @@ def build_app(run: str, n_gpus: int): ignore = [".git", ".venv", ".claude", "**/__pycache__", "syncer/target", "evidence*", "openspec", "docs", "tutorial", "data", "s1-runs"] image = (modal_image(modal, patched=False).run_commands(*image_commands()) - .pip_install("huggingface_hub[hf_transfer]==0.36.0") - .env({"PYTHONUNBUFFERED": "1", "HF_HUB_ENABLE_HF_TRANSFER": "1", "PYTHONPATH": "/root/yeto"}) + .env({"PYTHONUNBUFFERED": "1", "PYTHONPATH": "/root/yeto"}) .add_local_dir(str(REPO / "yeto"), "/root/yeto/yeto", copy=False)) vol = modal.Volume.from_name(VOLUME, create_if_missing=True, version=2) app = modal.App(APP + run) fns = { "prep": app.function(image=image, cpu=16.0, memory=65536, timeout=7200, volumes={MNT: vol}, ephemeral_disk=600 * 1024, name="prep", serialized=True)(prep), - "smoke": app.function(image=image, gpu=f"H200!:{n_gpus}", cpu=32.0, memory=786432, timeout=7200, + "smoke": app.function(image=image, gpu=f"H200:{n_gpus}", cpu=32.0, memory=786432, timeout=7200, retries=0, volumes={MNT: vol}, ephemeral_disk=1024 * 1024, name="smoke", serialized=True)(smoke), } diff --git a/yeto/rl/adapters/verl/megatron_recipe.py b/yeto/rl/adapters/verl/megatron_recipe.py index 27c485ff3..8d4298e6b 100644 --- a/yeto/rl/adapters/verl/megatron_recipe.py +++ b/yeto/rl/adapters/verl/megatron_recipe.py @@ -102,7 +102,7 @@ def build_overrides(p: NemotronMegatronPlan) -> list[str]: f"data.max_response_length={p.max_response_len}", f"data.train_batch_size={p.groups_per_round}", f"actor_rollout_ref.model.path={p.model_path}", - "actor_rollout_ref.model.trust_remote_code=True", + "actor_rollout_ref.model.trust_remote_code=False", "actor_rollout_ref.model.use_fused_kernels=False", "actor_rollout_ref.actor.use_kl_loss=False", "actor_rollout_ref.actor.clip_ratio_low=0.2", @@ -176,17 +176,21 @@ def build_overrides(p: NemotronMegatronPlan) -> list[str]: def image_commands(max_jobs: int = 16) -> list[str]: """Extra image layers on top of the verl uv venv: the megatron extra plus the Mamba training kernels the extra does not list (the example installs them by - hand: causal-conv1d and mamba-ssm). Hopper only; mamba-ssm's own CUDA - selective-scan is skipped (Mamba-2 training uses its Triton kernels).""" + hand: causal-conv1d and mamba-ssm). Hopper only; mamba-ssm builds its CUDA + loads its CUDA selective-scan module, so that build cannot be skipped.""" from .image import VENV_BIN from .pins import VERL_PYTHON, VERL_ROOT - env = f"MAX_JOBS={max_jobs} TORCH_CUDA_ARCH_LIST=9.0" + env = f"MAX_JOBS={max_jobs} TORCH_CUDA_ARCH_LIST=9.0 CC=gcc CXX=g++" return [ f"cd {VERL_ROOT} && UV_PYTHON=/usr/local/bin/python{VERL_PYTHON} uv sync --frozen --all-packages" " --extra vllm --extra megatron", f"uv pip install --python {VENV_BIN}/python ninja packaging wheel setuptools", f"{env} CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --python {VENV_BIN}/python --no-build-isolation" " 'causal-conv1d==1.6.2.post1'", - f"{env} MAMBA_SKIP_CUDA_BUILD=TRUE MAMBA_FORCE_BUILD=TRUE uv pip install --python {VENV_BIN}/python" + f"{env} MAMBA_FORCE_BUILD=TRUE uv pip install --python {VENV_BIN}/python" " --no-build-isolation --no-deps 'mamba-ssm==2.3.2.post1'", + # flash-attn 2.8.3's flash_attn.cute (FA4 preview) raises AttributeError (not + # ImportError) against the locked nvidia-cutlass-dsl 4.6.2, and Megatron-Core's + # attention.py only catches ImportError. FA4 targets Blackwell; Hopper does not use it. + f"rm -rf {VERL_ROOT}/.venv/lib/python{VERL_PYTHON}/site-packages/flash_attn/cute", ] From ac2a7dc6b1a7be744639e4532d964b1f4523caec Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 04:19:03 +0000 Subject: [PATCH 167/233] Miles recipe for the Nemotron-H family (hybrid Mamba-2 + MoE) The Megatron-Bridge hybrid provider selects the nemotron_h recipe; its router flags (sigmoid score, expert bias, top-k scaling, group routing, latent MoE) come from the provider, and RL fixes the routing statistics (bias update rate 0, aux loss 0). MTP checkpoints are refused for full-parameter RL until the island drops the MTP layer. Co-Authored-By: Claude Opus 5.5 --- s1-runs/nemotron/gr-prep4/app_id.txt | 1 + s1-runs/nemotron/gr-prep4/call_id.txt | 1 + s1-runs/nemotron/gr-prep5/app_id.txt | 1 + s1-runs/nemotron/gr-prep5/call_id.txt | 1 + s1-runs/nemotron/gr-prep6/app_id.txt | 1 + s1-runs/nemotron/gr-prep6/call_id.txt | 1 + tests/test_rl_nemotron_h_recipe.py | 87 +++++++++++++++++++++ yeto/rl/adapters/miles/config.py | 3 + yeto/rl/adapters/miles/models/nemotron_h.py | 87 +++++++++++++++++++++ yeto/rl/engine/run_config.py | 30 +++++++ 10 files changed, 213 insertions(+) create mode 100644 s1-runs/nemotron/gr-prep4/app_id.txt create mode 100644 s1-runs/nemotron/gr-prep4/call_id.txt create mode 100644 s1-runs/nemotron/gr-prep5/app_id.txt create mode 100644 s1-runs/nemotron/gr-prep5/call_id.txt create mode 100644 s1-runs/nemotron/gr-prep6/app_id.txt create mode 100644 s1-runs/nemotron/gr-prep6/call_id.txt create mode 100644 tests/test_rl_nemotron_h_recipe.py create mode 100644 yeto/rl/adapters/miles/models/nemotron_h.py diff --git a/s1-runs/nemotron/gr-prep4/app_id.txt b/s1-runs/nemotron/gr-prep4/app_id.txt new file mode 100644 index 000000000..4df455ead --- /dev/null +++ b/s1-runs/nemotron/gr-prep4/app_id.txt @@ -0,0 +1 @@ +ap-x5XCH9wc5uJnjEnPIDengA diff --git a/s1-runs/nemotron/gr-prep4/call_id.txt b/s1-runs/nemotron/gr-prep4/call_id.txt new file mode 100644 index 000000000..c5b5bb489 --- /dev/null +++ b/s1-runs/nemotron/gr-prep4/call_id.txt @@ -0,0 +1 @@ +fc-01M5KD7XH5EYKSAN0T4BSHJ1QM diff --git a/s1-runs/nemotron/gr-prep5/app_id.txt b/s1-runs/nemotron/gr-prep5/app_id.txt new file mode 100644 index 000000000..f6c705e91 --- /dev/null +++ b/s1-runs/nemotron/gr-prep5/app_id.txt @@ -0,0 +1 @@ +ap-vKl9KOvqfNesmNZEEZ9za1 diff --git a/s1-runs/nemotron/gr-prep5/call_id.txt b/s1-runs/nemotron/gr-prep5/call_id.txt new file mode 100644 index 000000000..e95bd6189 --- /dev/null +++ b/s1-runs/nemotron/gr-prep5/call_id.txt @@ -0,0 +1 @@ +fc-01M5KE4126TYRGPWY1GE39N5V6 diff --git a/s1-runs/nemotron/gr-prep6/app_id.txt b/s1-runs/nemotron/gr-prep6/app_id.txt new file mode 100644 index 000000000..1a22ab569 --- /dev/null +++ b/s1-runs/nemotron/gr-prep6/app_id.txt @@ -0,0 +1 @@ +ap-kaPxkOp5ljCobSWfDm00Ay diff --git a/s1-runs/nemotron/gr-prep6/call_id.txt b/s1-runs/nemotron/gr-prep6/call_id.txt new file mode 100644 index 000000000..69e4ec668 --- /dev/null +++ b/s1-runs/nemotron/gr-prep6/call_id.txt @@ -0,0 +1 @@ +fc-01M5KEYXTY8RJCTR8WGH9B46R8 diff --git a/tests/test_rl_nemotron_h_recipe.py b/tests/test_rl_nemotron_h_recipe.py new file mode 100644 index 000000000..8e0a97fe1 --- /dev/null +++ b/tests/test_rl_nemotron_h_recipe.py @@ -0,0 +1,87 @@ +"""Nemotron-H (Mamba-2 + MoE hybrid) Miles recipe: provider-driven model flags +reach the Miles argv; MTP is refused in full-parameter mode. CPU only; the +provider is a stand-in carrying the fields Megatron-Bridge 0.6.2 +``nemotron_h_bridge.py`` sets.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from test_s25_g1_stage1mini import REPO + +from yeto.rl.adapters.miles.models import nemotron_h + + +def _provider(**kw): + # Nemotron-3 Nano 30B-A3B (HF config.json) as the Bridge provider sees it + base = dict( + is_hybrid_model=True, hybrid_layer_pattern="MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME", + num_layers=52, hidden_size=2688, num_attention_heads=32, num_query_groups=2, kv_channels=128, + ffn_hidden_size=1856, seq_length=262144, vocab_size=131072, layernorm_epsilon=1e-5, + normalization="RMSNorm", position_embedding_type="none", share_embeddings_and_output_weights=False, + add_bias_linear=False, gated_linear_unit=False, + num_moe_experts=128, moe_ffn_hidden_size=1856, moe_router_topk=6, moe_layer_freq=1, + moe_shared_expert_intermediate_size=3712, moe_router_score_function="sigmoid", + moe_router_enable_expert_bias=True, moe_router_load_balancing_type="seq_aux_loss", + moe_router_dtype="fp32", moe_grouped_gemm=True, moe_router_topk_scaling_factor=2.5, + moe_router_num_groups=1, moe_router_group_topk=1, moe_latent_size=None, mtp_num_layers=0, + mamba_num_heads=64, mamba_head_dim=64, mamba_state_dim=128, mamba_num_groups=8) + return SimpleNamespace(**{**base, **kw}) + + +def _flag(argv, name): + return argv[argv.index(name) + 1] + + +def test_model_flags_from_provider(): + flags = list(nemotron_h.model_flags(_provider())) + assert _flag(flags, "--moe-router-score-function") == "sigmoid" + assert _flag(flags, "--moe-router-topk-scaling-factor") == "2.5" + assert _flag(flags, "--moe-router-bias-update-rate") == "0" + assert _flag(flags, "--moe-aux-loss-coeff") == "0" + assert "--moe-router-enable-expert-bias" in flags and "--moe-grouped-gemm" in flags + assert "--moe-latent-size" not in flags + # 120B Super: latent MoE and 512 experts come from the provider, not a table + super_flags = list(nemotron_h.model_flags(_provider(moe_latent_size=1024, num_moe_experts=512))) + assert _flag(super_flags, "--moe-latent-size") == "1024" + + +def test_check_provider_refuses_mtp_in_full_parameter_and_bad_pattern(): + nemotron_h.check_provider(_provider(), parameter_mode="full") + with pytest.raises(nemotron_h.NemotronHRecipeError, match="MTP"): + nemotron_h.check_provider(_provider(mtp_num_layers=1), parameter_mode="full") + nemotron_h.check_provider(_provider(mtp_num_layers=1), parameter_mode="lora") + with pytest.raises(nemotron_h.NemotronHRecipeError, match="experts"): + nemotron_h.check_provider(_provider(num_moe_experts=None), parameter_mode="full") + + +def test_nemotron_h_provider_reaches_the_miles_argv(tmp_path): + from yeto.rl.adapters.miles.island_entry import build_ports_launch, parse_args + from yeto.rl.engine.run_config import RECIPE_NEMOTRON_H, resolve_rl_run_config + from yeto.rl.stage1 import expand, load_plan + + model = tmp_path / "model" + model.mkdir() + (model / "config.json").write_text(json.dumps({"model_type": "nemotron_h"})) + prompts = tmp_path / "p.jsonl" + prompts.write_text('{"messages": [{"role": "user", "content": "hi"}], "label": "1"}\n') + argv = expand(load_plan(REPO / "examples/stage1/glm5next-mini-modal.yaml")).trainer["miles"]["island_argv"][3:] + argv = [str(model) if a == "/data/models/glm5next-mini-bf16-7fb7883" else a for a in argv] + argv = [str(REPO / a) if a.startswith("examples/rl_algorithms/") else a for a in argv] + if "--rl-model-recipe" in argv: + argv[argv.index("--rl-model-recipe") + 1] = "generic" + args = parse_args(argv) + config = resolve_rl_run_config(args, model_path=model, prompt_path=prompts, provider=_provider(), + target_modules=[], yeto_policy_sync=False, verify_ref_load=False) + assert config.model_recipe.name == RECIPE_NEMOTRON_H + assert config.model_recipe.training_attention_backend == "auto" + miles = list(build_ports_launch(args, config, []).argv) + assert _flag(miles, "--megatron-to-hf-mode") == "bridge" + assert _flag(miles, "--position-embedding-type") == "none" + assert _flag(miles, "--num-experts") == "128" + assert _flag(miles, "--moe-router-bias-update-rate") == "0" + assert "--swiglu" not in miles # relu2 experts + assert miles.count("--moe-router-score-function") == 1 diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index eb787d272..cbc65f521 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -334,6 +334,7 @@ def _check_sglang_tp(value, config): "model_recipe.gdn.gated_delta_net": _ok, "model_recipe.gdn.qkv_format": _ok, "model_recipe.transformer_impl": _ok, + "model_recipe.model_flags": _ok, **{f"geometry.{name}": _ok for name in _GEOMETRY_FIELDS}, **{ f"geometry.moe.{name}": _ok @@ -1141,6 +1142,8 @@ def translate_run_config( values.extend( ("--moe-shared-expert-intermediate-size", str(moe.shared_expert_intermediate_size)) ) + # recipe-owned model flags read from the provider (nemotron_h: router, latent MoE) + values.extend(getattr(recipe, "model_flags", ())) if geometry.multi_latent_attention: values.append("--multi-latent-attention") for name, value in geometry.mla_dims: diff --git a/yeto/rl/adapters/miles/models/nemotron_h.py b/yeto/rl/adapters/miles/models/nemotron_h.py new file mode 100644 index 000000000..61e9490da --- /dev/null +++ b/yeto/rl/adapters/miles/models/nemotron_h.py @@ -0,0 +1,87 @@ +"""Miles recipe for the NVIDIA Nemotron-H family (Mamba-2 + MoE + attention). + +Miles trains Nemotron-H through Megatron-Bridge (``--megatron-to-hf-mode +bridge``): the model is built from the Bridge ``HybridModelProvider`` of the HF +checkpoint, and only the training knobs that +``miles/backends/megatron_utils/model_provider.py:_apply_bridge_runtime_config`` +copies from args reach it. The Megatron argv still has to describe the same +model, because Megatron's argument checks and Miles' rollout/export code read +it. This module writes those flags from the provider, so one recipe covers +the family (pattern-driven, not one checkpoint): + +* router (Bridge ``nemotron_h_bridge.py`` provider_bridge): sigmoid score, + expert bias, ``seq_aux_loss``, fp32 router, grouped GEMM, the top-k scaling + factor and group routing from the checkpoint; +* RL keeps the routing statistics fixed: ``--moe-router-bias-update-rate 0`` + and ``--moe-aux-loss-coeff 0`` (Miles ``scripts/models/nemotron-3-nano-30b-a3b.py``). + Both are copied onto the provider by Miles, so the expert bias does not + drift on an island and every island trains the same router; +* latent MoE (``moe_latent_size``, Nemotron-3 Super / 3.5 Super VL) is passed + through. + +Differences that matter for the 120B Super (VL) and are NOT handled here: + +* MTP: the Super checkpoints carry one MTP layer; the Bridge provider builds it + (``mtp_num_layers``) and Miles does not copy an MTP override in bridge mode. + :func:`check_provider` refuses MTP in full-parameter RL until the island + drops it (NeMo-RL's Super-VL recipe sets ``num_nextn_predict_layers=0``). +* 512 experts with top-22: covered by the same flags; EP must divide 512. +* VL: the Bridge provider of ``NemotronH_Omni_Reasoning_V3`` is not this one; + the vision freeze rule is :mod:`yeto.rl.hybrid_params`. +""" + +from __future__ import annotations + +from typing import Any + + +class NemotronHRecipeError(ValueError): + pass + + +def _value(provider: Any, name: str, default: Any = None) -> Any: + value = getattr(provider, name, default) + return default if value is None else value + + +def check_provider(provider: Any, *, parameter_mode: str) -> None: + pattern = str(_value(provider, "hybrid_layer_pattern", "")) + if not pattern: + raise NemotronHRecipeError("Nemotron-H provider has no hybrid_layer_pattern") + if "E" in pattern and _value(provider, "num_moe_experts") is None: + raise NemotronHRecipeError("hybrid pattern has MoE layers (E) but the provider has no experts") + if int(_value(provider, "mtp_num_layers", 0) or 0) and parameter_mode == "full": + raise NemotronHRecipeError( + "Nemotron-H checkpoint has MTP layers; full-parameter RL would train and sync them, " + "and Miles bridge mode cannot drop them yet (120B Super gap)") + + +def model_flags(provider: Any) -> tuple[str, ...]: + """Miles/Megatron model flags of a Nemotron-H provider (module docstring).""" + flags: list[str] = ["--make-vocab-size-divisible-by", "128"] + if _value(provider, "num_moe_experts") is None: + return tuple(flags) + flags += [ + "--moe-router-score-function", str(_value(provider, "moe_router_score_function", "sigmoid")), + "--moe-router-load-balancing-type", str(_value(provider, "moe_router_load_balancing_type", "seq_aux_loss")), + "--moe-router-dtype", str(_value(provider, "moe_router_dtype", "fp32")), + "--moe-router-bias-update-rate", "0", + "--moe-aux-loss-coeff", "0", + ] + if _value(provider, "moe_router_enable_expert_bias", False): + flags.append("--moe-router-enable-expert-bias") + if _value(provider, "moe_grouped_gemm", False): + flags.append("--moe-grouped-gemm") + scaling = _value(provider, "moe_router_topk_scaling_factor") + if scaling is not None: + flags += ["--moe-router-topk-scaling-factor", repr(float(scaling))] + groups = _value(provider, "moe_router_num_groups") + if groups is not None: + flags += ["--moe-router-num-groups", str(int(groups)), + "--moe-router-group-topk", str(int(_value(provider, "moe_router_group_topk", 1)))] + if _value(provider, "moe_router_pre_softmax", False): + flags.append("--moe-router-pre-softmax") + latent = _value(provider, "moe_latent_size") + if latent is not None: + flags += ["--moe-latent-size", str(int(latent))] + return tuple(flags) diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index c0a3d4b09..7858496ed 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -40,6 +40,12 @@ # GLM5.3-Flash (glm5_next): KDA + DSA MLA + 288-expert MoE through the Miles # plugin spec (adapters/miles/models/glm5_next.py), full-parameter. RECIPE_GLM5_NEXT = "glm5_next" +# NVIDIA Nemotron-H family (Nemotron-3 Nano / Super, Nemotron-3.5 Super VL text +# model): Mamba-2 + MoE + attention hybrid built by the Megatron-Bridge +# HybridModelProvider from the HF checkpoint (Miles --megatron-to-hf-mode +# bridge). The recipe is driven by the provider's layer pattern and router +# fields, not by one checkpoint (adapters/miles/models/nemotron_h.py). +RECIPE_NEMOTRON_H = "nemotron_h" TRANSFORMER_IMPLS = ("transformer_engine", "local") GATED_DELTA_NET = "gated_delta_net" @@ -143,6 +149,9 @@ class ModelRecipe: # Megatron --transformer-impl of a recipe that chooses it (glm5_next: the # N7.5 TE-free fallback is "local"); the other recipes keep their own. transformer_impl: str = "transformer_engine" + # Recipe-owned Miles model flags read from the provider (nemotron_h router, + # latent MoE); empty for the other recipes. + model_flags: tuple[str, ...] = () @dataclass(frozen=True) @@ -535,6 +544,20 @@ def qwen3_8_next_variant(args, provider) -> str | None: return None +def _is_nemotron_h(provider) -> bool: + """A Megatron-Bridge hybrid (Mamba) provider built from a Nemotron-H checkpoint.""" + return bool(getattr(provider, "is_hybrid_model", False)) and bool( + getattr(provider, "hybrid_layer_pattern", None)) + + +def _recipe_model_flags(recipe_name: str, provider) -> tuple[str, ...]: + if recipe_name != RECIPE_NEMOTRON_H: + return () + from ..adapters.miles.models.nemotron_h import model_flags + + return model_flags(provider) + + def select_gdn_recipe(provider) -> GdnRecipe: return GdnRecipe( gated_delta_net=( @@ -836,6 +859,12 @@ def resolve_rl_run_config( "the Qwen3.8-Flash-Next recipe needs --no-sglang-deterministic-inference " "(FlashInfer GDN prefill is not supported with deterministic inference)" ) + elif _is_nemotron_h(provider): + from ..adapters.miles.models.nemotron_h import check_provider + + check_provider(provider, parameter_mode=parameter_mode) + # Miles run_nemotron_3_nano.py: --attention-backend auto + recipe_name, attention_backend = RECIPE_NEMOTRON_H, "auto" elif qwen35_recipe: recipe_name, attention_backend = RECIPE_QWEN3_5, "flash" else: @@ -973,6 +1002,7 @@ def resolve_rl_run_config( training_attention_backend=attention_backend, gdn=gdn, transformer_impl=transformer_impl, + model_flags=_recipe_model_flags(recipe_name, provider), ), geometry=geometry, parallel=ParallelLayout( From 391f7ed1e332bcf3df4b50beb0624d2e6c528251 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 04:52:41 +0000 Subject: [PATCH 168/233] Accept ungated (squared ReLU) MoE experts in the expert partition check Nemotron-H experts hold only the up projection in linear_fc1, so fc1 rows equal the expert FFN size instead of twice it. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_moe_full_parameter_layout.py | 18 ++++++++++++++++++ yeto/rl/moe_layout.py | 9 +++++++-- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/tests/test_rl_moe_full_parameter_layout.py b/tests/test_rl_moe_full_parameter_layout.py index 82108d0e4..60f2ecd4a 100644 --- a/tests/test_rl_moe_full_parameter_layout.py +++ b/tests/test_rl_moe_full_parameter_layout.py @@ -153,6 +153,24 @@ def test_missing_expert_and_router_mismatch_are_refused(): check_expert_partition(wrong_router) +def test_ungated_experts_are_accepted_and_bad_shapes_refused(): + # Nemotron-H experts are squared ReLU: fc1 holds only up, (ffn, hidden). + specs = qwen3_moe_specs(**TINY, tp=2, ep=2) + ungated = [ + ParameterSpec(s.role, s.name, (6, 16), s.dtype, 96, s.shard_id) + if ".linear_fc1." in s.name else s + for s in specs + ] + assert check_expert_partition(ungated).num_experts == 4 + odd = [ + ParameterSpec(s.role, s.name, (9, 16), s.dtype, 144, s.shard_id) + if ".linear_fc1." in s.name else s + for s in specs + ] + with pytest.raises(ValueError, match="not a gated or ungated MLP"): + check_expert_partition(odd) + + def owner_names_in_plan_order(lay: ParameterLayout) -> dict[str, list[str]]: names: dict[str, list[str]] = {} for fid in range(lay.fragments.num_fragments): diff --git a/yeto/rl/moe_layout.py b/yeto/rl/moe_layout.py index 94b622ffd..94b56aa8b 100644 --- a/yeto/rl/moe_layout.py +++ b/yeto/rl/moe_layout.py @@ -115,9 +115,14 @@ def check_expert_partition(specs: Iterable) -> ExpertPartition | None: if len(seen) != 1: raise ValueError(f"expert {projection} shapes differ across experts: {sorted(seen)}") (fc1,), (fc2,) = shapes["linear_fc1"], shapes["linear_fc2"] - if len(fc1) != 2 or len(fc2) != 2 or fc1[0] != 2 * fc2[1] or fc1[1] != fc2[0]: + # Gated experts (SwiGLU) stack gate and up in fc1; ungated ones (Nemotron-H + # squared ReLU) hold only up. + if ( + len(fc1) != 2 or len(fc2) != 2 + or fc1[0] not in (fc2[1], 2 * fc2[1]) or fc1[1] != fc2[0] + ): raise ValueError( - f"expert shapes are not a gated MLP: linear_fc1 {fc1}, linear_fc2 {fc2}" + f"expert shapes are not a gated or ungated MLP: linear_fc1 {fc1}, linear_fc2 {fc2}" ) by_layer: dict[tuple[str, int], dict[int, tuple[str, str]]] = {} From 7f5075334a30528ea72e0c1d8703c0cd74e08104 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 07:46:35 +0000 Subject: [PATCH 169/233] rl: per-fragment full-parameter apply and async pipelines up to 8 Apply: the Miles role stream applies one merged fragment at a time, but each apply prepared, hashed and copied every fragment of the model. With a Miles that takes skip_resident_fragments (Miles fork commit b13c170ce, michaellchung/miles branch yeto/s27-sync-apply), MilesChunkedFullParameterAdapter.apply asks the trainer to read and copy only the fragments whose bytes its optimizer masters do not already hold; the applied values and the commit count are unchanged. The capability is detected from the group's apply signature, so the pinned Miles (0b64b40c1) keeps the whole-model apply; YETO_FULL_PARAMETER_SKIP_RESIDENT=0 turns it off. A peer-snapshot restore drops the trainer's resident record before it writes. Pipeline: asynchronous full-parameter rounds accepted --rl-sync-fragment-pipeline 1 or 2. The syncer and the island pairing take any depth up to P, so the bound is now 8 (the default stays 2). Tests run the real syncer with 3, 4 and 8 rounds held open at once against the outer Nesterov reference, and the island pairing with depth 3, 4 and 8. No change to the parameter layout hash or the syncer wire protocol; the pipeline depth already enters the syncer profile hash. Co-Authored-By: Claude Opus 5.5 --- tests/test_e3_overlap_miles_role_stream.py | 90 +++++++++++++- tests/test_rl_full_parameter_async.py | 111 +++++++++++++++++- yeto/launch/rl_args.py | 13 +- yeto/rl/adapters/miles/cut_plugin.py | 15 +++ yeto/rl/adapters/miles/island_entry.py | 15 ++- .../miles/models/chunked_full_parameter.py | 29 +++++ yeto/rl/engine/run_flags.py | 10 +- 7 files changed, 269 insertions(+), 14 deletions(-) diff --git a/tests/test_e3_overlap_miles_role_stream.py b/tests/test_e3_overlap_miles_role_stream.py index 6a69f85fc..014c44cbb 100644 --- a/tests/test_e3_overlap_miles_role_stream.py +++ b/tests/test_e3_overlap_miles_role_stream.py @@ -327,7 +327,7 @@ def _config(sweeps: int) -> RoleStreamConfig: max_fragment_bytes=64, max_chunk_bytes=8, overlap=True) -def _run(monkeypatch, sweeps: int, *, cut=None): +def _run(monkeypatch, sweeps: int, *, cut=None, trainer_cls=None): """Run ``sweeps`` overlapped rounds; ``cut(rollout_id, syncer)`` runs while that round trains (the mid-overlap cut point).""" generator = torch.Generator().manual_seed(11) @@ -340,7 +340,7 @@ def rand(scale=1.0): peer = {k: rand(0.1) for k in range(sweeps)} ray = FakeRay() adapter = _adapter() - trainer = _Trainer(adapter, ray, initial, lambda k: steps[k]) + trainer = (trainer_cls or _Trainer)(adapter, ray, initial, lambda k: steps[k]) syncer = _Syncer(initial, sweeps, lambda k: peer[k]) client = _Client(syncer) @@ -502,3 +502,89 @@ def cut(rollout_id, syncer): run = _run(monkeypatch, sweeps, cut=cut) _check_exactly_once(run, sweeps) assert [step for step, attempt, _ in run.syncer.pushed if attempt == 2][0] == lost + + +class _ResidentTrainer(_Trainer): + """A trainer whose apply takes ``skip_resident_fragments`` like Miles. + + It records the payload hash of each fragment it last exported or applied + and skips a fragment with the same hash, after checking that the skipped + bytes really are the ones it holds (the role stream's precondition).""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.resident: dict[int, str] = {} + self.skipped = 0 + self.read = 0 + + def _remember(self, cut) -> None: + self.resident = {d.fragment_id: d.payload_hash for shard in cut.shards for d in shard.fragments} + + async def export_full_parameter_chunked_cut(self, *args, **kwargs): + cut = await super().export_full_parameter_chunked_cut(*args, **kwargs) + self._remember(cut) + return cut + + async def apply_full_parameter_chunked_cut(self, cut, *, commit_token=None, max_commit_attempts=3, + skip_resident_fragments=False): + assert commit_token and skip_resident_fragments + for shard in cut.shards: + for descriptor, refs in zip(shard.fragments, shard.chunk_refs, strict=True): + assert all(ref is not None for ref in refs), "applied a released chunk" + value = torch.cat([self.ray.get(ref) for ref in refs]).clone() + if self.resident.get(descriptor.fragment_id) == descriptor.payload_hash: + assert torch.equal(value.view(torch.int32), + self.values[descriptor.fragment_id].view(torch.int32)) + self.skipped += 1 + continue + self.read += 1 + self.values[descriptor.fragment_id] = value + self._remember(cut) + self.installed = cut.policy_version + self.generation = 0 + return self.adapter.expected_parameter_tensor_count + + +def test_skip_resident_is_requested_and_keeps_every_value_bit_identical(monkeypatch): + sweeps = 4 + plain = _run(monkeypatch, sweeps) + run = _run(monkeypatch, sweeps, trainer_cls=_ResidentTrainer) + _check_exactly_once(run, sweeps) + assert run.trainer.read > 0 + for f in range(P): + assert torch.equal(run.trainer.values[f].view(torch.int32), plain.trainer.values[f].view(torch.int32)) + assert torch.equal(run.syncer.global_[f].view(torch.int32), plain.syncer.global_[f].view(torch.int32)) + + +def test_skip_resident_follows_the_group_signature_and_the_kill_switch(monkeypatch): + from yeto.rl.adapters.miles.models.chunked_full_parameter import SKIP_RESIDENT_ENV, skip_resident_supported + + trainer = _ResidentTrainer(_adapter(), FakeRay(), {f: torch.zeros(n) for f, n in enumerate(SIZES)}, + lambda _k: {}) + assert skip_resident_supported(trainer) + assert not skip_resident_supported(_Trainer(_adapter(), FakeRay(), {}, lambda _k: {})) + assert not skip_resident_supported(object()) + monkeypatch.setenv(SKIP_RESIDENT_ENV, "0") + assert not skip_resident_supported(trainer) + + +def test_a_one_fragment_mixed_target_keeps_the_other_rows_resident(): + asyncio.run(_one_fragment_mixed_target()) + + +async def _one_fragment_mixed_target(): + """``assemble_mixed_target`` reuses the current cut's rows for every + fragment it does not replace, so only the replaced one is read.""" + adapter = _adapter() + ray = FakeRay() + initial = {f: torch.arange(n, dtype=torch.float32) + 10 * f for f, n in enumerate(SIZES)} + trainer = _ResidentTrainer(adapter, ray, initial, lambda _k: {}) + current = adapter.validate_cut( + await trainer.export_full_parameter_chunked_cut(0, 0, max_chunk_bytes=adapter.max_chunk_bytes), 0, 0) + replacement = torch.tensor([0.5, -0.0, 3.0e-39]) + stored = adapter.fragment_sink(current, ray_module=ray)(0, 1, replacement.numpy().tobytes()) + target = adapter.assemble_mixed_target(current, target_policy_version=1, authoritative_fragments=(stored,)) + assert await adapter.apply(trainer, target, commit_token="t") == adapter.expected_parameter_tensor_count + assert (trainer.read, trainer.skipped) == (1, P - 1) + assert torch.equal(trainer.values[0].view(torch.int32), replacement.view(torch.int32)) + assert torch.equal(trainer.values[1], initial[1]) diff --git a/tests/test_rl_full_parameter_async.py b/tests/test_rl_full_parameter_async.py index c83a9b587..0dc26fab3 100644 --- a/tests/test_rl_full_parameter_async.py +++ b/tests/test_rl_full_parameter_async.py @@ -37,6 +37,7 @@ def _launcher_args(engine, extra=(), gpu="aws:1xa100@us-east-1"): return _provenance_launcher_args(engine, ("--rl-sync-mode", "strict", *extra), gpu=gpu) from test_rl_full_parameter_sync import _candidate, _PairedRole, _role_stream_config from yeto.protocol import PullRequest +from yeto.rl.adapters.miles.full_parameter.role_stream import RoleSubmissionCandidate from yeto.rl.engine.full_parameter_sync import FullParameterSyncProfile, RolePairing LAYOUT = "e" * 64 @@ -452,7 +453,7 @@ def test_async_checkpoint_interval_reaches_the_syncer(monkeypatch): @pytest.mark.parametrize( ("extra", "message"), [ - (ASYNC + ("--rl-sync-fragment-pipeline", "3"), r"between 1 and min\(2, --fragments\)"), + (ASYNC + ("--rl-sync-fragment-pipeline", "5"), r"between 1 and min\(8, --fragments\)"), (ASYNC + ("--rl-sync-checkpoint-every", "0"), "checkpoint-every >= 1"), (ASYNC + ("--rl-sync-async-checkpoint",), "no --rl-sync-async-checkpoint"), (ASYNC + ("--rl-sync-overlap", "--rl-sync-fragment-pipeline", "2"), "no --rl-sync-overlap"), @@ -488,8 +489,15 @@ def test_island_checks_the_async_profile(): _check_dense_full_sweep(_island_args(), yeto_policy_sync=True) _check_dense_full_sweep(_island_args(rl_sync_matrix_merge="rda"), yeto_policy_sync=True) with pytest.raises(ValueError, match="--pipeline 1..min"): - _check_dense_full_sweep(_island_args(pipeline=3, fragments=4, total_fragment_steps=12), + _check_dense_full_sweep(_island_args(pipeline=5, fragments=4, total_fragment_steps=12), yeto_policy_sync=True) + with pytest.raises(ValueError, match=r"--pipeline 1..min\(8, P\)"): + _check_dense_full_sweep(_island_args(pipeline=9, fragments=60, total_fragment_steps=60, global_rounds=1), + yeto_policy_sync=True) + _check_dense_full_sweep(_island_args(pipeline=4, fragments=4, total_fragment_steps=12), + yeto_policy_sync=True) + _check_dense_full_sweep(_island_args(pipeline=8, fragments=60, total_fragment_steps=60, global_rounds=1), + yeto_policy_sync=True) with pytest.raises(ValueError, match="strict sweep requires H=1"): _check_dense_full_sweep(_island_args(full_parameter_rounds="sweep"), yeto_policy_sync=True) with pytest.raises(ValueError, match="plain averaging"): @@ -794,3 +802,102 @@ def test_async_streaming_weighted_fold_reaches_the_rl_syncer(monkeypatch): default = _launcher_args("ports", ASYNC, gpu=TWO_ISLANDS) _prepare_rl_args(default) assert "--weighted-fold" not in syncer_command(default, 2).split() + + +# -- deeper asynchronous pipelines (up to 8 fragment rounds in flight) -------- + + +def test_async_pipeline_default_stays_two_and_accepts_up_to_eight(): + _, plan = _plan(ASYNC) + assert _flag(plan["role_syncers"][0]["command"], "--pipeline") == "1" # strict-mode control: serial + _, deep = _plan(ASYNC + ("--rl-sync-fragment-pipeline", "4")) + assert _flag(deep["role_syncers"][0]["command"], "--pipeline") == "4" + for island in deep["island_requests"]: + assert _flag(island["learner_command"], "--pipeline") == "4" + + +def _deep_candidate(global_step: int, fragments: int): + return RoleSubmissionCandidate(role="actor", permit=PullRequest((global_step - 1) % fragments, global_step, 1), + base_version=max(0, global_step - fragments), local_step=2, c_steps=2, + c_tokens=11, payload_bytes=8, submission_cut=object(), anchor=object()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("depth", [3, 4, 8]) +async def test_deep_pipeline_submits_depth_rounds_and_applies_in_version_order(depth): + fragments = 8 + actor = _PairedRole("actor", [_deep_candidate(step, fragments) for step in range(1, fragments + 1)], [], + release_broadcasts=False) + actor.fragment_versions = (0,) * fragments + profile = FullParameterSyncProfile(roles=("actor",), total_fragment_steps=2 * fragments, fragments=fragments, + local_horizon=2, pipeline_depth=depth, wait_timeout=900.0, + poll_seconds=0.001, partial_quorum=True) + pairing = RolePairing({"actor": actor}, profile) + + first = _submitted(pairing.submit_ready()) + assert first == list(range(1, depth + 1)) + assert len(pairing.inflight) == depth + assert _submitted(pairing.submit_ready()) == [] # the pipeline is full + # The rounds close out of order; one closes without this island's push + # (a round it never pushed to is overwritten, never double-applied). + closed = list(reversed(first)) + actor._staged_steps.extend(closed) + actor._queued_steps.clear() + committed = await pairing.commit_staged_broadcasts() + assert committed["actor"].applied_fragments == tuple(range(depth)) + assert actor.fragment_versions[:depth] == tuple(range(1, depth + 1)) + assert actor.fragment_versions[depth:] == (0,) * (fragments - depth) + assert pairing.inflight == {} and len(pairing.journal) == depth and pairing.overwritten_rounds == 0 + # The next window opens with the following rounds. + assert _submitted(pairing.submit_ready()) == list(range(depth + 1, min(2 * depth, fragments) + 1)) + # A round this island never pushed (absent) is an overwrite in the deep window too. + absent = 2 * depth + 1 + if absent <= fragments: + actor.candidates = [c for c in actor.candidates if c.permit.global_step != absent] + actor._staged_steps.append(absent) + actor._staged_steps.extend(step for step in range(depth + 1, min(2 * depth, fragments) + 1)) + actor._queued_steps.clear() + await pairing.commit_staged_broadcasts() + assert actor.fragment_versions[absent - 1] == absent + assert pairing.overwritten_rounds == 1 + + +@needs_cargo +@pytest.mark.usefixtures("_hmac_key") +@pytest.mark.parametrize(("fragments", "depth"), [(4, 4), (8, 8), (8, 3)]) +def test_real_syncer_keeps_depth_rounds_open_and_merges_like_the_reference( + syncer_binary, tmp_path, monkeypatch, fragments, depth): + """Both learners hold their pushes until ``depth`` PULLs are open at once, + so the run only finishes if the syncer really keeps ``depth`` rounds in + flight; the merged weights match the outer Nesterov reference.""" + + import sys + + monkeypatch.setattr(sys.modules[__name__], "P", fragments) + rounds: dict[int, dict] = {} + held: dict[int, list] = {0: [], 1: []} + widest = [0, 0] + total = 2 * fragments + + def push(learner, permit): + delta = _delta(learner, permit.global_step) + rounds.setdefault(permit.global_step, {})[learner] = (1.0, delta) + held[learner].append((learner, permit, (_base(permit.global_step), 1, 7, delta))) + + def late(_versions): + out = [] + for learner in (0, 1): + rows = held[learner] + remaining = total - max((p.global_step for _l, p, _a in rows), default=0) + if rows and (len(rows) >= depth or remaining == 0): + widest[learner] = max(widest[learner], len(rows)) + out.extend(rows) + rows.clear() + return out + + finals, log = _drive(syncer_binary, tmp_path, ("--rl-sync-fragment-pipeline", str(depth)), + total_steps=2, push=push, late=late) + assert sorted(rounds) == list(range(1, total + 1)) + assert widest == [depth, depth] + _assert_matches(finals, _nesterov_reference(rounds, lr=0.7, mu=0.9)) + assert "push rejected" not in log diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 4ffb64915..97efe89a8 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -9,6 +9,7 @@ from pathlib import Path from ..gpu_spec import parse_gpu_spec +from ..rl.engine.run_flags import FULL_PARAMETER_ASYNC_DEFAULT_PIPELINE, FULL_PARAMETER_ASYNC_MAX_PIPELINE from ..rl.engine.selection import ( ASYNC_GRACE_CAP_MS, ASYNC_NESTEROV_MIN_H, @@ -199,8 +200,9 @@ def _check_bcast_delta(args) -> None: "pool options) and no --experimental-rl-sync") -# S27 P2: asynchronous rounds keep at most two fragment rounds in flight. -FULL_PARAMETER_ASYNC_MAX_PIPELINE = 2 +# S27 P2: asynchronous rounds keep FULL_PARAMETER_ASYNC_DEFAULT_PIPELINE (2) +# fragment rounds in flight by default and at most +# FULL_PARAMETER_ASYNC_MAX_PIPELINE (8); both live in rl.engine.run_flags. def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: @@ -238,7 +240,8 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: if async_rounds: max_pipeline = min(max_pipeline, FULL_PARAMETER_ASYNC_MAX_PIPELINE) if not 1 <= fragment_pipeline <= max_pipeline: - problems.append("--rl-sync-fragment-pipeline between 1 and min(2, --fragments)" if async_rounds + problems.append(f"--rl-sync-fragment-pipeline between 1 and min({FULL_PARAMETER_ASYNC_MAX_PIPELINE}, " + "--fragments)" if async_rounds else "--rl-sync-fragment-pipeline between 1 and --fragments") checkpoint_every = int(getattr(args, "rl_sync_checkpoint_every", 1)) if async_rounds: @@ -381,9 +384,9 @@ def _apply_rl_sync_mode(args) -> None: if (args.rl_sync_preset == FULL_PARAMETER_SYNC_PRESET and args.rl_sync_full_parameter_rounds == "async" and is_sync_mode_default(getattr(args, "rl_sync_fragment_pipeline", None))): - # At most two fragment rounds in flight; the fragment count is + # Two fragment rounds in flight by default; the fragment count is # checked (and may be 1) in _prepare_dense_full_sweep. - args.rl_sync_fragment_pipeline = min(FULL_PARAMETER_ASYNC_MAX_PIPELINE, + args.rl_sync_fragment_pipeline = min(FULL_PARAMETER_ASYNC_DEFAULT_PIPELINE, max(int(getattr(args, "fragments", 1) or 1), 1)) else: mode = "strict" diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 2fc2bde35..d190fe9c1 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -1385,10 +1385,24 @@ def rebase_full_parameter_tracking(actor: Any) -> None: rebase_full_parameter_scheduler_anchor(actor, reason="peer snapshot restore") +def forget_full_parameter_resident(actor: Any) -> None: + """A restore writes the optimizer masters: Miles must not treat any + fragment as already held by them (the resident-fragment apply skip). + A Miles without that skip has nothing to forget.""" + try: + from miles.backends.megatron_utils import full_parameter_state + except ImportError: + return + forget = getattr(full_parameter_state, "forget_full_parameter_resident_fragments", None) + if forget is not None: + forget(actor) + + def restore_cut_shard(actor: Any, *, directory: str, files: list[Mapping[str, Any]], cut_id: str) -> dict[str, Any]: """Verify and load this rank's shard; return the post-restore summary (identical digests expected).""" install_grad_norm_recorder() # a rebuilt trainer process has no recorder yet bump_weights_version() # rl-publish-fastpath: restore writes trainable weights + forget_full_parameter_resident(actor) _inject_restore_sleep(actor) progress = {"written": False} try: @@ -1579,6 +1593,7 @@ def restore_resharded_shard( """ install_grad_norm_recorder() bump_weights_version() # rl-publish-fastpath: restore writes trainable weights + forget_full_parameter_resident(actor) with trainer_resident(actor): return _restore_resharded(actor, directory=directory, files=files, cut_id=cut_id, source_dp=int(source_dp), rng_policy=rng_policy) diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 7c3a7bc6d..2779a014f 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -19,7 +19,12 @@ from yeto.rl.bridge import BridgeConfig from yeto.rl.decoupled import DecoupledBridgeConfig from yeto.rl.engine.resume import DEFAULT_CUT_EVERY -from yeto.rl.engine.run_flags import ISLAND, add_rl_flags, rl_option_pairs +from yeto.rl.engine.run_flags import ( + FULL_PARAMETER_ASYNC_MAX_PIPELINE, + ISLAND, + add_rl_flags, + rl_option_pairs, +) from yeto.rl.adapters.miles.models.deepseek_v4.expert_full import expert_full_specs from yeto.rl.export import adapter_targets, derive_peft_lora_specs from yeto.rl.adapters.miles.revision import verify_miles_revision @@ -801,12 +806,14 @@ def _check_dense_full_sweep(args, *, yeto_policy_sync: bool) -> None: _check_full_parameter_sync(args, yeto_policy_sync=yeto_policy_sync) if _full_parameter_async(args): # S27 P2 asynchronous rounds: H local rounds per fragment push, at - # most two fragment rounds in flight; no sweep-only options. + # most FULL_PARAMETER_ASYNC_MAX_PIPELINE fragment rounds in flight; + # no sweep-only options. if args.optimizer_steps != 1 or getattr(args, "local_horizon", 1) < 1: raise ValueError("dense full-parameter asynchronous rounds need H >= 1 " "and one optimizer step") - if not 1 <= getattr(args, "pipeline", 1) <= min(2, args.fragments): - raise ValueError("dense full-parameter asynchronous rounds need --pipeline 1..min(2, P)") + if not 1 <= getattr(args, "pipeline", 1) <= min(FULL_PARAMETER_ASYNC_MAX_PIPELINE, args.fragments): + raise ValueError("dense full-parameter asynchronous rounds need --pipeline " + f"1..min({FULL_PARAMETER_ASYNC_MAX_PIPELINE}, P)") if getattr(args, "sweep_overlap", False): raise ValueError("--sweep-overlap needs the strict sweep") elif args.optimizer_steps != 1 or getattr(args, "local_horizon", 1) != 1: diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index c1f1a56df..77adfb107 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -979,9 +979,15 @@ async def apply( ) -> int: if target.layout_hash != self.layout.layout_hash: raise ValueError("Miles target layout changed") + # The role stream builds every target from its current cut, which is + # the cut the trainer last exported or committed. A Miles that can + # skip resident fragments then reads and copies only the fragments + # this apply replaces; the applied values are the same. + extra = {"skip_resident_fragments": True} if skip_resident_supported(group) else {} applied = await group.apply_full_parameter_chunked_cut( target.transport_cut, commit_token=commit_token, + **extra, ) if applied != self.expected_parameter_tensor_count: raise RuntimeError("Miles applied an incomplete reference-backed cut") @@ -1002,6 +1008,29 @@ def _fragment_version(self, policy_version: int, fragment_id: int) -> int: ) +SKIP_RESIDENT_ENV = "YETO_FULL_PARAMETER_SKIP_RESIDENT" + + +def skip_resident_supported(group: Any) -> bool: + """Whether ``group``'s apply takes ``skip_resident_fragments`` (and it is not disabled). + + ``YETO_FULL_PARAMETER_SKIP_RESIDENT=0`` keeps the whole-model apply.""" + + import inspect + import os + + if os.environ.get(SKIP_RESIDENT_ENV, "1") == "0": + return False + apply = getattr(group, "apply_full_parameter_chunked_cut", None) + if apply is None: + return False + try: + parameters = inspect.signature(apply).parameters + except (TypeError, ValueError): + return False + return "skip_resident_fragments" in parameters + + def _fragment_owner_shard(cut: Any, fragment_id: int) -> Any: matches = [ shard diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 9a441b641..f11dfe294 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -46,6 +46,14 @@ ) CLI = "cli" +# S27 asynchronous full-parameter rounds: the default and the largest number +# of fragment rounds in flight (--rl-sync-fragment-pipeline / island +# --pipeline). The syncer and the island pairing take any depth up to P; the +# bound limits the retained submissions and staged broadcasts (one fragment +# each) an island holds at once. +FULL_PARAMETER_ASYNC_DEFAULT_PIPELINE = 2 +FULL_PARAMETER_ASYNC_MAX_PIPELINE = 8 + ISLAND = "island" LAYERS = (CLI, ISLAND) @@ -942,7 +950,7 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. type=int, default=sync_mode_default(1), help="full-parameter only: fragment rounds the syncer gathers at once " - "(1..--fragments; asynchronous rounds: 1 or 2). Above 1, fragment k+1 uploads " + "(1..--fragments; asynchronous rounds: 1..min(8, --fragments)). Above 1, fragment k+1 uploads " "while fragment k merges and downloads (default: min(2, --fragments) with " "asynchronous rounds, else 1 = serial)", ), From abcd7b7d28a7dbeb5093aebc5e8642bc4415e47e Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 08:18:28 +0000 Subject: [PATCH 170/233] Ports image: Miles b13c170ce (per-fragment full-parameter apply) Re-pin MILES_NEXT_COMMIT to b13c170ce and append it to MILES_FORK_LINEAGE. b13c170ce = 0b64b40c1 + a chunked full-parameter apply skips fragments the masters already hold (skip_resident). Image ghcr.io/michaellchung/yeto-miles-ports:b13c170-2fa8801 @sha256:814a7e276d160b5febdf0f64f526c2a6426dd47115e1abdec3db183fa18e05c0 (previous 0b64b40-2fa8801 @sha256:2b8c36f6...). Co-Authored-By: Claude Opus 5.5 --- docker/miles-ports/Dockerfile | 2 +- .../build-record.json | 12 +++++ .../image-manifest.json | 50 +++++++++++++++++++ tests/test_rl_ports_image.py | 2 +- yeto/rl/adapters/miles/pins.py | 12 +++-- 5 files changed, 72 insertions(+), 6 deletions(-) create mode 100644 openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/build-record.json create mode 100644 openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/image-manifest.json diff --git a/docker/miles-ports/Dockerfile b/docker/miles-ports/Dockerfile index 11f4c7d44..8392ea724 100644 --- a/docker/miles-ports/Dockerfile +++ b/docker/miles-ports/Dockerfile @@ -19,7 +19,7 @@ FROM docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c ARG MILES_REPOSITORY=https://github.com/michaellchung/miles -ARG MILES_COMMIT=0b64b40c15fa972d20bd1393aca6067dfd2492a4 +ARG MILES_COMMIT=b13c170ce7611104aee28e38ed01c188e95417a2 ARG SGLANG_REPOSITORY=https://github.com/michaellchung/sglang ARG SGLANG_COMMIT=2fa880182eefdb31e64d8ea70317be27f9f32121 # setuptools-scm version of SGLANG_COMMIT (describe v0.5.20-69-g4e4148f1b); diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/build-record.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/build-record.json new file mode 100644 index 000000000..4bbe73542 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/build-record.json @@ -0,0 +1,12 @@ +{ + "tag": "ghcr.io/michaellchung/yeto-miles-ports:b13c170-2fa8801", + "digest": "sha256:814a7e276d160b5febdf0f64f526c2a6426dd47115e1abdec3db183fa18e05c0", + "layer_tar_sha256": "251f83dfac97969fff844e620dafb092a84f178a62ebb160f84bb191bef551c7", + "base": "docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "base_index": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "miles_commit": "b13c170ce7611104aee28e38ed01c188e95417a2", + "sglang_commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "sglang_version": "0.5.21.dev70+g2fa8801", + "peft_version": "0.20.0", + "peft_wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/image-manifest.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/image-manifest.json new file mode 100644 index 000000000..574e3fb51 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/image-manifest.json @@ -0,0 +1,50 @@ +{ + "base": { + "index_digest": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "manifest_digest": "sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "miles_commit": "9e4260de047a704208535c0e90c531929879ab40", + "platform": "linux/amd64", + "repository": "docker.io/radixark/miles", + "sglang_commit": "e5bba1f9c433da5352e521a8bb826bab960e7836" + }, + "builder": "scripts/build_miles_ports_image.sh (crane append); docker/miles-ports/Dockerfile equivalent", + "image": "ghcr.io/michaellchung/yeto-miles-ports", + "megatron_bridge": { + "commit": "8cd3466d14d2337c8492827b3712482c2b3e4866", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_bridge_init.py)" + ] + }, + "megatron_lm": { + "commit": "a84b105473eca9c9c6e49dbf1c8aa460c5cf796e", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_lm_lazy_modelopt.patch)" + ], + "path": "/root/Megatron-LM", + "repository": "https://github.com/radixark/Megatron-LM" + }, + "miles": { + "commit": "b13c170ce7611104aee28e38ed01c188e95417a2", + "install": "base's editable install (pip install -e /root/miles --no-deps) kept (same path); /root is on the base's sys.path, so without PYTHONPATH import miles is a namespace package over the /root/miles checkout -- yeto runs with PYTHONPATH=$HOME/miles (regular package)", + "path": "/root/miles", + "repository": "https://github.com/michaellchung/miles", + "upstream_commit": "9e4260de047a704208535c0e90c531929879ab40" + }, + "peft": { + "install": "wheel unpacked into the base's site-packages (pip install --no-deps equivalent)", + "version": "0.20.0", + "wheel": "peft-0.20.0-py3-none-any.whl", + "wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" + }, + "pure_python_overlay": true, + "schema": 1, + "sglang": { + "commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "install": "pip install -e /sgl-workspace/sglang/python[all] --no-deps (base's editable install, same path)", + "path": "/sgl-workspace/sglang", + "repository": "https://github.com/michaellchung/sglang", + "upstream_commit": "571212b636baca45e10fa3b4da11a289123f3235", + "version": "0.5.21.dev70+g2fa8801" + }, + "tag": "b13c170-2fa8801" +} \ No newline at end of file diff --git a/tests/test_rl_ports_image.py b/tests/test_rl_ports_image.py index f0e38dd77..a23990181 100644 --- a/tests/test_rl_ports_image.py +++ b/tests/test_rl_ports_image.py @@ -70,7 +70,7 @@ def test_build_inputs_match_the_pins(): def test_commits_and_digest_are_pinned_together_with_the_build_record(): """Commit pins, tag, Dockerfile ARGs and the image digest move together: the latest build record must name exactly the pinned commits and digest.""" - record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801" + record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801" record = json.loads((record_dir / "build-record.json").read_text()) assert record["miles_commit"] == rl.MILES_NEXT_COMMIT assert record["sglang_commit"] == rl.SGLANG_NEXT_COMMIT diff --git a/yeto/rl/adapters/miles/pins.py b/yeto/rl/adapters/miles/pins.py index f46a0df99..9be42278d 100644 --- a/yeto/rl/adapters/miles/pins.py +++ b/yeto/rl/adapters/miles/pins.py @@ -57,8 +57,9 @@ class MilesRevisionPins(NamedTuple): # distributed-optimizer masters, incl. CPU offload; S27 G-B); cc693c693 = + the # Adam proof leaves the hybrid optimizer's master_param mirror out (S27 IMG2); # a477bcbe0 = + one parameter may span several fragments via flat ranges (S27); -# 0b64b40c1 = + rebase_full_parameter_scheduler_anchor for peer-snapshot recovery. -MILES_NEXT_COMMIT = "0b64b40c15fa972d20bd1393aca6067dfd2492a4" +# 0b64b40c1 = + rebase_full_parameter_scheduler_anchor for peer-snapshot recovery; +# b13c170ce = + a chunked apply skips fragments the masters already hold (S27 G-X). +MILES_NEXT_COMMIT = "b13c170ce7611104aee28e38ed01c188e95417a2" # sgl-project/sglang ``sglang-miles`` head when radixark/miles@9e4260d was # committed (upstream's Dockerfile follows that branch unpinned). SGLANG_NEXT_REPOSITORY = "https://github.com/michaellchung/sglang" @@ -129,6 +130,9 @@ class MilesRevisionPins(NamedTuple): # 0b64b40c1 = + rebase_full_parameter_scheduler_anchor after a peer snapshot # restore (S27); no loss path. "0b64b40c15fa972d20bd1393aca6067dfd2492a4", + # b13c170ce = + per-fragment full-parameter apply: a chunked apply skips + # fragments the masters already hold (S27 G-X). + "b13c170ce7611104aee28e38ed01c188e95417a2", ) # First Miles commit whose full-parameter sync reads distributed-optimizer @@ -173,7 +177,7 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # every SHA. PUBLIC on ghcr.io/michaellchung since 2026-10-07 (anonymous # pull); a private image needs SKYPILOT_DOCKER_{USERNAME,PASSWORD,SERVER} or # --rl-image-private (yeto.launcher.registry_login_for; read:packages token). -# Tag 0b64b40-2fa8801; (S27 IMG3: Miles 0b64b40c1 = cc693c693 + flat-range fragment split + scheduler anchor rebase; previous cc693c6-2fa8801 @sha256:2011cbf9... (a477bcb-2fa8801 built but voided); build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/). Earlier: tag cc693c6-2fa8801; (S27 IMG2: Miles cc693c693 = 56190a5a8 (full-parameter sync with distributed optimizer + CPU offload, G-B) + Adam proof without the hybrid master mirror; previous 56190a5-2fa8801 @sha256:5254e34e... (failed on GPU: global apply changed local Adam state), ed5f64f-2fa8801 @sha256:a4769254...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/). Earlier: tag ed5f64f-2fa8801; (S27 IMG: Miles ed5f64fd7, host-Adam round cuts (B8); previous 44d41ef-2fa8801 @sha256:5b9dc2a9...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-ed5f64f-2fa8801/). Earlier: tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt +# Tag b13c170-2fa8801; (S27 IMG4: Miles b13c170ce = 0b64b40c1 + per-fragment full-parameter apply (G-X); previous 0b64b40-2fa8801 @sha256:2b8c36f6...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/). Earlier: tag 0b64b40-2fa8801; (S27 IMG3: Miles 0b64b40c1 = cc693c693 + flat-range fragment split + scheduler anchor rebase; previous cc693c6-2fa8801 @sha256:2011cbf9... (a477bcb-2fa8801 built but voided); build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/). Earlier: tag cc693c6-2fa8801; (S27 IMG2: Miles cc693c693 = 56190a5a8 (full-parameter sync with distributed optimizer + CPU offload, G-B) + Adam proof without the hybrid master mirror; previous 56190a5-2fa8801 @sha256:5254e34e... (failed on GPU: global apply changed local Adam state), ed5f64f-2fa8801 @sha256:a4769254...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/). Earlier: tag ed5f64f-2fa8801; (S27 IMG: Miles ed5f64fd7, host-Adam round cuts (B8); previous 44d41ef-2fa8801 @sha256:5b9dc2a9...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-ed5f64f-2fa8801/). Earlier: tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt # Megatron-LM patch in the layer; previous 17aeb5c-2fa8801 @sha256:b304f857...; # build record openspec/changes/rl-infra-spec/evidence/ports-image/ # 2026-10-14-44d41ef-2fa8801/). Earlier: tag 17aeb5c-2fa8801; (S24 L2.1-L2.4: Miles s24-miles-import, peft 0.20.0 and the @@ -182,7 +186,7 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # ports-image/2026-10-14-17aeb5c-2fa8801/). Earlier: tag 386ae3b-2fa8801; (S23 P0-1.11 GPU run: Miles s23-sao-streaming-roles; previous 64b591a-2fa8801 @sha256:fa2413be...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-386ae3b-2fa8801/). Earlier: (S19 #1: critic family + algo-supplement metrics baked in; previous ddce209-2fa8801 @sha256:9c252c38...) (S18 ARU-3: Miles s18-agentic-suspend = 2f7871fb2 + --agentic-suspend-between-turns; previous 2f7871f-2fa8801 @sha256:62b4f164..., efbbc63-2fa8801 @sha256:a7990076..., 8bc5223-4e4148f @sha256:4aeafd77...) ghcr.io/michaellchung (public package); linux/amd64 only. MILES_NEXT_IMAGE = ( "docker:ghcr.io/michaellchung/yeto-miles-ports@sha256:" - "2b8c36f6cbdeddcbc3d867d23808b94a12a20217ce9fd2eb55ba94aa2c12f7f0" + "814a7e276d160b5febdf0f64f526c2a6426dd47115e1abdec3db183fa18e05c0" ) MILES_NEXT_IMAGE_MANIFEST = "/opt/yeto/image-manifest.json" # Nebius VM images whose /var/lib/docker already holds a docker image's From 8848adb1ba9d0b1b2607cd52809e23b5e86f9972 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:57:47 +0000 Subject: [PATCH 171/233] rl: detect the resident-fragment apply through Miles' forwarding group handle s27-gr-c ran Miles b13c170ce but every apply still prepared and copied the whole model (prepare 150-245 s, commit 95-160 s per merged fragment). The role stream's group is Miles' RayWorkerHandle, whose __getattr__ returns a generic call(*args, **kwargs); the signature check saw no skip_resident_fragments parameter and kept the whole-model apply. With pipeline 8 an island then needed about 8 x 5-6 min to apply one round's broadcasts and pushed nothing for the next rounds before the stall watchdog. A forwarding signature is now decided by the Miles installed in the island process (the trainer's image): full_parameter_transport.apply_chunked_cut must name the parameter. The decision is printed once per reason ("[yeto-full-parameter] skip resident fragments: ..."). Regression tests drive the apply through a RayWorkerHandle-like forwarding group. Co-Authored-By: Claude Opus 5.5 --- tests/test_e3_overlap_miles_role_stream.py | 53 ++++++++++++++++++ .../miles/models/chunked_full_parameter.py | 56 ++++++++++++++++--- 2 files changed, 102 insertions(+), 7 deletions(-) diff --git a/tests/test_e3_overlap_miles_role_stream.py b/tests/test_e3_overlap_miles_role_stream.py index 014c44cbb..9cc9e55f5 100644 --- a/tests/test_e3_overlap_miles_role_stream.py +++ b/tests/test_e3_overlap_miles_role_stream.py @@ -588,3 +588,56 @@ async def _one_fragment_mixed_target(): assert (trainer.read, trainer.skipped) == (1, P - 1) assert torch.equal(trainer.values[0].view(torch.int32), replacement.view(torch.int32)) assert torch.equal(trainer.values[1], initial[1]) + + +class _ForwardingHandle: + """Like Miles' RayWorkerHandle: every method is a generic ``call(*args, **kwargs)``.""" + + def __init__(self, target) -> None: + self._target = target + + def __getattr__(self, name): + async def call(*args, **kwargs): + return await getattr(self._target, name)(*args, **kwargs) + return call + + +def _local_miles_skips() -> bool: + import inspect + + from miles.ray import full_parameter_transport + + return "skip_resident_fragments" in inspect.signature(full_parameter_transport.apply_chunked_cut).parameters + + +def test_a_forwarding_group_uses_the_local_miles_to_decide_the_skip(monkeypatch): + """s27-gr-c: the production group forwards ``*args, **kwargs``, so the + signature check alone kept the whole-model apply on every rank.""" + from yeto.rl.adapters.miles.models.chunked_full_parameter import SKIP_RESIDENT_ENV, skip_resident_supported + + trainer = _ResidentTrainer(_adapter(), FakeRay(), {f: torch.zeros(n) for f, n in enumerate(SIZES)}, + lambda _k: {}) + handle = _ForwardingHandle(trainer) + assert skip_resident_supported(handle) is _local_miles_skips() + monkeypatch.setenv(SKIP_RESIDENT_ENV, "0") + assert not skip_resident_supported(handle) + + +def test_one_fragment_apply_through_a_forwarding_group_reads_one_fragment(): + if not _local_miles_skips(): + pytest.skip("this Miles has no skip_resident_fragments") + + async def body(): + adapter = _adapter() + ray = FakeRay() + initial = {f: torch.arange(n, dtype=torch.float32) + 10 * f for f, n in enumerate(SIZES)} + trainer = _ResidentTrainer(adapter, ray, initial, lambda _k: {}) + handle = _ForwardingHandle(trainer) + current = adapter.validate_cut( + await handle.export_full_parameter_chunked_cut(0, 0, max_chunk_bytes=adapter.max_chunk_bytes), 0, 0) + stored = adapter.fragment_sink(current, ray_module=ray)(1, 2, torch.tensor([4.0, -2.0]).numpy().tobytes()) + target = adapter.assemble_mixed_target(current, target_policy_version=1, authoritative_fragments=(stored,)) + await adapter.apply(handle, target, commit_token="t") + assert (trainer.read, trainer.skipped) == (1, P - 1) + + asyncio.run(body()) diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index 77adfb107..95a1bbe41 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -1011,24 +1011,66 @@ def _fragment_version(self, policy_version: int, fragment_id: int) -> int: SKIP_RESIDENT_ENV = "YETO_FULL_PARAMETER_SKIP_RESIDENT" +_SKIP_RESIDENT_REPORTED: set[str] = set() + + +def _takes_skip_resident(function: Any) -> bool | None: + """True/False when ``function``'s signature names or lacks the + parameter; None when it only forwards ``*args, **kwargs``.""" + + import inspect + + try: + parameters = inspect.signature(function).parameters + except (TypeError, ValueError): + return False + if "skip_resident_fragments" in parameters: + return True + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()): + return None + return False + + def skip_resident_supported(group: Any) -> bool: """Whether ``group``'s apply takes ``skip_resident_fragments`` (and it is not disabled). + The production group is Miles' ``RayWorkerHandle`` around the + TrainerController actor: its ``__getattr__`` returns a generic + ``call(*args, **kwargs)`` that forwards to the remote method, so its + signature says nothing (s27-gr-c ran the whole-model apply for this + reason). A forwarding signature is then decided by the Miles installed in + this process, which is the trainer's Miles (same image): + ``full_parameter_transport.apply_chunked_cut`` must name the parameter. + ``YETO_FULL_PARAMETER_SKIP_RESIDENT=0`` keeps the whole-model apply.""" - import inspect import os if os.environ.get(SKIP_RESIDENT_ENV, "1") == "0": - return False + return _report_skip_resident(False, f"{SKIP_RESIDENT_ENV}=0") apply = getattr(group, "apply_full_parameter_chunked_cut", None) if apply is None: - return False + return _report_skip_resident(False, "group has no chunked apply") + takes = _takes_skip_resident(apply) + if takes is not None: + return _report_skip_resident(takes, "group apply signature") try: - parameters = inspect.signature(apply).parameters - except (TypeError, ValueError): - return False - return "skip_resident_fragments" in parameters + from miles.ray import full_parameter_transport + except ImportError: + return _report_skip_resident(False, "forwarding group and no local Miles transport") + local = _takes_skip_resident(full_parameter_transport.apply_chunked_cut) is True + return _report_skip_resident(local, "forwarding group; local Miles transport") + + +def _report_skip_resident(enabled: bool, reason: str) -> bool: + key = f"{enabled}:{reason}" + if key not in _SKIP_RESIDENT_REPORTED: + _SKIP_RESIDENT_REPORTED.add(key) + import sys + + print(f"[yeto-full-parameter] skip resident fragments: " + f"{'enabled' if enabled else 'disabled'} ({reason})", file=sys.stderr, flush=True) + return enabled def _fragment_owner_shard(cut: Any, fragment_id: int) -> Any: From c4186d5320080cdb958743222417a128c2dd3be9 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:23:13 +0000 Subject: [PATCH 172/233] rl/miles: allow-listed trainer-speed flags for Nemotron-H MFU runs YETO_MILES_TRAIN_EXTRA_ARGV (island env, e.g. launcher --modal-env) appends allow-listed Megatron flags (recompute granularity/modules, MoE dispatcher, permute fusion, CE fusion, timers) to the Nemotron-H trainer argv. None of them changes TP/EP/PP/CP/DP, so the full-parameter layout stays the same. Empty or unset: no change. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_miles_train_extra_argv.py | 27 +++++++++++++++++++++++ yeto/rl/adapters/miles/config.py | 29 +++++++++++++++++++++++++ 2 files changed, 56 insertions(+) create mode 100644 tests/test_rl_miles_train_extra_argv.py diff --git a/tests/test_rl_miles_train_extra_argv.py b/tests/test_rl_miles_train_extra_argv.py new file mode 100644 index 000000000..1749cff2b --- /dev/null +++ b/tests/test_rl_miles_train_extra_argv.py @@ -0,0 +1,27 @@ +"""S27 G-Y: YETO_MILES_TRAIN_EXTRA_ARGV allow-list (MFU measurement flags, Nemotron-H only).""" +import pytest + +from yeto.rl.adapters.miles.config import MilesConfigError, train_extra_argv + + +def test_empty_is_noop_for_any_recipe(): + assert train_extra_argv("generic", "") == () + assert train_extra_argv("nemotron_h", " ") == () + + +def test_allowed_flags_pass_through_in_order(): + raw = ("--recompute-granularity selective --recompute-modules moe layernorm " + "--moe-token-dispatcher-type alltoall --moe-permute-fusion") + assert train_extra_argv("nemotron_h", raw) == tuple(raw.split()) + + +@pytest.mark.parametrize("raw", ["--tensor-model-parallel-size 1", "selective", + "--moe-permute-fusion --expert-model-parallel-size 4"]) +def test_layout_or_unknown_flags_are_refused(raw): + with pytest.raises(MilesConfigError): + train_extra_argv("nemotron_h", raw) + + +def test_other_recipes_refuse_a_set_value(): + with pytest.raises(MilesConfigError): + train_extra_argv("generic", "--moe-permute-fusion") diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index cbc65f521..57379b22f 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -1175,6 +1175,7 @@ def translate_run_config( values = glm5_next.apply_recipe(values, glm5_next.model_argv( **shape, transformer_impl=recipe.transformer_impl)) values.extend(extra_argv) + values.extend(train_extra_argv(recipe.name, os.environ.get(TRAIN_EXTRA_ENV, ""))) reject_fault_tolerance_flags(values[1:]) # generated argv never carries FT flags _refuse_decoupled_ppo_breakers(values) # also when the spec emitted the flag # IR-1: the final argv never pairs the session server with partial rollout @@ -1217,6 +1218,34 @@ def translate_run_config( "--optimizer-cpu-offload", "--overlap-cpu-optimizer-d2h-h2d") +# S27 G-Y: trainer-speed flags for MFU measurement on the Nemotron-H recipe, from the +# island environment (launcher --modal-env). Only allow-listed Megatron flags pass; +# none of them changes the parallel layout (TP/EP/PP/CP/DP). +TRAIN_EXTRA_ENV = "YETO_MILES_TRAIN_EXTRA_ARGV" +TRAIN_EXTRA_FLAGS = frozenset({ + "--recompute-granularity", "--recompute-modules", "--recompute-method", + "--recompute-num-layers", "--moe-token-dispatcher-type", "--moe-permute-fusion", + "--cross-entropy-loss-fusion", "--timing-log-level", "--timing-log-option", +}) + + +def train_extra_argv(recipe_name: str, raw: str) -> tuple[str, ...]: + """Parse ``YETO_MILES_TRAIN_EXTRA_ARGV``; empty unless set. Nemotron-H only.""" + import shlex + + from yeto.rl.engine.run_config import RECIPE_NEMOTRON_H + + tokens = shlex.split(raw or "") + if not tokens: + return () + if recipe_name != RECIPE_NEMOTRON_H: + raise MilesConfigError(f"{TRAIN_EXTRA_ENV} is only for the {RECIPE_NEMOTRON_H} recipe") + bad = [t for t in tokens if t.startswith("--") and t not in TRAIN_EXTRA_FLAGS] + if bad or not tokens[0].startswith("--"): + raise MilesConfigError(f"{TRAIN_EXTRA_ENV}: flags not allowed: {bad or tokens[:1]}") + return tuple(tokens) + + def train_memory_argv(memory: Any) -> tuple[str, ...]: """Megatron full uniform recompute (one layer per chunk), Miles dynamic batching (the same values the validated Flash-Next recipe uses) and the From 388bfe2508ca2f372212ded6d4bea8a405dc016a Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:23:31 +0000 Subject: [PATCH 173/233] rl/miles: train extra argv also splits on commas (one launcher word) Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_miles_train_extra_argv.py | 5 +++++ yeto/rl/adapters/miles/config.py | 4 ++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_rl_miles_train_extra_argv.py b/tests/test_rl_miles_train_extra_argv.py index 1749cff2b..69d2ffa4d 100644 --- a/tests/test_rl_miles_train_extra_argv.py +++ b/tests/test_rl_miles_train_extra_argv.py @@ -25,3 +25,8 @@ def test_layout_or_unknown_flags_are_refused(raw): def test_other_recipes_refuse_a_set_value(): with pytest.raises(MilesConfigError): train_extra_argv("generic", "--moe-permute-fusion") + + +def test_commas_split_like_spaces(): + assert train_extra_argv("nemotron_h", "--recompute-modules,moe,layernorm") == ( + "--recompute-modules", "moe", "layernorm") diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 57379b22f..2dc9a0f4a 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -1230,12 +1230,12 @@ def translate_run_config( def train_extra_argv(recipe_name: str, raw: str) -> tuple[str, ...]: - """Parse ``YETO_MILES_TRAIN_EXTRA_ARGV``; empty unless set. Nemotron-H only.""" + """Parse ``YETO_MILES_TRAIN_EXTRA_ARGV`` (space- or comma-separated); empty unless set. Nemotron-H only.""" import shlex from yeto.rl.engine.run_config import RECIPE_NEMOTRON_H - tokens = shlex.split(raw or "") + tokens = shlex.split((raw or "").replace(",", " ")) # commas: one shell word if not tokens: return () if recipe_name != RECIPE_NEMOTRON_H: From 681b41449a299e7b3319a3189c2a39d17cc6a443 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:35:41 +0000 Subject: [PATCH 174/233] rl: island wait timeouts, async pipelines up to 60, terminal retire - RolePairing._check_timeouts: a paired round whose broadcast has arrived (staged now, before the check) or was applied is not a timeout; every push of a pair (a replay on a reopened round included) restarts its wait. s27-gr-d island 0 failed on a step whose broadcast arrived during its serial pushes. - --rl-sync-wait-timeout-s (default 900, whole seconds): the island's paired-broadcast wait, terminal barrier and FINAL wait (island --sweep-wait-timeout-s -> RoleStreamConfig.wait_timeout). A value other than 900 also sets the full-parameter syncer's --final-ack-timeout-s, which is part of the syncer profile hash; layout hash, session contract and training contract are unchanged. - FULL_PARAMETER_ASYNC_MAX_PIPELINE 8 -> 60: one boundary can push a whole P=60 model round (--total-steps 1 --rl-sync-fragment-pipeline 60). - Terminal retire: once the island trains no more (any_finalizing), a fragment's last round is closed without a trainer apply (retire_staged_broadcast records the anchor and version as an apply would); the FINAL apply overwrites every fragment, so the final parameters are unchanged. Off with peer recovery, overlap or YETO_FULL_PARAMETER_TERMINAL_RETIRE=0. Tests: timeout rules, retire rules, plan of the whole-round fixed roster options, a real Rust syncer merging a 60-fragment round at once (quorum 2, grace 0, max base lag 0), and the real Miles role stream with and without the retire (bit-identical final values, one trainer apply instead of P). Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_full_parameter_async.py | 222 +++++++++++++++++- .../test_s27_whole_round_miles_role_stream.py | 166 +++++++++++++ yeto/launch/rl_args.py | 6 +- yeto/launch/rl_island_task.py | 2 + yeto/launch/syncer.py | 23 +- .../miles/full_parameter/role_stream.py | 110 +++++++-- yeto/rl/adapters/miles/island_entry.py | 9 + yeto/rl/adapters/miles/run_flags.py | 7 + yeto/rl/engine/full_parameter_sync.py | 60 ++++- yeto/rl/engine/run_flags.py | 20 +- 10 files changed, 591 insertions(+), 34 deletions(-) create mode 100644 tests/test_s27_whole_round_miles_role_stream.py diff --git a/tests/test_rl_full_parameter_async.py b/tests/test_rl_full_parameter_async.py index 0dc26fab3..7448f2bb2 100644 --- a/tests/test_rl_full_parameter_async.py +++ b/tests/test_rl_full_parameter_async.py @@ -23,6 +23,7 @@ import time from dataclasses import replace from pathlib import Path +from types import SimpleNamespace import numpy as np import pytest @@ -453,7 +454,10 @@ def test_async_checkpoint_interval_reaches_the_syncer(monkeypatch): @pytest.mark.parametrize( ("extra", "message"), [ - (ASYNC + ("--rl-sync-fragment-pipeline", "5"), r"between 1 and min\(8, --fragments\)"), + (ASYNC + ("--rl-sync-fragment-pipeline", "5"), r"between 1 and min\(60, --fragments\)"), + (ASYNC + ("--rl-sync-wait-timeout-s", "0"), "wait-timeout-s in whole seconds"), + (ASYNC + ("--rl-sync-wait-timeout-s", "12.5"), "wait-timeout-s in whole seconds"), + (("--rl-sync-wait-timeout-s", "3600"), "only applies to --rl-sync-preset full-parameter"), (ASYNC + ("--rl-sync-checkpoint-every", "0"), "checkpoint-every >= 1"), (ASYNC + ("--rl-sync-async-checkpoint",), "no --rl-sync-async-checkpoint"), (ASYNC + ("--rl-sync-overlap", "--rl-sync-fragment-pipeline", "2"), "no --rl-sync-overlap"), @@ -491,9 +495,11 @@ def test_island_checks_the_async_profile(): with pytest.raises(ValueError, match="--pipeline 1..min"): _check_dense_full_sweep(_island_args(pipeline=5, fragments=4, total_fragment_steps=12), yeto_policy_sync=True) - with pytest.raises(ValueError, match=r"--pipeline 1..min\(8, P\)"): - _check_dense_full_sweep(_island_args(pipeline=9, fragments=60, total_fragment_steps=60, global_rounds=1), + with pytest.raises(ValueError, match=r"--pipeline 1..min\(60, P\)"): + _check_dense_full_sweep(_island_args(pipeline=61, fragments=64, total_fragment_steps=64, global_rounds=1), yeto_policy_sync=True) + _check_dense_full_sweep(_island_args(pipeline=60, fragments=60, total_fragment_steps=60, global_rounds=1), + yeto_policy_sync=True) _check_dense_full_sweep(_island_args(pipeline=4, fragments=4, total_fragment_steps=12), yeto_policy_sync=True) _check_dense_full_sweep(_island_args(pipeline=8, fragments=60, total_fragment_steps=60, global_rounds=1), @@ -901,3 +907,213 @@ def late(_versions): assert widest == [depth, depth] _assert_matches(finals, _nesterov_reference(rounds, lr=0.7, mu=0.9)) assert "push rejected" not in log + + +# -- S27 G-AA: island wait timeouts, pipeline 60, terminal retire --------------- + + +def _timed_pairing(*streams, wait_timeout, total=6, fragments=2, depth=2): + profile = FullParameterSyncProfile(roles=tuple(stream.role for stream in streams), total_fragment_steps=total, + fragments=fragments, local_horizon=2, pipeline_depth=depth, + wait_timeout=wait_timeout, poll_seconds=0.001, partial_quorum=True) + return RolePairing({stream.role: stream for stream in streams}, profile) + + +class _ArrivingRole(_PairedRole): + """Broadcasts arrive on the client (``arrived``) before the pairing stages them.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, release_broadcasts=False, **kwargs) + self.arrived: list[int] = [] + + def stage_available_broadcasts_now(self): + self._staged_steps.extend(self.arrived) + self.arrived.clear() + + async def stage_available_broadcasts(self): + self.stage_available_broadcasts_now() + return await super().stage_available_broadcasts() + + +def test_a_broadcast_that_arrived_before_the_check_is_not_a_timeout(): + # s27-gr-d island 0: step 8's broadcast arrived while the island pushed + # its next round; the check ran before the apply and failed the island. + actor = _ArrivingRole("actor", [_candidate("actor", 1)], []) + pairing = _timed_pairing(actor, wait_timeout=1.0) + assert _submitted(pairing.submit_ready()) == [1] + (pair,) = pairing.inflight.values() + pair.retained_at -= 5.0 # the island spent longer than the wait on local work + actor.arrived.append(1) + + pairing.submit_ready() # no TimeoutError: the round closed + + assert actor.staged_fragment_versions == ((0, 1),) + assert len(pairing.inflight) == 1 # applied at the next commit, not here + + +def test_a_missing_broadcast_still_times_out(): + actor = _ArrivingRole("actor", [_candidate("actor", 1)], []) + pairing = _timed_pairing(actor, wait_timeout=1.0) + pairing.submit_ready() + next(iter(pairing.inflight.values())).retained_at -= 5.0 + with pytest.raises(TimeoutError, match="paired WAN broadcast"): + pairing.submit_ready() + assert actor.closed + + +def test_a_replayed_push_restarts_the_wait(): + # The syncer reopened the round under a new attempt (OutOfRound for the + # old push); the replay is a new push and the wait starts again. + actor = _ArrivingRole("actor", [_candidate("actor", 1)], []) + pairing = _timed_pairing(actor, wait_timeout=1.0) + pairing.submit_ready() + (pair,) = pairing.inflight.values() + pair.retained_at -= 5.0 + newer = _candidate("actor", 1) + actor.candidates.append(replace(newer, permit=replace(newer.permit, round_attempt=2))) + + produced = pairing.submit_ready() + + assert _submitted(produced) == [1] + assert pair.candidates["actor"].permit.round_attempt == 2 + assert time.monotonic() - pair.retained_at < 1.0 + + +class _RetiringRole(_PairedRole): + retires_terminal_broadcasts = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.retired: list[tuple[int, int]] = [] + self.applied: list[tuple[int, int]] = [] + + async def apply_staged_broadcast(self, fragment_id, version, *, catch_up=False): + self.applied.append((fragment_id, version)) + return await super().apply_staged_broadcast(fragment_id, version, catch_up=catch_up) + + async def retire_staged_broadcast(self, fragment_id, version, *, catch_up=False): + self.retired.append((fragment_id, version)) + return await super().apply_staged_broadcast(fragment_id, version, catch_up=catch_up) + + +@pytest.mark.asyncio +async def test_the_terminal_phase_retires_last_rounds_instead_of_applying_them(): + # total 6, P 2: steps 5 and 6 are the last rounds; 6 arrives as FINAL. + actor = _RetiringRole("actor", [], []) + actor.fragment_versions = (3, 2) + pairing = _timed_pairing(actor, wait_timeout=900.0) + actor._staged_steps.extend([4, 5]) + await pairing.commit_staged_broadcasts() + assert actor.applied == [(1, 4), (0, 5)] and actor.retired == [] # still training: apply + + actor.fragment_versions = (3, 2) + actor.applied.clear() + pairing.terminal_seen = True # no more training on this island + actor._staged_steps.extend([4, 5]) + await pairing.commit_staged_broadcasts() + assert actor.applied == [(1, 4)] # not a last round: applied as before + assert actor.retired == [(0, 5)] + assert actor.fragment_versions == (5, 4) and pairing.retired_rounds == 1 + + +@pytest.mark.asyncio +async def test_no_retire_unless_every_role_supports_it(): + actor, critic = _RetiringRole("actor", [], []), _PairedRole("critic", [], []) + for role in (actor, critic): + role.fragment_versions = (3, 4) + role._staged_steps.append(5) + pairing = _timed_pairing(actor, critic, wait_timeout=900.0) + pairing.terminal_seen = True + await pairing.commit_staged_broadcasts() + assert actor.retired == [] and actor.applied == [(0, 5)] + + +def test_wait_timeout_reaches_the_island_and_the_final_ack(): + from yeto import launcher + + base = ASYNC + ("--rl-sync-fragment-pipeline", "2") + args, plan = _plan(base) + _, longer = _plan(base + ("--rl-sync-wait-timeout-s", "3600")) + island = longer["island_requests"][0]["learner_command"] + assert _flag(island, "--sweep-wait-timeout-s") == "3600" + assert "--sweep-wait-timeout-s" not in plan["island_requests"][0]["learner_command"] + syncer = longer["role_syncers"][0]["command"] + assert _flag(syncer, "--final-ack-timeout-s") == "3600" + assert "--quorum-timeout-s" not in syncer # the quorum timeout keeps its default + assert "--final-ack-timeout-s" not in plan["role_syncers"][0]["command"] + # The final-ACK timeout is part of the syncer profile, so its hash moves. + assert _flag(island, "--syncer-profile-sha256") != _flag( + plan["island_requests"][0]["learner_command"], "--syncer-profile-sha256") + assert launcher.rl_syncer_profile(args, 2).sha256 == _flag( + plan["island_requests"][0]["learner_command"], "--syncer-profile-sha256") + + +def test_the_island_passes_the_wait_timeout_to_its_role_stream(): + from yeto.rl.adapters.miles.island_entry import _sweep_wait_timeout_s + + assert _sweep_wait_timeout_s(SimpleNamespace()) == 900.0 + assert _sweep_wait_timeout_s(SimpleNamespace(sweep_wait_timeout_s=3600.0)) == 3600.0 + with pytest.raises(ValueError, match="sweep-wait-timeout-s"): + _sweep_wait_timeout_s(SimpleNamespace(sweep_wait_timeout_s=-1.0)) + + +def test_whole_round_fixed_roster_plan_is_accepted(): + """The s27-gr-f options: one boundary pushes all 60 fragments to a fixed + roster syncer (quorum 2 of 2, no grace, exact base).""" + extra = ("--fragments", "60", "--rl-sync-full-parameter-rounds", "async", "--total-steps", "1", + "--rl-sync-fragment-pipeline", "60", "--rl-sync-quorum", "2", "--rl-sync-grace-ms", "0", + "--rl-sync-max-base-lag", "0", "--rl-sync-wait-timeout-s", "3600") + args, plan = _plan(FULL[:-2] + extra) + syncer = plan["role_syncers"][0]["command"] + for flag, value in (("--pipeline", "60"), ("--quorum", "2"), ("--grace-ms", "0"), ("--max-base-lag", "0"), + ("--total-steps", "60"), ("--final-ack-timeout-s", "3600")): + assert _flag(syncer, flag) == value, flag + for island in plan["island_requests"]: + command = island["learner_command"] + assert _flag(command, "--pipeline") == "60" + assert _flag(command, "--sweep-wait-timeout-s") == "3600" + assert args.rl_total_fragment_steps == 60 + + +@needs_cargo +@pytest.mark.usefixtures("_hmac_key") +def test_real_fixed_roster_syncer_merges_a_whole_60_fragment_round_at_once(syncer_binary, tmp_path, monkeypatch): + """s27-gr-f: ``--total-steps 1 --rl-sync-fragment-pipeline 60`` with a fixed + roster (quorum 2 of 2, grace 0, max base lag 0). Both learners hold every + push until all 60 PULLs are open, so the run finishes only if the syncer + opens the whole round at once; then every step merges both learners, no + push is refused and no round is reopened.""" + + import sys + + fragments = 60 + monkeypatch.setattr(sys.modules[__name__], "P", fragments) + rounds: dict[int, dict] = {} + held: dict[int, list] = {0: [], 1: []} + widest = [0, 0] + + def push(learner, permit): + assert permit.round_attempt == 1 + delta = _delta(learner, permit.global_step) + rounds.setdefault(permit.global_step, {})[learner] = (1.0, delta) + held[learner].append((learner, permit, (0, 1, 7, delta))) + + def late(_versions): + out = [] + for learner in (0, 1): + if len(held[learner]) == fragments: + widest[learner] = fragments + out.extend(held[learner]) + held[learner].clear() + return out + + extra = ("--rl-sync-fragment-pipeline", "60", "--rl-sync-quorum", "2", "--rl-sync-grace-ms", "0", + "--rl-sync-max-base-lag", "0", "--rl-sync-learner-weight", "equal", + "--rl-sync-outer-lr", "1.0", "--rl-sync-outer-momentum", "0", "--rl-sync-wait-timeout-s", "3600") + finals, log = _drive(syncer_binary, tmp_path, extra, total_steps=1, push=push, late=late) + assert sorted(rounds) == list(range(1, fragments + 1)) + assert all(sorted(contributions) == [0, 1] for contributions in rounds.values()) + assert widest == [fragments, fragments] + _assert_matches(finals, _nesterov_reference(rounds, lr=1.0, mu=0.0)) + assert "push rejected" not in log and "OutOfRound" not in log + assert "new frozen-membership attempt" not in log diff --git a/tests/test_s27_whole_round_miles_role_stream.py b/tests/test_s27_whole_round_miles_role_stream.py new file mode 100644 index 000000000..262700be6 --- /dev/null +++ b/tests/test_s27_whole_round_miles_role_stream.py @@ -0,0 +1,166 @@ +"""S27 G-AA: one boundary pushes a whole model round on the real Miles role stream. + +``--total-steps 1 --rl-sync-fragment-pipeline P`` (s27-gr-f): the syncer opens +every fragment round at once, the island pushes all P fragments at its first +boundary and then trains no more. Every broadcast is then a fragment's last +round, followed only by the FINAL cut, so the pairing retires those rounds +instead of applying each one to the trainer. The final trainer values must be +bit-identical to the run that applies every broadcast. + +The coordinator drives the real ``MilesFullParameterRoleStream`` and +``MilesChunkedFullParameterAdapter``; only the trainer group, Ray and the +syncer are fakes (``test_e3_overlap_miles_role_stream``). + +Needs Miles (``importorskip``). +""" + +from __future__ import annotations + +# ruff: noqa: I001 -- the Miles test helpers install a bridge stub before import. + +import asyncio +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("miles") + +from test_e3_overlap_miles_role_stream import ( + OTHER_WEIGHT, + P, + SIZES, + _adapter, + _Client, + _Driver, + _Syncer, + _Trainer, +) +from test_rl_miles_chunked_full_parameter import FakeRay, _evidence, component + +from yeto.protocol import PullRequest +from yeto.rl.adapters.miles.full_parameter import role_stream as role_stream_module +from yeto.rl.adapters.miles.full_parameter.role_stream import ( + TERMINAL_RETIRE_ENV, + MilesFullParameterRoleStream, + RoleStreamConfig, +) +from yeto.rl.adapters.miles.models.chunked_full_parameter import ( + MilesChunkedFullParameterAdapter, +) +from yeto.rl.engine.full_parameter_sync import ( + FullParameterSync, + FullParameterSyncProfile, + RolePairing, +) + + +class _AsyncSyncer(_Syncer): + """Asynchronous fragment rounds, fixed roster of this island and one + virtual peer: a round merges as soon as this island pushed it and is + broadcast at once; all rounds open at INIT (pipeline P); the last step + arrives as FINAL only.""" + + def commit_ready(self) -> None: + for step in sorted(self.received): + delta = self.received.pop(step) + f, k = (step - 1) % P, (step - 1) // P + self.global_[f] = self.global_[f] + OTHER_WEIGHT * (delta + self.peer(k)[f]) + self.versions[f] = step + self.committed.append(step) + if step < self.sweeps * P: + self.client.broadcast(f, step, self.global_[f]) + if len(self.committed) == self.sweeps * P and not self.client.finalizing.is_set(): + self.client.publish_final() + + +class _WholeRoundClient(_Client): + def pull_sweep(self, sweep, attempt) -> None: + # pipeline P: every round of the (only) model round opens at INIT. + if sweep == 0: + self.pulls.extend(PullRequest(f, f + 1, attempt) for f in range(P)) + + +class _CountingTrainer(_Trainer): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.applies: list[tuple[int, ...]] = [] + + async def apply_full_parameter_chunked_cut(self, cut, *, commit_token=None, max_commit_attempts=3): + self.applies.append(tuple(sorted(d.fragment_id for shard in cut.shards for d in shard.fragments))) + return await super().apply_full_parameter_chunked_cut( + cut, commit_token=commit_token, max_commit_attempts=max_commit_attempts) + + +def _config() -> RoleStreamConfig: + adapter = _adapter() + return RoleStreamConfig( + algorithm="grpo", role="actor", component=component(), syncer_addr=("127.0.0.1", 30100), + learner_id=0, learner_generation=0, learner_generations=(0, 0), + total_fragment_steps=P, expected_fragments=P, + expected_layout_hash=adapter.layout.layout_hash, local_horizon=1, optimizer_steps_per_round=1, + training_contract_hash="2" * 64, syncer_profile_hash="3" * 64, pipeline_depth=P, + algorithm_spec_sha256="a1" * 32, wait_timeout=5.0, poll_seconds=0.001, + max_fragment_bytes=64, max_chunk_bytes=8, partial_quorum=True) + + +def _run(monkeypatch, *, retire: bool): + monkeypatch.setenv(TERMINAL_RETIRE_ENV, "1" if retire else "0") + generator = torch.Generator().manual_seed(23) + + def rand(scale=1.0): + return {f: torch.randn(n, generator=generator) * scale for f, n in enumerate(SIZES)} + + initial, step, peer = rand(), rand(0.1), rand(0.1) + ray = FakeRay() + adapter = _adapter() + trainer = _CountingTrainer(adapter, ray, initial, lambda _k: step) + syncer = _AsyncSyncer(initial, 1, lambda _k: peer) + client = _WholeRoundClient(syncer) + + async def initialize(cls, group, *, policy_version, **_kwargs): + cut = await group.export_full_parameter_chunked_cut(policy_version, 0, + max_chunk_bytes=adapter.max_chunk_bytes) + return adapter, adapter.validate_cut(cut, policy_version, 0) + + monkeypatch.setattr(MilesChunkedFullParameterAdapter, "initialize", classmethod(initialize)) + monkeypatch.setattr(role_stream_module, "SyncerClient", lambda *_args, **_kwargs: client) + stream = MilesFullParameterRoleStream(_config(), ray_module=ray) + profile = FullParameterSyncProfile.from_configs({"actor": stream.config}) + assert profile.partial_quorum and profile.effective_pipeline_depth == P + session = FullParameterSync( + RolePairing({"actor": stream}, profile), groups={"actor": trainer}, run=asyncio.run, + trajectories=lambda rollout_id: _evidence(rollout_id=rollout_id, + behavior_policy_hash=stream.current.policy_hash)) + driver = _Driver() + session.start(driver) + boundary = session.boundary(driver, rollout_id=0, stats=None) + assert boundary.stop # one boundary: every last round was ready + session.published(driver, rollout_id=0, policy_hash=boundary.state.policy_tensor_hash()) + session.finish(driver) + return SimpleNamespace(trainer=trainer, syncer=syncer, stream=stream, driver=driver, + pairing=session.pairing, initial=initial, step=step, peer=peer) + + +def test_the_terminal_retire_keeps_the_final_parameters_bit_identical(monkeypatch): + applied = _run(monkeypatch, retire=False) + retired = _run(monkeypatch, retire=True) + + for run in (applied, retired): + assert run.syncer.committed == list(range(1, P + 1)) + assert run.stream.final_applied and run.stream.final_acknowledged + for f in range(P): + assert torch.equal(run.trainer.values[f], run.syncer.global_[f]) # FINAL applied + for f in range(P): + assert torch.equal(retired.trainer.values[f].view(torch.int32), applied.trainer.values[f].view(torch.int32)) + assert torch.equal(retired.syncer.global_[f], applied.syncer.global_[f]) + # The pushed deltas are the same bytes in both runs. + assert [(s, a, d.view(torch.int32).tolist()) for s, a, d in retired.syncer.pushed] == [ + (s, a, d.view(torch.int32).tolist()) for s, a, d in applied.syncer.pushed] + # Applied: each non-terminal broadcast, then FINAL. Retired: FINAL only. + assert applied.trainer.applies == [(0, 1)] * P + assert retired.trainer.applies == [(0, 1)] + assert retired.pairing.retired_rounds == P - 1 and applied.pairing.retired_rounds == 0 + assert retired.stream.retired_fragments == set(range(P - 1)) + rounds = [fields for name, fields in retired.driver.events if name == "rl_full_parameter_round"] + assert rounds[-1]["sync/retired_rounds"] == P - 1 diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 97efe89a8..3d8e177d1 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -173,6 +173,7 @@ def sao_contract_mounts(args, learner_id: int) -> dict[str, str]: ("rl_sync_full_parameter_rounds", "sweep", "--rl-sync-full-parameter-rounds"), ("rl_sync_expert_token_weights", False, "--rl-sync-expert-token-weights"), ("rl_sync_init_mode", "auto", "--rl-sync-init-mode"), + ("rl_sync_wait_timeout_s", 900.0, "--rl-sync-wait-timeout-s"), ) @@ -202,7 +203,7 @@ def _check_bcast_delta(args) -> None: # S27 P2: asynchronous rounds keep FULL_PARAMETER_ASYNC_DEFAULT_PIPELINE (2) # fragment rounds in flight by default and at most -# FULL_PARAMETER_ASYNC_MAX_PIPELINE (8); both live in rl.engine.run_flags. +# FULL_PARAMETER_ASYNC_MAX_PIPELINE (60); both live in rl.engine.run_flags. def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: @@ -243,6 +244,9 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: problems.append(f"--rl-sync-fragment-pipeline between 1 and min({FULL_PARAMETER_ASYNC_MAX_PIPELINE}, " "--fragments)" if async_rounds else "--rl-sync-fragment-pipeline between 1 and --fragments") + wait_timeout = float(getattr(args, "rl_sync_wait_timeout_s", 900.0)) + if not 0 < wait_timeout <= 7 * 86400 or wait_timeout != int(wait_timeout): + problems.append("--rl-sync-wait-timeout-s in whole seconds within 1..604800") checkpoint_every = int(getattr(args, "rl_sync_checkpoint_every", 1)) if async_rounds: # The async checkpoint writer and the sync/train overlap are diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 97c76089d..1f11bf6a3 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -121,6 +121,8 @@ def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: sweep += " --sweep-overlap" if getattr(args, "rl_sync_push_no_tensor_blocks", False): sweep += " --sweep-push-no-tensor-blocks" + if float(getattr(args, "rl_sync_wait_timeout_s", 900.0)) != 900.0: + sweep += f" --sweep-wait-timeout-s {float(args.rl_sync_wait_timeout_s):g}" if getattr(args, "rl_sync_fragment_max_bytes", 2 << 30) != 2 << 30: sweep += f" --sweep-fragment-max-bytes {int(args.rl_sync_fragment_max_bytes)}" if getattr(args, "rl_sync_push_dtype", "f32") == "int8": diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index fe9d3f6d5..9cc96a5f5 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -170,9 +170,26 @@ def _rl_quorum_timeout_s(args) -> int | None: def _syncer_quorum_timeout(args) -> str: """The syncer's --quorum-timeout-s, the same value the island's pause - budget uses; "" by default (syncer 900 s).""" + budget uses; "" by default (syncer 900 s). A full-parameter + ``--rl-sync-wait-timeout-s`` other than 900 also sets --final-ack-timeout-s.""" value = _rl_quorum_timeout_s(args) - return "" if value is None else f" --quorum-timeout-s {value}" + final_ack = _rl_final_ack_timeout_s(args) + return (("" if value is None else f" --quorum-timeout-s {value}") + + ("" if final_ack is None else f" --final-ack-timeout-s {final_ack}")) + + +def _rl_final_ack_timeout_s(args) -> int | None: + """The full-parameter syncer's final-ACK wait: ``--rl-sync-wait-timeout-s`` + when set (the islands wait as long for FINAL and apply it before they ACK); + None keeps the syncer default (its quorum timeout).""" + if not rl_dense_sweep(args): + return None + value = float(getattr(args, "rl_sync_wait_timeout_s", SYNCER_DEFAULT_QUORUM_TIMEOUT_S)) + if value == SYNCER_DEFAULT_QUORUM_TIMEOUT_S: + return None + if value != int(value): + raise ValueError("--rl-sync-wait-timeout-s must be whole seconds (syncer u64)") + return int(value) def _syncer_checkpoint_every(args) -> int: @@ -263,7 +280,7 @@ def rl_syncer_profile(args, num_learners: int): if rl_dense_sweep(args): quorum_timeout = _rl_quorum_timeout_s(args) or SYNCER_DEFAULT_QUORUM_TIMEOUT_S - final_ack_timeout = quorum_timeout + final_ack_timeout = _rl_final_ack_timeout_s(args) or quorum_timeout sweep_fragments = None if rl_full_parameter_async(args) else args.fragments mark_final = False elif rl_sao_streaming(args): diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 3ed3a0511..b3d57c1de 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -13,6 +13,7 @@ import asyncio import hashlib import json +import os import sys import time from collections.abc import Mapping, Sequence @@ -82,6 +83,9 @@ def _never_materialize(_policy_version: int): raise UnsupportedLayoutError(LAYOUT_FRAGMENTS) +TERMINAL_RETIRE_ENV = "YETO_FULL_PARAMETER_TERMINAL_RETIRE" + + def require_sha256(name: str, value: object) -> str: if ( type(value) is not str @@ -434,6 +438,9 @@ def __init__(self, config: RoleStreamConfig, *, ray_module=None) -> None: self.submission_rounds = 0 self.submission_tokens = 0 self.transport_version = 0 + # Terminal phase: fragments whose last round closed without a trainer + # apply (``retire_staged_broadcast``); the FINAL apply covers them. + self.retired_fragments: set[int] = set() self.local_generation = 0 self.local_rounds = 0 self.action_tokens = 0 @@ -978,6 +985,12 @@ async def complete_local_round( async def stage_available_broadcasts(self) -> tuple[tuple[int, int], ...]: """Drain complete broadcasts without mutating the trainable model.""" + return self.stage_available_broadcasts_now() + + def stage_available_broadcasts_now(self) -> tuple[tuple[int, int], ...]: + """:meth:`stage_available_broadcasts` without the event loop (the + pairing's timeout check stages arrived broadcasts before it judges).""" + if self.client is None: raise RuntimeError(f"{self.config.algorithm} {self.config.role} client is unavailable") self.client.check_health() @@ -1036,8 +1049,31 @@ async def apply_staged_broadcast( island never received (a re-JOIN catch-up); the global value overwrites it like any closed round.""" - key = (fragment_id, version) - stored = self._staged.get(key) + stored = self._staged_for_apply(fragment_id, version, rebase=rebase, catch_up=catch_up) + if not rebase: + await self._apply_authoritative({fragment_id: stored}) + else: + if not self.config.overlap or self._round_base is None or self.current is None: + raise RuntimeError(f"{self.config.role} rebase without an overlap round base") + rebased = self.adapter.rebase_fragment( + stored, self.current, self._round_base, ray_module=self.ray_module) + try: + await self._apply_authoritative( + {fragment_id: rebased}, + anchors={fragment_id: stored}, + counters=(self._rounds_at_base, self._tokens_at_base, self._routed_at_base), + ) + finally: + rebased.release() # the applied target holds its own references + self._close_staged_round(fragment_id, version) + return (fragment_id,) + + def _staged_for_apply( + self, fragment_id: int, version: int, *, rebase: bool, catch_up: bool + ) -> StoredAuthoritativeFragment: + """The staged broadcast ``(fragment_id, version)`` once its base checks pass.""" + + stored = self._staged.get((fragment_id, version)) if stored is None: raise RuntimeError( f"{self.config.algorithm} {self.config.role} matching staged fragment disappeared" @@ -1054,22 +1090,10 @@ async def apply_staged_broadcast( raise RuntimeError( f"{self.config.algorithm} {self.config.role} staged fragment skipped its base version" ) - if not rebase: - await self._apply_authoritative({fragment_id: stored}) - else: - if not self.config.overlap or self._round_base is None or self.current is None: - raise RuntimeError(f"{self.config.role} rebase without an overlap round base") - rebased = self.adapter.rebase_fragment( - stored, self.current, self._round_base, ray_module=self.ray_module) - try: - await self._apply_authoritative( - {fragment_id: rebased}, - anchors={fragment_id: stored}, - counters=(self._rounds_at_base, self._tokens_at_base, self._routed_at_base), - ) - finally: - rebased.release() # the applied target holds its own references - self._staged.pop(key, None) + return stored + + def _close_staged_round(self, fragment_id: int, version: int) -> None: + self._staged.pop((fragment_id, version), None) for older in [k for k in self._staged if k[0] == fragment_id and k[1] < version]: self._staged.pop(older).release() if self.config.partial_quorum: @@ -1077,6 +1101,56 @@ async def apply_staged_broadcast( permit = self._pending.get(fragment_id) if permit is not None and permit.global_step <= version: self._pending.pop(fragment_id) + + @property + def retires_terminal_broadcasts(self) -> bool: + """Whether :meth:`retire_staged_broadcast` may skip a terminal apply. + + Only for a partial quorum without overlap and without peer recovery + (a donor would serve trainer state behind its fragment versions); + ``YETO_FULL_PARAMETER_TERMINAL_RETIRE=0`` keeps every apply.""" + + return ( + self.config.partial_quorum + and not self.config.overlap + and self.config.peer_recovery is None + and os.environ.get(TERMINAL_RETIRE_ENV, "1") != "0" + ) + + async def retire_staged_broadcast( + self, + fragment_id: int, + version: int, + *, + catch_up: bool = False, + ) -> tuple[int, ...]: + """Terminal phase (partial quorum): close a fragment's last round + without writing it into the trainer. + + The pairing calls this only after the island stopped training and + only for a round with no later round of that fragment: nothing reads + the trainer's value of the fragment again before the FINAL cut, and + the FINAL apply overwrites every fragment. The broadcast becomes the + fragment's anchor and version exactly as an apply would record them, + so pushes, closed-round checks and the FINAL manifest see the same + state; only the trainer apply (whole-model costs per fragment) is + left out. The final parameters are the FINAL bytes either way.""" + + if not self.retires_terminal_broadcasts: + raise RuntimeError(f"{self.config.role} cannot retire a broadcast without applying it") + if version < self.config.total_fragment_steps - self.config.expected_fragments + 1: + raise RuntimeError(f"{self.config.role} retires only a fragment's last round") + stored = self._staged_for_apply(fragment_id, version, rebase=False, catch_up=catch_up) + previous_anchor = self._anchors[fragment_id] + if previous_anchor is not None and previous_anchor is not stored: + self._release_anchor(previous_anchor) + self._anchors[fragment_id] = stored + self._fragment_versions[fragment_id] = stored.version + self._rounds_at_anchor[fragment_id] = self.local_rounds + self._tokens_at_anchor[fragment_id] = self.action_tokens + self._routed_at_anchor[fragment_id] = self.action_routed + self.retired_fragments.add(fragment_id) + self._close_staged_round(fragment_id, version) return (fragment_id,) def finish_rebase(self) -> None: diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 2779a014f..dd4586520 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -931,6 +931,7 @@ def _dense_full_sync_config( syncer_profile_hash=args.syncer_profile_sha256, pipeline_depth=int(args.pipeline) if async_rounds else 1, algorithm_spec_sha256=algorithm_spec_sha256, + wait_timeout=_sweep_wait_timeout_s(args), wan_streams=args.wan_streams, wan_stream_schedule=getattr(args, "wan_stream_schedule", "rr"), wan_stall_timeout=float(getattr(args, "wan_stall_timeout", 0.0) or 0.0), @@ -956,6 +957,14 @@ def _dense_full_sync_config( return FullParameterSyncConfig((actor,), sweep=not async_rounds) +def _sweep_wait_timeout_s(args) -> float: + """``--sweep-wait-timeout-s`` (launcher ``--rl-sync-wait-timeout-s``); 900 s by default.""" + value = float(getattr(args, "sweep_wait_timeout_s", 900.0) or 900.0) + if not 0 < value <= 7 * 86400: + raise ValueError("--sweep-wait-timeout-s must be within (0, 604800] seconds") + return value + + PEER_RECOVER_ENV = "YETO_RL_PEER_RECOVER" # "1": the launcher relaunched this island (launch.peer_recovery) diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 1036225ea..4f93307fb 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -357,6 +357,13 @@ help="full-parameter: f32 byte bound of one sync fragment; a larger " "parameter spans several fragments as dimension-0 row slices", ), + island_flag( + "--sweep-wait-timeout-s", + type=float, + default=900.0, + help="full-parameter sync: seconds to wait for a paired broadcast after its " + "last push, for the terminal barrier and for the FINAL cut", + ), island_flag( "--sweep-push-int8-block", type=int, diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index f876c6fd8..f88f83be7 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -277,6 +277,8 @@ def __init__(self, streams: Mapping[str, RoleStream], profile: FullParameterSync self.overwritten_rounds = 0 self.caught_up_rounds = 0 self.caught_up_fragments = 0 # catch-up jumps (one per fragment and jump) + # Terminal phase: last rounds closed without a trainer apply (_retires). + self.retired_rounds = 0 # -- fail-stop --------------------------------------------------------- @@ -440,6 +442,9 @@ def _push_once( if submission is not None: pair.submissions[role] = submission produced[role].append(submission) + # The wait for this round's broadcast starts at its last push: + # a replay answers a reopened round (new attempt). + pair.retained_at = time.monotonic() def _check_timeouts(self, by_role: Mapping[str, Mapping[LogicalKey, Any]], now: float) -> None: seen = set().union(*(set(keys) for keys in by_role.values())) @@ -449,13 +454,28 @@ def _check_timeouts(self, by_role: Mapping[str, Mapping[LogicalKey, Any]], now: timeout = self.profile.wait_timeout if any(now - started >= timeout for started in self._pull_skew_since.values()): raise TimeoutError("timed out waiting for a peer role's matching PULL") - for pair in self.inflight.values(): - if now - pair.retained_at < timeout: + expired = [(key, pair) for key, pair in self.inflight.items() if now - pair.retained_at >= timeout] + if not expired: + return + # A broadcast that arrived while this island trained, captured or + # pushed closes its round; it is applied at the next commit, so it + # is not a missing broadcast. Stage what has arrived, then decide. + self._stage_arrived_broadcasts() + for (step, fragment_id, _base), pair in expired: + if self._round_closed(fragment_id, step): continue if pair.complete: raise TimeoutError("timed out waiting for a paired WAN broadcast") raise TimeoutError("timed out with a one-sided paired submission in flight") + def _stage_arrived_broadcasts(self) -> None: + """Stage arrived broadcasts now (no apply) on streams that can do it + synchronously (``stage_available_broadcasts_now``).""" + for stream in self.streams.values(): + stage = getattr(stream, "stage_available_broadcasts_now", None) + if stage is not None: + stage() + # -- apply (rule 3) ------------------------------------------------------ def _round_closed(self, fragment_id: int, version: int) -> bool: @@ -536,16 +556,39 @@ async def _commit_partial_quorum(self, applied: dict[str, list[int]]) -> None: self.caught_up_rounds += fragment_rounds(fragment_id, base, fragments) - ( fragment_rounds(fragment_id, current, fragments)) + def _retires(self, version: int) -> bool: + """Terminal phase: close ``version`` without a trainer apply. + + Once ``any_finalizing`` holds the island trains no more, and a + fragment's last round (``version >= total - P + 1``) is followed only + by the FINAL cut, which overwrites every fragment. Nothing reads the + trainer's value of that fragment in between, so its apply cannot + change the final parameters; streams record the round as closed + (``retire_staged_broadcast``) instead. Every role must support it.""" + + return ( + self.profile.partial_quorum + and self.any_finalizing + and version >= self.profile.total_fragment_steps - self.profile.fragments + 1 + and all(getattr(stream, "retires_terminal_broadcasts", False) + for stream in self.streams.values()) + ) + async def _overwrite_or_apply_pair( self, fragment_id: int, version: int, applied: dict[str, list[int]], *, catch_up: bool ) -> None: closed = [(key, pair) for key, pair in sorted(self.inflight.items()) if key[1] == fragment_id and key[0] <= version] exact = next((entry for entry in closed if entry[0][0] == version), None) + retire = self._retires(version) if not catch_up and exact is not None and exact[1].complete: - await self._apply_pair(exact[0], exact[1], applied) + await self._apply_pair(exact[0], exact[1], applied, retire=retire) return - if catch_up: + if retire: + self.retired_rounds += 1 + calls = (stream.retire_staged_broadcast(fragment_id, version, catch_up=catch_up) + for stream in self.streams.values()) + elif catch_up: calls = (stream.apply_staged_broadcast(fragment_id, version, catch_up=True) for stream in self.streams.values()) else: @@ -575,12 +618,16 @@ def _is_whole_sweep(self, ready: list[LogicalKey]) -> bool: and {key[1] for key in sweep} == set(range(self.profile.fragments))) async def _apply_pair( - self, key: LogicalKey, pair: PairedSubmission, applied: dict[str, list[int]] + self, key: LogicalKey, pair: PairedSubmission, applied: dict[str, list[int]], + *, retire: bool = False, ) -> None: step, fragment_id, _base = key if not pair.complete: raise RuntimeError("broadcast arrived before every paired PUSH was retained") - if self.profile.overlap: + if retire: + self.retired_rounds += 1 + calls = (stream.retire_staged_broadcast(fragment_id, step) for stream in self.streams.values()) + elif self.profile.overlap: calls = (stream.apply_staged_broadcast(fragment_id, step, rebase=True) for stream in self.streams.values()) else: @@ -778,6 +825,7 @@ def _partial_quorum_fields(self, versions: tuple[int, ...]) -> dict[str, int]: "sync/overwritten_fragments": overwritten, "sync/caught_up_rounds": pairing.caught_up_rounds, "sync/caught_up_fragments": pairing.caught_up_fragments, + "sync/retired_rounds": pairing.retired_rounds, "sync/outer_version": merged_version(versions), } diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index f11dfe294..4f1c9103c 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -50,9 +50,14 @@ # of fragment rounds in flight (--rl-sync-fragment-pipeline / island # --pipeline). The syncer and the island pairing take any depth up to P; the # bound limits the retained submissions and staged broadcasts (one fragment -# each) an island holds at once. +# each) an island holds at once. 60 lets the probed P=60 layouts push a whole +# model round at once (``--total-steps 1``: one boundary submits every +# fragment). FULL_PARAMETER_ASYNC_DEFAULT_PIPELINE = 2 -FULL_PARAMETER_ASYNC_MAX_PIPELINE = 8 +FULL_PARAMETER_ASYNC_MAX_PIPELINE = 60 +# The island's wait for a paired broadcast, the terminal barrier and the FINAL +# cut (--rl-sync-wait-timeout-s / island --sweep-wait-timeout-s). +FULL_PARAMETER_DEFAULT_WAIT_TIMEOUT_S = 900.0 ISLAND = "island" LAYERS = (CLI, ISLAND) @@ -950,7 +955,7 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. type=int, default=sync_mode_default(1), help="full-parameter only: fragment rounds the syncer gathers at once " - "(1..--fragments; asynchronous rounds: 1..min(8, --fragments)). Above 1, fragment k+1 uploads " + "(1..--fragments; asynchronous rounds: 1..min(60, --fragments)). Above 1, fragment k+1 uploads " "while fragment k merges and downloads (default: min(2, --fragments) with " "asynchronous rounds, else 1 = serial)", ), @@ -984,6 +989,15 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "dimension-0 row slices, each merged as its own tensor (RDA per slice). " "A smaller bound changes the layout hash, so the probe must use the same value", ), + cli_flag( + "--rl-sync-wait-timeout-s", + type=float, + default=FULL_PARAMETER_DEFAULT_WAIT_TIMEOUT_S, + help="full-parameter only: seconds an island waits for a paired broadcast " + "after its last push of that round, for the terminal barrier and for the " + "FINAL cut; a value other than 900 also sets the syncer's final-ACK " + "timeout (which changes the syncer profile hash, not the layout hash)", + ), cli_flag( "--rl-sync-push-int8-block", type=int, From b8f41f60d56c480753aa7bae3cbf0e3d607c596f Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:37:25 +0000 Subject: [PATCH 175/233] rl ports: --context-parallel for the Miles trainer (S27 G-AC long sequences) Launcher flag --context-parallel N is forwarded to the island, enters the trainer layout (DP = world / (TP x PP x CP)), the placement request and Miles' --context-parallel-size. --rl-max-tokens-per-gpu x CP must cover --seq-len (Miles packs max_tokens_per_gpu x CP tokens per micro batch). Co-Authored-By: Claude Opus 5.5 --- yeto/cli.py | 3 +++ yeto/launch/rl_args.py | 4 +++- yeto/launch/rl_island_task.py | 2 ++ yeto/launch/rl_ports.py | 2 +- yeto/rl/adapters/miles/config.py | 4 ++-- yeto/rl/adapters/miles/run_flags.py | 1 + yeto/rl/engine/run_config.py | 15 ++++++++++++--- 7 files changed, 24 insertions(+), 7 deletions(-) diff --git a/yeto/cli.py b/yeto/cli.py index 4ba41308e..a23278bb7 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -190,6 +190,9 @@ def loss_spec(value: str) -> str: ) tune.add_argument("--tensor-parallel", type=int, default=1, help="megatron backend: TP degree") tune.add_argument("--pipeline-parallel", type=int, default=1, help="megatron backend: PP degree") + tune.add_argument("--context-parallel", type=int, default=1, + help="RL ports (Miles) trainer: context-parallel degree (S27 G-AC long sequences); " + "the trainer world is TP x PP x CP x DP") tune.add_argument( "--train-on", choices=["assistant", "all"], diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 97efe89a8..0059112dd 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -926,7 +926,9 @@ def _prepare_rl_args( raise ValueError("RL v0 does not support external learner slots") if args.tensor_parallel <= 0 or args.pipeline_parallel <= 0: raise ValueError("RL tensor and pipeline parallelism must be positive") - model_parallel = args.tensor_parallel * args.pipeline_parallel + if int(getattr(args, "context_parallel", 1) or 1) <= 0: + raise ValueError("RL context parallelism must be positive") + model_parallel = args.tensor_parallel * args.pipeline_parallel * int(getattr(args, "context_parallel", 1) or 1) if args.expert_parallel is not None: if args.expert_parallel <= 0: raise ValueError("RL expert parallelism must be positive") diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 97c76089d..5fbee656b 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -272,6 +272,8 @@ def make_miles_island_task( flags += algorithm_flags if args.expert_parallel is not None: flags += f" --expert-parallel {args.expert_parallel}" + if int(getattr(args, "context_parallel", 1) or 1) > 1: + flags += f" --context-parallel {int(args.context_parallel)}" for flag, name in ( ("--rl-prompt-column", "rl_prompt_column"), ("--rl-label-column", "rl_label_column"), diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index 37e208ee2..672246bf5 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -64,7 +64,7 @@ def _prepare_ports_algorithm(args) -> None: problems += launch_problems(spec, { "rollout_batch_size": getattr(args, "rollout_batch_size", None), "rollout_max_response_len": getattr(args, "rollout_max_response_len", None), - "context_parallel_size": 1, # ports emits --context-parallel-size 1 + "context_parallel_size": int(getattr(args, "context_parallel", 1) or 1), "multi_lora": False, # rl-algo-critic-family 2.3 (critic run-level rejections) "sync_preset": getattr(args, "rl_sync_preset", "strict-avg"), diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index cbc65f521..79fc213d9 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -752,7 +752,7 @@ def translate_run_config( problems = launch_problems(algorithm, { "rollout_batch_size": config.batch.groups_per_round, "rollout_max_response_len": config.batch.rollout_max_response_len, - "context_parallel_size": 1, # ports emits --context-parallel-size 1 + "context_parallel_size": int(getattr(config.parallel, "context_parallel", 1) or 1), "multi_lora": any(t.split("=", 1)[0] == "--multi-lora" for t in extra_argv), # rl-algo-critic-family 2.3: critic GPU counts / deploy-component / indep-dp "extra_argv": tuple(extra_argv), @@ -917,7 +917,7 @@ def translate_run_config( "--sglang-mem-fraction-static", str(serving.mem_fraction_static), "--tensor-model-parallel-size", str(parallel.tensor_parallel), "--pipeline-model-parallel-size", str(parallel.pipeline_parallel), - "--context-parallel-size", "1", + "--context-parallel-size", str(int(getattr(parallel, "context_parallel", 1) or 1)), "--expert-model-parallel-size", str(parallel.expert_parallel), "--expert-tensor-parallel-size", "1", # data diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 1036225ea..08100c3e5 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -436,6 +436,7 @@ island_flag("--tensor-parallel", type=int, default=1), island_flag("--pipeline-parallel", type=int, default=1), island_flag("--expert-parallel", type=int, default=None), + island_flag("--context-parallel", type=int, default=1), island_flag("--rollout-num-gpus", type=int, default=None), # fleet-dashboard 2.1/2.2 (ports only): rl_heartbeat / rl_resource_sample # periods in seconds; default None = 30 / 60 with --rl-observe-timeline, diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 7858496ed..a02b6522b 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -185,6 +185,9 @@ class ParallelLayout: # --rl-allow-cross-node-engine-tp); default False = TP stays inside a node. allow_cross_node_tp: bool = False allow_cross_node_engine_tp: bool = False + # S27 G-AC: Megatron context parallelism of the trainer (Miles splits each + # packed micro batch of max_tokens_per_gpu * CP tokens across the CP ranks). + context_parallel: int = 1 @property def colocated(self) -> bool: @@ -707,9 +710,12 @@ def resolve_rl_run_config( actor_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node tensor_parallel = getattr(args, "tensor_parallel", 1) pipeline_parallel = getattr(args, "pipeline_parallel", 1) - model_parallel = tensor_parallel * pipeline_parallel + context_parallel = int(getattr(args, "context_parallel", 1) or 1) + model_parallel = tensor_parallel * pipeline_parallel * context_parallel rules = _rules() rules.check_trainer_parallel(actor_gpus, tensor_parallel, pipeline_parallel) + if context_parallel < 1 or actor_gpus % model_parallel: + raise ValueError("Miles actor world must be divisible by TP*PP*CP") is_moe = getattr(provider, "num_moe_experts", None) is not None expert_parallel = getattr(args, "expert_parallel", None) or ( actor_gpus if is_moe else 1 @@ -1025,6 +1031,7 @@ def resolve_rl_run_config( bundle_map=_island_bundle_map(args, dedicated_rollout_gpus), allow_cross_node_tp=bool(getattr(args, "rl_allow_cross_node_tp", False)), allow_cross_node_engine_tp=bool(getattr(args, "rl_allow_cross_node_engine_tp", False)), + context_parallel=context_parallel, ), trainable=TrainableConfig( parameter_mode=parameter_mode, @@ -1142,10 +1149,12 @@ def resolve_rl_run_config( def _train_memory(args) -> TrainMemoryConfig: max_tokens = getattr(args, "rl_max_tokens_per_gpu", None) - if max_tokens is not None and int(max_tokens) < int(args.seq_len): + cp = int(getattr(args, "context_parallel", 1) or 1) + if max_tokens is not None and int(max_tokens) * cp < int(args.seq_len): # a sample longer than the budget would not fit in any micro batch + # (Miles packs max_tokens_per_gpu * CP tokens per micro batch: dp_schedule.py) raise ValueError( - f"--rl-max-tokens-per-gpu {max_tokens} is below --seq-len {args.seq_len}" + f"--rl-max-tokens-per-gpu {max_tokens} x CP {cp} is below --seq-len {args.seq_len}" ) offload = float(getattr(args, "rl_optimizer_offload_fraction", 0.0) or 0.0) if not 0.0 <= offload <= 1.0: From fdd468ab66953f341dda4398b11a2fa53033b130 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:38:51 +0000 Subject: [PATCH 176/233] tests: fixed-roster quorum timeout keeps pushes and attempt on the real syncer Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_full_parameter_async.py | 48 +++++++++++++++++++++++++-- 1 file changed, 45 insertions(+), 3 deletions(-) diff --git a/tests/test_rl_full_parameter_async.py b/tests/test_rl_full_parameter_async.py index 7448f2bb2..3a9a7f244 100644 --- a/tests/test_rl_full_parameter_async.py +++ b/tests/test_rl_full_parameter_async.py @@ -617,7 +617,7 @@ def _nesterov_reference(rounds, *, lr: float, mu: float) -> dict[int, np.ndarray return params -def _drive(binary, tmp_path, extra, *, total_steps, push, late=None): +def _drive(binary, tmp_path, extra, *, total_steps, push, late=None, syncer_args=()): """Run the launcher's async syncer with two learners; return (finals, log). ``push(learner, permit)`` is called once per PULL a learner sees and returns @@ -626,8 +626,8 @@ def _drive(binary, tmp_path, extra, *, total_steps, push, late=None): versions and returns ``(learner, permit, answer)`` pushes to send now.""" log_file = tmp_path / "syncer.log" - process, port, profile_hash = _start_async_syncer(binary, tmp_path, *extra, - total_steps=total_steps, log_file=log_file) + process, port, profile_hash = _start_async_syncer(binary, tmp_path, *extra, total_steps=total_steps, + syncer_args=tuple(syncer_args), log_file=log_file) clients = [_client(learner, port, profile_hash) for learner in (0, 1)] seen: list[set] = [set(), set()] versions = [0] * P @@ -1117,3 +1117,45 @@ def late(_versions): _assert_matches(finals, _nesterov_reference(rounds, lr=1.0, mu=0.0)) assert "push rejected" not in log and "OutOfRound" not in log assert "new frozen-membership attempt" not in log + + +@needs_cargo +@pytest.mark.usefixtures("_hmac_key") +def test_real_fixed_roster_quorum_timeout_keeps_the_pushes_and_the_attempt(syncer_binary, tmp_path, monkeypatch): + """With a fixed roster the quorum timeout only re-sends the PULL to the + missing learner: the pushes that arrived stay, the attempt stays 1, and + the late learner's push (same attempt) completes the round.""" + + from yeto.launch import syncer as launch_syncer + + monkeypatch.setattr(launch_syncer, "SYNCER_DEFAULT_QUORUM_TIMEOUT_S", 1) # profile hash of the 1 s syncer + rounds: dict[int, dict] = {} + held: list = [] + started = time.monotonic() + + def push(learner, permit): + assert permit.round_attempt == 1 + delta = _delta(learner, permit.global_step) + rounds.setdefault(permit.global_step, {})[learner] = (1.0, delta) + answer = (0, 1, 7, delta) + if learner == 1: + held.append((learner, permit, answer)) # late: past the 1 s quorum timeout + return None + return answer + + def late(_versions): + if held and time.monotonic() - started > 3.5: + out = list(held) + held.clear() + return out + return [] + + extra = ("--rl-sync-fragment-pipeline", str(P), "--rl-sync-quorum", "2", "--rl-sync-grace-ms", "0", + "--rl-sync-max-base-lag", "0", "--rl-sync-learner-weight", "equal", + "--rl-sync-outer-lr", "1.0", "--rl-sync-outer-momentum", "0") + finals, log = _drive(syncer_binary, tmp_path, extra, total_steps=1, push=push, late=late, + syncer_args=("--quorum-timeout-s", "1")) + assert "fixed roster incomplete" in log + assert "new frozen-membership attempt" not in log and "push rejected" not in log + assert all(sorted(contributions) == [0, 1] for contributions in rounds.values()) + _assert_matches(finals, _nesterov_reference(rounds, lr=1.0, mu=0.0)) From aa04c6a23b5d5d185dba9d6e64d2f5ebdd761835 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:39:21 +0000 Subject: [PATCH 177/233] multinode: the expert group shares TP x CP x DP ranks (CP-aware trainer replica size) Co-Authored-By: Claude Opus 5.5 --- yeto/rl/engine/multinode.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/yeto/rl/engine/multinode.py b/yeto/rl/engine/multinode.py index b891c174d..72c0c790b 100644 --- a/yeto/rl/engine/multinode.py +++ b/yeto/rl/engine/multinode.py @@ -46,9 +46,11 @@ def trainer_replica_gpus(dims: Mapping[str, int]) -> int: ranks (expert tensor parallel ``etp`` defaults to 1), so ``ep*etp <= tp`` needs no extra GPU and a larger EP needs ``ep*etp/tp`` data-parallel replicas of the dense ``tp*cp*pp`` group (the old ``tp*cp*ep*pp`` - over-estimated by a factor of up to tp).""" + over-estimated by a factor of up to tp). Megatron-Core carves the expert + group out of TP x CP x DP (ETP x EP x EDP = TP x CP x DP), so with context + parallelism ``ep*etp <= tp*cp`` needs no extra GPU either (S27 G-AC).""" tp, cp, pp, ep, etp = (max(1, int(dims.get(dim, 1) or 1)) for dim in ("tp", "cp", "pp", "ep", "etp")) - return tp * cp * pp * max(1, -(-ep * etp // tp)) + return tp * cp * pp * max(1, -(-ep * etp // (tp * cp))) class TopologyError(ValueError): From a62580ed6e014f843dfcf8aa3b761dbbb11d2091 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:42:27 +0000 Subject: [PATCH 178/233] miles config: map parallel.context_parallel (S27 G-AC) Co-Authored-By: Claude Opus 5.5 --- yeto/rl/adapters/miles/config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 79fc213d9..2fe38457b 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -351,6 +351,7 @@ def _check_sglang_tp(value, config): "parallel.tensor_parallel": _ok, "parallel.pipeline_parallel": _ok, "parallel.expert_parallel": _ok, + "parallel.context_parallel": _ok, # S27 G-AC: --context-parallel-size (translate_run_config) # Derived by Megatron from world size / TP / PP; consumed, not passed. "parallel.data_parallel": _ok, "parallel.rollout_num_gpus_per_engine": _ok, From 3c1ad34c4906b7f9e0761ec3df46044f89d4a996 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:36:32 +0000 Subject: [PATCH 179/233] rl: deferred capture, batched partial-quorum applies, faster pushes - Deferred capture (asynchronous rounds, no overlap): a local round's capture asks Miles for descriptors only (export defer_payload=True); the pairing materializes exactly the fragments the next submit_ready pushes (RolePairing.prepare_submissions -> stream.materialize_submission) while the trainer still holds them. ready_submission_candidates skips a fragment until it is stored. Needs the Miles deferred export and the skip-resident apply; YETO_FULL_PARAMETER_DEFER_CAPTURE=0 keeps the full capture. Cut identities are unchanged (same descriptors). - Batched applies: _commit_partial_quorum plans its decisions in the same order and applies the broadcasts of distinct fragments in one trainer apply per role (apply_staged_broadcasts); a repeated fragment flushes the batch first. YETO_FULL_PARAMETER_BATCH_APPLY=0 keeps one apply each. - The optimizer proofs around an apply skip the whole-model model/master comparison when Miles offers check_model=False (the commit ends with it). - protocol._block_aligned_chunks yields zero-copy views: slicing bytes copied the rest of a part at every tensor-segment cut (about 14 s/GiB with 10 MiB tensors); the encoded bytes are identical. - --observe-training-entropy and --mamba-training-ssm-states-dtype reach the Miles argv. Tests: the real Miles role stream with every feature on and off gives bit-identical trainer values, syncer values and pushed deltas over three asynchronous rounds (6 -> 3 trainer applies, 21 -> 6 whole-model checks); flag plumbing tests. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_full_parameter_async.py | 15 + tests/test_rl_miles_adapter_config.py | 18 ++ tests/test_s27_sync_fast_miles_role_stream.py | 289 ++++++++++++++++++ yeto/launch/rl_island_task.py | 4 + yeto/protocol.py | 17 +- yeto/rl/adapters/miles/config.py | 6 + .../miles/full_parameter/role_stream.py | 89 +++++- .../miles/models/chunked_full_parameter.py | 124 +++++++- yeto/rl/engine/full_parameter_sync.py | 107 ++++++- yeto/rl/engine/run_config.py | 6 + yeto/rl/engine/run_flags.py | 14 + 11 files changed, 674 insertions(+), 15 deletions(-) create mode 100644 tests/test_s27_sync_fast_miles_role_stream.py diff --git a/tests/test_rl_full_parameter_async.py b/tests/test_rl_full_parameter_async.py index 3a9a7f244..6189339a2 100644 --- a/tests/test_rl_full_parameter_async.py +++ b/tests/test_rl_full_parameter_async.py @@ -1159,3 +1159,18 @@ def late(_versions): assert "new frozen-membership attempt" not in log and "push rejected" not in log assert all(sorted(contributions) == [0, 1] for contributions in rounds.values()) _assert_matches(finals, _nesterov_reference(rounds, lr=1.0, mu=0.0)) + + +def test_observe_training_entropy_reaches_every_island(): + _, plain = _plan(ASYNC) + _, observed = _plan(ASYNC + ("--observe-training-entropy",)) + for island in plain["island_requests"]: + assert "--observe-training-entropy" not in island["learner_command"] + for island in observed["island_requests"]: + assert " --observe-training-entropy" in island["learner_command"] + + +def test_mamba_ssm_state_dtype_reaches_every_island(): + _, plan = _plan(ASYNC + ("--mamba-training-ssm-states-dtype", "fp32")) + for island in plan["island_requests"]: + assert _flag(island["learner_command"], "--mamba-training-ssm-states-dtype") == "fp32" diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index 6de68515c..2775c290e 100644 --- a/tests/test_rl_miles_adapter_config.py +++ b/tests/test_rl_miles_adapter_config.py @@ -839,3 +839,21 @@ def test_generic_recipe_trains_without_dropout(): cfg = sub(cfg, "model_recipe", name=rc.RECIPE_GENERIC) joined = " ".join(mc.translate_run_config(cfg, AlgorithmSpec()).argv) assert "--attention-dropout 0.0 --hidden-dropout 0.0" in joined + + +def test_observe_training_entropy_reaches_the_miles_argv(): + plain = list(mc.translate_run_config(make_config(), AlgorithmSpec()).argv) + assert "--observe-training-entropy" not in plain + cfg = sub(make_config(), "algorithm", observe_training_entropy=True) + argv = list(mc.translate_run_config(cfg, AlgorithmSpec()).argv) + assert argv.count("--observe-training-entropy") == 1 + assert [value for value in argv if value != "--observe-training-entropy"] == plain + + +def test_mamba_training_ssm_state_dtype_reaches_the_miles_argv(): + plain = list(mc.translate_run_config(make_config(), AlgorithmSpec()).argv) + assert "--mamba-training-ssm-states-dtype" not in plain + cfg = sub(make_config(), "algorithm", mamba_training_ssm_states_dtype="fp32") + argv = list(mc.translate_run_config(cfg, AlgorithmSpec()).argv) + at = argv.index("--mamba-training-ssm-states-dtype") + assert argv[at + 1] == "fp32" and argv[:at] + argv[at + 2:] == plain diff --git a/tests/test_s27_sync_fast_miles_role_stream.py b/tests/test_s27_sync_fast_miles_role_stream.py new file mode 100644 index 000000000..4fcdf684a --- /dev/null +++ b/tests/test_s27_sync_fast_miles_role_stream.py @@ -0,0 +1,289 @@ +"""S27 G-AA: deferred capture and batched applies on the real Miles role stream. + +Asynchronous rounds (partial quorum), several local rounds per island: + +* deferred capture: each local round's capture hashes every fragment but + stores none in Ray; the pairing materializes only the fragments it pushes; +* batched apply: the broadcasts one commit drains are applied in one trainer + apply (one pair of optimizer proofs, one whole-model commit); +* optimizer proofs without the whole-model model/master comparison. + +The trainer fake models the Miles contract (``full_parameter_state``): a +deferred row has no bytes, a skip-resident apply skips a fragment whose +payload hash the trainer recorded at its last export or commit, a local step +drops that record, and materialization refuses a fragment whose bytes +changed. Miles' own tests check that side against the real state code. + +The run with every feature on must give bit-identical trainer values, syncer +values and pushed deltas to the run with every feature off. + +Needs Miles (``importorskip``). +""" + +from __future__ import annotations + +# ruff: noqa: I001 -- the Miles test helpers install a bridge stub before import. + +import asyncio +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("miles") + +from test_e3_overlap_miles_role_stream import ( + OTHER_WEIGHT, + P, + SIZES, + _adapter, + _build_cut, + _Client, + _Driver, + _Syncer, + _Trainer, +) +from test_rl_miles_chunked_full_parameter import FakeRay, _evidence, component + +_state = pytest.importorskip("miles.backends.megatron_utils.full_parameter_state") +if not hasattr(_state, "DEFERRED_FULL_PARAMETER_CHUNK"): + pytest.skip("needs a Miles with the deferred export (yeto/s27-sync-fast)", allow_module_level=True) +DEFERRED_FULL_PARAMETER_CHUNK = _state.DEFERRED_FULL_PARAMETER_CHUNK +is_deferred_full_parameter_chunk = _state.is_deferred_full_parameter_chunk +from yeto.protocol import PullRequest +from yeto.rl.adapters.miles.full_parameter import role_stream as role_stream_module +from yeto.rl.adapters.miles.full_parameter.role_stream import ( + TERMINAL_RETIRE_ENV, + MilesFullParameterRoleStream, + RoleStreamConfig, +) +from yeto.rl.adapters.miles.models import chunked_full_parameter as adapter_module +from yeto.rl.adapters.miles.models.chunked_full_parameter import ( + DEFER_CAPTURE_ENV, + SKIP_RESIDENT_ENV, + MilesChunkedFullParameterAdapter, +) +from yeto.rl.engine.full_parameter_sync import ( + BATCH_APPLY_ENV, + FullParameterSync, + FullParameterSyncProfile, + RolePairing, +) + +ROUNDS = 3 # local rounds per island; every fragment merges ROUNDS times +DEPTH = 2 + + +class _AsyncSyncer(_Syncer): + """Asynchronous fragment rounds with this island and one virtual peer: + a round merges when this island pushed it and is broadcast at once; the + round ``DEPTH`` steps later opens then; the last step arrives as FINAL.""" + + def commit_ready(self) -> None: + total = self.sweeps * P + for step in sorted(self.received): + delta = self.received.pop(step) + f, k = (step - 1) % P, (step - 1) // P + self.global_[f] = self.global_[f] + OTHER_WEIGHT * (delta + self.peer(k)[f]) + self.versions[f] = step + self.committed.append(step) + if step < total: + self.client.broadcast(f, step, self.global_[f]) + if step + DEPTH <= total: + self.client.pulls.append(PullRequest((step + DEPTH - 1) % P, step + DEPTH, self.attempt)) + if len(self.committed) == total and not self.client.finalizing.is_set(): + self.client.publish_final() + + +class _AsyncClient(_Client): + def pull_sweep(self, sweep, attempt) -> None: + if sweep == 0: + self.pulls.extend(PullRequest(step - 1, step, attempt) for step in range(1, DEPTH + 1)) + + +def _hashes(cut): + return {d.fragment_id: d.payload_hash for shard in cut.shards for d in shard.fragments} + + +class _MilesContractTrainer(_Trainer): + """``_Trainer`` plus the Miles deferred-export / skip-resident contract.""" + + def __init__(self, *args, fast: bool, **kwargs): + super().__init__(*args, **kwargs) + self.fast = fast + self.resident: dict[int, str] | None = None + self.stored_bytes = 0 + self.applies = 0 + self.read_fragments = 0 + self.model_checks = 0 + self.deferred_exports = 0 + + def _stored(self, cut): + self.stored_bytes += sum(4 * d.numel for shard in cut.shards for d in shard.fragments) + return cut + + async def record_full_parameter_local_step(self, **kwargs): + self.resident = None # the masters changed + self.model_checks += 1 + return await super().record_full_parameter_local_step(**kwargs) + + async def export_full_parameter_chunked_cut(self, policy_version, local_step_generation=0, *, + max_chunk_bytes, defer_payload=False): + assert defer_payload == self.fast or local_step_generation == 0 + if not defer_payload: + cut = self._stored(await super().export_full_parameter_chunked_cut( + policy_version, local_step_generation, max_chunk_bytes=max_chunk_bytes)) + else: + self.deferred_exports += 1 + cut = _build_cut(self.adapter, self.values, policy_version=policy_version, + generation=local_step_generation, ray=FakeRay()) # hashes only + for shard in cut.shards: + for refs in shard.chunk_refs: + refs[:] = [DEFERRED_FULL_PARAMETER_CHUNK] * len(refs) + self.resident = _hashes(cut) + return cut + + async def materialize_full_parameter_fragments(self, cut, fragment_ids): + assert self.fast + now = _hashes(_build_cut(self.adapter, self.values, policy_version=0, generation=0, ray=FakeRay())) + wanted = _hashes(cut) + stored = {} + fresh = _build_cut(self.adapter, self.values, policy_version=cut.policy_version, + generation=cut.local_step_generation, ray=self.ray) + for fragment_id in fragment_ids: + assert self.resident is not None and self.resident[fragment_id] == wanted[fragment_id], \ + "materialized a fragment the masters no longer hold" + assert now[fragment_id] == wanted[fragment_id] + ((refs, numel),) = [(r, d.numel) for shard in fresh.shards + for d, r in zip(shard.fragments, shard.chunk_refs, strict=True) + if d.fragment_id == fragment_id] + stored[fragment_id] = list(refs) + self.stored_bytes += 4 * numel + return stored + + async def apply_full_parameter_chunked_cut(self, cut, *, commit_token=None, max_commit_attempts=3, + skip_resident_fragments=False): + assert commit_token + assert skip_resident_fragments == self.fast + self.applies += 1 + for shard in cut.shards: + for descriptor, refs in zip(shard.fragments, shard.chunk_refs, strict=True): + if skip_resident_fragments and self.resident is not None and \ + self.resident.get(descriptor.fragment_id) == descriptor.payload_hash: + continue # the masters hold these bytes + assert not any(is_deferred_full_parameter_chunk(ref) for ref in refs), "read a deferred row" + assert all(ref is not None for ref in refs), "applied a released chunk" + self.read_fragments += 1 + self.values[descriptor.fragment_id] = torch.cat([self.ray.get(ref) for ref in refs]).clone() + self.installed = cut.policy_version + self.generation = 0 + self.resident = _hashes(cut) + self.model_checks += 1 + return self.adapter.expected_parameter_tensor_count + + async def full_parameter_optimizer_states(self, *, check_model=True): + assert check_model is False and self.fast + return await super().full_parameter_optimizer_states() + + +class _OriginalContractTrainer(_MilesContractTrainer): + """The trainer before S27 G-AA: every optimizer proof compares the model.""" + + async def full_parameter_optimizer_states(self): + self.model_checks += 1 + return await _Trainer.full_parameter_optimizer_states(self) + + +def _config() -> RoleStreamConfig: + adapter = _adapter() + return RoleStreamConfig( + algorithm="grpo", role="actor", component=component(), syncer_addr=("127.0.0.1", 30100), + learner_id=0, learner_generation=0, learner_generations=(0, 0), + total_fragment_steps=ROUNDS * P, expected_fragments=P, + expected_layout_hash=adapter.layout.layout_hash, local_horizon=1, optimizer_steps_per_round=1, + training_contract_hash="2" * 64, syncer_profile_hash="3" * 64, pipeline_depth=DEPTH, + algorithm_spec_sha256="a1" * 32, wait_timeout=5.0, poll_seconds=0.001, + max_fragment_bytes=64, max_chunk_bytes=8, partial_quorum=True) + + +def _run(monkeypatch, *, fast: bool): + switch = "1" if fast else "0" + for env in (DEFER_CAPTURE_ENV, BATCH_APPLY_ENV, SKIP_RESIDENT_ENV, TERMINAL_RETIRE_ENV): + monkeypatch.setenv(env, switch) + monkeypatch.setattr(adapter_module, "_local_transport_features", + lambda: frozenset({"skip_resident", "deferred_export", "optimizer_state_check_model"})) + generator = torch.Generator().manual_seed(29) + + def rand(scale=1.0): + return {f: torch.randn(n, generator=generator) * scale for f, n in enumerate(SIZES)} + + initial = rand() + steps = {k: rand(0.1) for k in range(ROUNDS)} + peer = {k: rand(0.1) for k in range(ROUNDS)} + ray = FakeRay() + adapter = _adapter() + trainer = (_MilesContractTrainer if fast else _OriginalContractTrainer)( + adapter, ray, initial, lambda k: steps[k], fast=fast) + syncer = _AsyncSyncer(initial, ROUNDS, lambda k: peer[k]) + client = _AsyncClient(syncer) + + async def initialize(cls, group, *, policy_version, **_kwargs): + cut = await group.export_full_parameter_chunked_cut(policy_version, 0, + max_chunk_bytes=adapter.max_chunk_bytes) + return adapter, adapter.validate_cut(cut, policy_version, 0) + + monkeypatch.setattr(MilesChunkedFullParameterAdapter, "initialize", classmethod(initialize)) + monkeypatch.setattr(role_stream_module, "SyncerClient", lambda *_args, **_kwargs: client) + stream = MilesFullParameterRoleStream(_config(), ray_module=ray) + profile = FullParameterSyncProfile.from_configs({"actor": stream.config}) + session = FullParameterSync( + RolePairing({"actor": stream}, profile), groups={"actor": trainer}, run=asyncio.run, + trajectories=lambda rollout_id: _evidence(rollout_id=rollout_id, + behavior_policy_hash=stream.current.policy_hash)) + driver = _Driver() + session.start(driver) + stops = [] + for rollout_id in range(ROUNDS + 2): + boundary = session.boundary(driver, rollout_id=rollout_id, stats=None) + stops.append(boundary.stop) + if boundary.stop: + session.published(driver, rollout_id=rollout_id, policy_hash=boundary.state.policy_tensor_hash()) + break + session.finish(driver) + return SimpleNamespace(trainer=trainer, syncer=syncer, stream=stream, stops=stops, + pairing=session.pairing, driver=driver) + + +def _bits(values): + return {f: value.view(torch.int32).tolist() for f, value in values.items()} + + +def test_deferred_capture_and_batched_applies_keep_every_value_bit_identical(monkeypatch): + slow = _run(monkeypatch, fast=False) + fast = _run(monkeypatch, fast=True) + + for run in (slow, fast): + assert run.stops == [False] * (ROUNDS - 1) + [True] + assert run.syncer.committed == list(range(1, ROUNDS * P + 1)) + assert run.stream.final_applied and run.stream.final_acknowledged + assert _bits(fast.trainer.values) == _bits(slow.trainer.values) + assert _bits(fast.syncer.global_) == _bits(slow.syncer.global_) + assert [(s, a, d.view(torch.int32).tolist()) for s, a, d in fast.syncer.pushed] == [ + (s, a, d.view(torch.int32).tolist()) for s, a, d in slow.syncer.pushed] + # Every trainer value went through the same local rounds. + assert len(fast.trainer.trained_from) == len(slow.trainer.trained_from) == ROUNDS + for got, want in zip(fast.trainer.trained_from, slow.trainer.trained_from, strict=True): + assert _bits(got) == _bits(want) + + # Deferred capture: the round cuts store only what is pushed (once each). + assert fast.trainer.deferred_exports == ROUNDS + model_bytes = 4 * sum(SIZES) + assert slow.trainer.stored_bytes == (1 + ROUNDS) * model_bytes + assert fast.trainer.stored_bytes == (1 + ROUNDS) * model_bytes # INIT + every pushed fragment once + print(f"applies {slow.trainer.applies}->{fast.trainer.applies} " + f"fragment reads {slow.trainer.read_fragments}->{fast.trainer.read_fragments} " + f"whole-model checks {slow.trainer.model_checks}->{fast.trainer.model_checks}") + # Batched applies: one apply per drained commit instead of one per fragment. + assert fast.trainer.applies < slow.trainer.applies + assert fast.trainer.read_fragments < slow.trainer.read_fragments + assert fast.trainer.model_checks < slow.trainer.model_checks diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 1f11bf6a3..984c8f298 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -290,6 +290,10 @@ def make_miles_island_task( flags += f" {flag} {shlex.quote(str(value))}" if args.use_rollout_routing_replay: flags += " --use-rollout-routing-replay" + if getattr(args, "observe_training_entropy", False): + flags += " --observe-training-entropy" + if getattr(args, "mamba_training_ssm_states_dtype", None): + flags += f" --mamba-training-ssm-states-dtype {args.mamba_training_ssm_states_dtype}" if getattr(args, "apply_chat_template_kwargs", None): chat_kwargs = json.dumps( args.apply_chat_template_kwargs, sort_keys=True, separators=(",", ":") diff --git a/yeto/protocol.py b/yeto/protocol.py index 1002f11e6..4447c2a75 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -281,20 +281,25 @@ def _block_aligned_chunks( left = next(ahead, 0) if ahead is not None else float("inf") carry = b"" for part in parts: - data = carry + bytes(part) - while data: + # Zero-copy: chunks are read-only views of the part (a carried + # partial block is joined once). Slicing bytes here copied the rest + # of the part at every segment cut, quadratic in the segments per + # part (S27 G-AA: ~14 s/GiB with 10 MiB tensors). + data = memoryview(carry + bytes(part) if carry else part).cast("B") + start = 0 + while start < len(data): if left == 0: raise ValueError("f32 push delta is longer than its tensor segments") - take = min(len(data) // 4, left) + take = min((len(data) - start) // 4, left) cut = take if take == left else take - take % block if not cut: break - yield data[: 4 * cut] - data = data[4 * cut :] + yield data[start : start + 4 * cut] + start += 4 * cut left -= cut if left == 0 and ahead is not None: left = next(ahead, 0) - carry = data + carry = bytes(data[start:]) if len(carry) % 4: raise ValueError("f32 push delta ends mid-value") if carry: diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index cbc65f521..2a6756150 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -402,6 +402,8 @@ def _check_sglang_tp(value, config): "algorithm.lr_schedule.decay_style": _ok, "algorithm.lr_schedule.decay_iters": _ok, "algorithm.seed": _ok, + "algorithm.observe_training_entropy": _ok, + "algorithm.mamba_training_ssm_states_dtype": _ok, "algorithm.rollout_seed": _ok, # rl-algo-critic-family 2.4: translated with the AlgorithmSpec (critic_load_argv) "algorithm.critic": _ok, @@ -963,6 +965,10 @@ def translate_run_config( # Explicit LR schedule: Miles' default horizon is --num-rollout (global # rounds), which under decoupled is not the island's local step count. values.extend(lr_schedule_argv(config.algorithm.lr_schedule)) + if config.algorithm.observe_training_entropy: + values.append("--observe-training-entropy") + if config.algorithm.mamba_training_ssm_states_dtype is not None: + values.extend(("--mamba-training-ssm-states-dtype", config.algorithm.mamba_training_ssm_states_dtype)) if algorithm.kl_coef is not None: values.extend(("--kl-coef", str(algorithm.kl_coef))) if dynamic_filter is not None: diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index b3d57c1de..35b8b66d1 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -42,6 +42,8 @@ MilesChunkedFullParameterAdapter, ReferencedPolicyCut, StoredAuthoritativeFragment, + defer_capture_supported, + optimizer_proof_without_model_check_supported, ) from yeto.rl.engine.policy_digest import ( PolicyDigest, @@ -441,6 +443,10 @@ def __init__(self, config: RoleStreamConfig, *, ray_module=None) -> None: # Terminal phase: fragments whose last round closed without a trainer # apply (``retire_staged_broadcast``); the FINAL apply covers them. self.retired_fragments: set[int] = set() + # Decided once the trainer group is known (``_defers_capture``, + # ``_proof_kwargs``). + self._defer_capture: bool | None = None + self._proof_without_model_check: bool | None = None self.local_generation = 0 self.local_rounds = 0 self.action_tokens = 0 @@ -950,6 +956,7 @@ async def complete_local_round( self.group, policy_version=self.transport_version, local_step_generation=expected_generation, + defer_payload=self._defers_capture(), ) previous = self.current previous_submission = self.submission_cut @@ -982,6 +989,45 @@ async def complete_local_round( False, ) + def _defers_capture(self) -> bool: + """S27 G-AA deferred capture: a local round's capture hashes every + fragment but stores none in Ray; :meth:`materialize_submission` stores + the fragments a push needs. Only for asynchronous rounds without + overlap (every reader of the cut bytes is then a push of a + materialized fragment or a skip-resident apply that never reads a + deferred row) and only when the trainer's Miles supports it.""" + + if self._defer_capture is None: + self._defer_capture = bool( + self.config.partial_quorum + and not self.config.overlap + and self.group is not None + and defer_capture_supported(self.group) + ) + return self._defer_capture + + async def materialize_submission(self, fragment_ids: Sequence[int]) -> tuple[int, ...]: + """Store the deferred fragments ``fragment_ids`` of the submission cut. + + Called by the pairing before it pushes them, while the trainer still + holds their bytes (no local step since the capture, and an apply + never writes a fragment this island still has to push from this + cut: such a fragment's version would move past the submission's).""" + + if self.adapter is None or self.group is None or self.submission_cut is None: + return () + return await self.adapter.materialize(self.group, self.submission_cut, fragment_ids) + + def _proof_kwargs(self) -> dict[str, bool]: + """The optimizer proofs around an apply skip their whole-model + model/master comparison when Miles supports it: the commit itself + copies every master into the model and ends with that comparison.""" + + if self._proof_without_model_check is None: + self._proof_without_model_check = bool( + self.group is not None and optimizer_proof_without_model_check_supported(self.group)) + return {"check_model": False} if self._proof_without_model_check else {} + async def stage_available_broadcasts(self) -> tuple[tuple[int, int], ...]: """Drain complete broadcasts without mutating the trainable model.""" @@ -1068,6 +1114,32 @@ async def apply_staged_broadcast( self._close_staged_round(fragment_id, version) return (fragment_id,) + async def apply_staged_broadcasts( + self, items: Sequence[tuple[int, int, bool]] + ) -> tuple[int, ...]: + """Apply several staged broadcasts ``(fragment_id, version, catch_up)`` + of distinct fragments in one trainer apply (partial quorum). + + Same checks and bookkeeping as one :meth:`apply_staged_broadcast` per + item; the fragments occupy disjoint master ranges, so one apply of all + of them writes the same bytes as one apply each, and the whole-model + costs of an apply (optimizer proofs, master-to-model copy, checks) are + paid once.""" + + if self.config.overlap: + raise RuntimeError(f"{self.config.role} batches applies only without overlap") + fragment_ids = [fragment_id for fragment_id, _version, _catch_up in items] + if not items or len(set(fragment_ids)) != len(fragment_ids): + raise RuntimeError(f"{self.config.role} batched apply needs distinct fragments") + selected = { + fragment_id: self._staged_for_apply(fragment_id, version, rebase=False, catch_up=catch_up) + for fragment_id, version, catch_up in items + } + await self._apply_authoritative(selected) + for fragment_id, version, _catch_up in items: + self._close_staged_round(fragment_id, version) + return tuple(fragment_ids) + def _staged_for_apply( self, fragment_id: int, version: int, *, rebase: bool, catch_up: bool ) -> StoredAuthoritativeFragment: @@ -1206,7 +1278,7 @@ async def _apply_authoritative( target_policy_version=target_version, authoritative_fragments=tuple(selected.values()), ) - before = (tuple(await self.group.full_parameter_optimizer_states()) + before = (tuple(await self.group.full_parameter_optimizer_states(**self._proof_kwargs())) if optimizer_proof else None) await self.adapter.apply( self.group, @@ -1214,7 +1286,7 @@ async def _apply_authoritative( commit_token=self._apply_token(target_version, selected), ) if before is not None: - after = tuple(await self.group.full_parameter_optimizer_states()) + after = tuple(await self.group.full_parameter_optimizer_states(**self._proof_kwargs())) validate_optimizer_apply_transition( before, after, @@ -1270,8 +1342,13 @@ def _apply_token( ).encode() return f"{self.config.algorithm}-{self.config.role}:{hashlib.sha256(payload).hexdigest()}" - def ready_submission_candidates(self) -> tuple[RoleSubmissionCandidate, ...]: - """Return H-qualified retained permits without releasing any bytes.""" + def ready_submission_candidates( + self, *, include_deferred: bool = False + ) -> tuple[RoleSubmissionCandidate, ...]: + """Return H-qualified retained permits without releasing any bytes. + + A fragment of a deferred capture is a candidate only once it is + materialized (``include_deferred``: also before, to plan which).""" if ( self.client is None @@ -1283,8 +1360,12 @@ def ready_submission_candidates(self) -> tuple[RoleSubmissionCandidate, ...]: return () self.client.check_health() self._queue_pulls() + deferred = (self.adapter.deferred_fragment_ids(self.submission_cut) + if self._defer_capture and not include_deferred else frozenset()) candidates = [] for fragment_id in sorted(self._pending): + if fragment_id in deferred: + continue permit = self._pending[fragment_id] version = self.fragment_versions[fragment_id] expected_base = max( diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index 95a1bbe41..2cd701261 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -471,14 +471,67 @@ async def capture( *, policy_version: int, local_step_generation: int, + defer_payload: bool = False, ) -> ReferencedPolicyCut: + """``defer_payload``: Miles hashes every fragment (same descriptors, + same cut identity) but stores none in Ray; :meth:`materialize` stores + the fragments a push needs while the masters still hold them.""" + cut = await group.export_full_parameter_chunked_cut( policy_version, local_step_generation, max_chunk_bytes=self.max_chunk_bytes, + **({"defer_payload": True} if defer_payload else {}), ) return self.validate_cut(cut, policy_version, local_step_generation) + @staticmethod + def deferred_fragment_ids(cut: ReferencedPolicyCut) -> frozenset[int]: + """Fragments of ``cut`` whose chunk refs are deferred-export markers.""" + + try: + from miles.backends.megatron_utils.full_parameter_state import ( + is_deferred_full_parameter_chunk, + ) + except ImportError: # an old Miles never defers + return frozenset() + return frozenset( + descriptor.fragment_id + for shard in cut.transport_cut.shards + for descriptor, references in zip(shard.fragments, shard.chunk_refs, strict=True) + if any(is_deferred_full_parameter_chunk(reference) for reference in references) + ) + + async def materialize( + self, + group: MilesChunkedFullParameterGroup, + cut: ReferencedPolicyCut, + fragment_ids: Sequence[int], + ) -> tuple[int, ...]: + """Store the deferred fragments ``fragment_ids`` of ``cut`` in Ray. + + The owner ranks export them again from their masters, refuse when the + masters changed since the cut, and check the bytes against the cut's + descriptors; the refs then replace the markers in ``cut``'s rows (the + cut identity does not change). Returns the fragments stored.""" + + deferred = self.deferred_fragment_ids(cut) + wanted = tuple(sorted({fragment_id for fragment_id in fragment_ids if fragment_id in deferred})) + if not wanted: + return () + stored = await group.materialize_full_parameter_fragments(cut.transport_cut, wanted) + if set(stored) != set(wanted): + raise RuntimeError("Miles materialized the wrong full-parameter fragments") + for shard in cut.transport_cut.shards: + for descriptor, references in zip(shard.fragments, shard.chunk_refs, strict=True): + refs = stored.get(descriptor.fragment_id) + if refs is None: + continue + if len(refs) != len(references): + raise RuntimeError("Miles materialized a fragment with the wrong chunk count") + references[:] = refs + return wanted + def validate_cut( self, cut: Any, @@ -1015,6 +1068,10 @@ def _fragment_version(self, policy_version: int, fragment_id: int) -> int: def _takes_skip_resident(function: Any) -> bool | None: + return _takes_parameter(function, "skip_resident_fragments") + + +def _takes_parameter(function: Any, name: str) -> bool | None: """True/False when ``function``'s signature names or lacks the parameter; None when it only forwards ``*args, **kwargs``.""" @@ -1024,13 +1081,78 @@ def _takes_skip_resident(function: Any) -> bool | None: parameters = inspect.signature(function).parameters except (TypeError, ValueError): return False - if "skip_resident_fragments" in parameters: + if name in parameters: return True if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()): return None return False +def _local_transport_features() -> frozenset[str]: + try: + from miles.ray import full_parameter_transport + except ImportError: + return frozenset() + return frozenset(getattr(full_parameter_transport, "FULL_PARAMETER_TRANSPORT_FEATURES", ())) + + +def transport_supports(group: Any, method: str, parameter: str, feature: str) -> bool: + """Whether ``group.`` takes ``parameter`` (Miles transport ``feature``). + + A signature that names or lacks the parameter decides; a forwarding + handle (``*args, **kwargs``, Miles' ``RayWorkerHandle``) is decided by the + Miles in this process, which is the trainer's Miles (same image).""" + + function = getattr(group, method, None) + if function is None: + return False + takes = _takes_parameter(function, parameter) + if takes is not None: + return takes + return feature in _local_transport_features() + + +DEFER_CAPTURE_ENV = "YETO_FULL_PARAMETER_DEFER_CAPTURE" + + +def defer_capture_supported(group: Any) -> bool: + """Deferred capture needs the deferred export, on-demand materialization + and the skip-resident apply (a deferred row is only ever skipped). + ``YETO_FULL_PARAMETER_DEFER_CAPTURE=0`` keeps the full capture.""" + + import os + + if os.environ.get(DEFER_CAPTURE_ENV, "1") == "0": + return _report_feature("deferred capture", False, f"{DEFER_CAPTURE_ENV}=0") + enabled = ( + skip_resident_supported(group) + and transport_supports(group, "export_full_parameter_chunked_cut", "defer_payload", "deferred_export") + and getattr(group, "materialize_full_parameter_fragments", None) is not None + and ("deferred_export" in _local_transport_features() + or _takes_parameter(group.materialize_full_parameter_fragments, "fragment_ids") is True) + ) + return _report_feature("deferred capture", enabled, "group and local Miles transport") + + +def optimizer_proof_without_model_check_supported(group: Any) -> bool: + return transport_supports(group, "full_parameter_optimizer_states", "check_model", + "optimizer_state_check_model") + + +_FEATURES_REPORTED: set[str] = set() + + +def _report_feature(name: str, enabled: bool, reason: str) -> bool: + key = f"{name}:{enabled}:{reason}" + if key not in _FEATURES_REPORTED: + _FEATURES_REPORTED.add(key) + import sys + + print(f"[yeto-full-parameter] {name}: {'enabled' if enabled else 'disabled'} ({reason})", + file=sys.stderr, flush=True) + return enabled + + def skip_resident_supported(group: Any) -> bool: """Whether ``group``'s apply takes ``skip_resident_fragments`` (and it is not disabled). diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index f88f83be7..ba139fcb8 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -58,6 +58,7 @@ from __future__ import annotations import asyncio +import os import time from collections.abc import Mapping, Sequence from dataclasses import dataclass, field, replace @@ -368,6 +369,36 @@ def submit_ready(self) -> dict[str, RoleProgress]: raise return {role: RoleProgress(submissions=tuple(values)) for role, values in produced.items()} + async def prepare_submissions(self) -> None: + """Materialize the deferred fragments the next :meth:`submit_ready` + would push (S27 G-AA deferred capture). + + The plan repeats ``submit_ready``'s choice (shared, not closed, within + the free pipeline capacity) over candidates that include deferred + fragments. A PULL that arrives after this plan stays pending until a + later call, exactly as if it had arrived after ``submit_ready``.""" + + if self.failure is not None or not all( + callable(getattr(stream, "materialize_submission", None)) for stream in self.streams.values() + ): + return + try: + by_role = { + role: _candidate_map(stream.ready_submission_candidates(include_deferred=True), role=role) + for role, stream in self.streams.items() + } + capacity = self.profile.effective_pipeline_depth - len(self.inflight) + shared = sorted(set.intersection(*(set(keys) for keys in by_role.values())) - set(self.inflight)) + if self.profile.partial_quorum: + shared = [key for key in shared if not self._round_closed(key[1], key[0])] + fragment_ids = [key[1] for key in shared[: max(0, capacity)]] + if fragment_ids: + await asyncio.gather(*(stream.materialize_submission(fragment_ids) + for stream in self.streams.values())) + except BaseException as error: + self.fail_stop(error) + raise + def _note_final_rounds(self, by_role: Mapping[str, Mapping[LogicalKey, Any]]) -> None: """Set ``terminal_seen`` once every fragment's last round is ready. @@ -538,23 +569,63 @@ async def _commit_partial_quorum(self, applied: dict[str, list[int]]) -> None: *(set(stream.staged_fragment_versions) for stream in self.streams.values()) ) monotonic = all(getattr(stream, "monotonic_broadcasts", False) for stream in self.streams.values()) + # One trainer apply per batch of distinct fragments (S27 G-AA). The + # order of decisions is unchanged: a fragment seen twice first flushes + # the batch, so its second version is judged against the applied one. + batch: list[_PlannedApply] | None = [] if self._batches_applies() else None for fragment_id, version in sorted(staged, key=lambda item: (item[1], item[0])): if version == self.profile.total_fragment_steps: continue # the terminal cut arrives as FINAL + if batch and any(planned.fragment_id == fragment_id for planned in batch): + await self._flush_applies(batch, applied) base = max(0, version - fragments) held = {stream.fragment_versions[fragment_id] for stream in self.streams.values()} if len(held) != 1: continue (current,) = held if current == base: - await self._overwrite_or_apply_pair(fragment_id, version, applied, catch_up=False) + await self._overwrite_or_apply_pair(fragment_id, version, applied, catch_up=False, batch=batch) elif current < base and monotonic and not any( f == fragment_id and v < version for f, v in staged ): - await self._overwrite_or_apply_pair(fragment_id, version, applied, catch_up=True) + await self._overwrite_or_apply_pair(fragment_id, version, applied, catch_up=True, batch=batch) self.caught_up_fragments += 1 self.caught_up_rounds += fragment_rounds(fragment_id, base, fragments) - ( fragment_rounds(fragment_id, current, fragments)) + if batch: + await self._flush_applies(batch, applied) + + def _batches_applies(self) -> bool: + """Every role applies several fragments at once + (``apply_staged_broadcasts``); ``YETO_FULL_PARAMETER_BATCH_APPLY=0`` + keeps one apply per fragment.""" + + return (not self.profile.overlap + and os.environ.get(BATCH_APPLY_ENV, "1") != "0" + and all(callable(getattr(stream, "apply_staged_broadcasts", None)) + for stream in self.streams.values())) + + async def _flush_applies(self, batch: list[_PlannedApply], applied: dict[str, list[int]]) -> None: + """Apply every planned fragment in one call per role, then do each + fragment's bookkeeping exactly as its own apply would.""" + + items = [(planned.fragment_id, planned.version, planned.catch_up) for planned in batch] + results = await asyncio.gather(*(stream.apply_staged_broadcasts(items) for stream in self.streams.values())) + for role, fragment_ids in zip(self.streams, results, strict=True): + applied[role].extend(fragment_ids) + for planned in batch: + for role in self.streams: + if self.streams[role].fragment_versions[planned.fragment_id] != planned.version: + raise RuntimeError("batched apply committed mismatched versions") + if planned.pair is not None: + self._journal(planned.pair) + self.inflight.pop(planned.key) + continue + for key, pair in planned.closed: + self._release(pair) + self.inflight.pop(key) + self.overwritten_rounds += 1 + batch.clear() def _retires(self, version: int) -> bool: """Terminal phase: close ``version`` without a trainer apply. @@ -575,13 +646,22 @@ def _retires(self, version: int) -> bool: ) async def _overwrite_or_apply_pair( - self, fragment_id: int, version: int, applied: dict[str, list[int]], *, catch_up: bool + self, fragment_id: int, version: int, applied: dict[str, list[int]], *, catch_up: bool, + batch: list[_PlannedApply] | None = None, ) -> None: closed = [(key, pair) for key, pair in sorted(self.inflight.items()) if key[1] == fragment_id and key[0] <= version] exact = next((entry for entry in closed if entry[0][0] == version), None) retire = self._retires(version) - if not catch_up and exact is not None and exact[1].complete: + paired = not catch_up and exact is not None and exact[1].complete + if batch is not None and not retire: + # The paired apply and the overwrite differ only in bookkeeping. + batch.append(_PlannedApply(fragment_id, version, catch_up, + key=exact[0] if paired else None, + pair=exact[1] if paired else None, + closed=() if paired else tuple(closed))) + return + if paired: await self._apply_pair(exact[0], exact[1], applied, retire=retire) return if retire: @@ -652,6 +732,7 @@ async def finish_without_more_training(self) -> dict[str, RoleProgress]: progress = {role: RoleProgress() for role in self.streams} deadline = time.monotonic() + self.profile.wait_timeout while not self.finalizing: + await self.prepare_submissions() submitted = self.submit_ready() committed = await self.commit_staged_broadcasts() progress = { @@ -704,6 +785,23 @@ def _candidate_map(candidates: Sequence[Any], *, role: str) -> dict[LogicalKey, return result +BATCH_APPLY_ENV = "YETO_FULL_PARAMETER_BATCH_APPLY" + + +@dataclass(frozen=True) +class _PlannedApply: + """One broadcast of a batched partial-quorum apply and its bookkeeping: + a paired round (``key``/``pair``, journaled) or an overwrite (``closed`` + pairs released).""" + + fragment_id: int + version: int + catch_up: bool + key: LogicalKey | None = None + pair: PairedSubmission | None = None + closed: tuple[tuple[LogicalKey, PairedSubmission], ...] = () + + def _retained_replay(retained: Any, observed: Any | None) -> Any | None: """Consume only a newly observed permit; keep the retained operands.""" @@ -862,6 +960,7 @@ async def _boundary(self, rollout_id: int) -> bool: # Rounds that closed while this island trained are applied first, # so the island never uploads a delta the syncer would refuse. await pairing.commit_staged_broadcasts() + await pairing.prepare_submissions() pairing.submit_ready() submit_ts = unix_now() await pairing.commit_staged_broadcasts() diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 7858496ed..5691f2390 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -284,6 +284,10 @@ class AlgorithmConfig: # state given). The algorithm itself (gamma, value_clip, init ...) is the # AlgorithmSpec; this carries only what one run resolves. critic: "CriticRunConfig | None" = None + # Log the training-side entropy (metrics only while entropy_coef is 0). + observe_training_entropy: bool = False + # Mamba training inter-chunk SSM state dtype ("fp32"/"bf16"; None = Megatron default). + mamba_training_ssm_states_dtype: str | None = None @dataclass(frozen=True) @@ -1074,6 +1078,8 @@ def resolve_rl_run_config( requested=getattr(args, "rl_lr_schedule", None) or "auto", full_parameter_async=getattr(args, "full_parameter_rounds", "sweep") == "async", ), + observe_training_entropy=bool(getattr(args, "observe_training_entropy", False)), + mamba_training_ssm_states_dtype=getattr(args, "mamba_training_ssm_states_dtype", None) or None, seed=args.seed, critic=resolve_critic_run_config(args), rollout_seed=getattr(args, "rollout_seed", args.seed + args.learner_id), diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 4f1c9103c..dc1ca535f 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -879,6 +879,20 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. _ROLLOUT_ENGINE = ( shared_flag("--use-rollout-routing-replay", action="store_true"), + shared_flag( + "--observe-training-entropy", + action="store_true", + help="log the training-side policy entropy (Miles --observe-training-entropy); " + "metrics only, the loss is unchanged while the entropy coefficient is 0", + ), + shared_flag( + "--mamba-training-ssm-states-dtype", + choices=("fp32", "bf16"), + default=None, + help="Mamba/hybrid models: dtype of the inter-chunk SSM states in training " + "(Megatron default: the params dtype, bf16; SGLang rollouts keep fp32 states). " + "Changes the training numerics, so it is part of the training contract", + ), shared_flag("--rollout-max-response-len", type=int, default=32768, island={"default": None, "required": True}), ) From b2521533b819244c4948b73b04ec000f5cd5a40c Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:41:47 +0000 Subject: [PATCH 180/233] run_flags: keep the new flag help backend-neutral Co-Authored-By: Claude Opus 5.5 --- yeto/rl/engine/run_flags.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index dc1ca535f..dff8b850e 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -882,15 +882,15 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. shared_flag( "--observe-training-entropy", action="store_true", - help="log the training-side policy entropy (Miles --observe-training-entropy); " - "metrics only, the loss is unchanged while the entropy coefficient is 0", + help="log the training-side policy entropy; metrics only, the loss is " + "unchanged while the entropy coefficient is 0", ), shared_flag( "--mamba-training-ssm-states-dtype", choices=("fp32", "bf16"), default=None, help="Mamba/hybrid models: dtype of the inter-chunk SSM states in training " - "(Megatron default: the params dtype, bf16; SGLang rollouts keep fp32 states). " + "(default: the params dtype, bf16; rollouts keep fp32 states). " "Changes the training numerics, so it is part of the training contract", ), shared_flag("--rollout-max-response-len", type=int, default=32768, From 2b25c30f47e87d8830249ed7d622c70d1e00a003 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:58:02 +0000 Subject: [PATCH 181/233] Ports image: Miles aa6b887d8 (deferred full-parameter export) Re-pin MILES_NEXT_COMMIT to aa6b887d8 and append 23d4de975 and aa6b887d8 to MILES_FORK_LINEAGE. aa6b887d8 = b13c170ce + deferred full-parameter export with on-demand materialize, faster finite/hash checks, push checks of small chunks inline, and the hybrid R3 layer mapping. Image ghcr.io/michaellchung/yeto-miles-ports:aa6b887-2fa8801 @sha256:2c4bd1dbea4d4f72ff2365a968dc32929e7372d2613d68498e64ed00c82eaa41 (previous b13c170-2fa8801 @sha256:814a7e27...). Co-Authored-By: Claude Opus 5.5 --- docker/miles-ports/Dockerfile | 2 +- .../build-record.json | 12 +++++ .../image-manifest.json | 50 +++++++++++++++++++ tests/test_rl_ports_image.py | 2 +- yeto/rl/adapters/miles/pins.py | 15 ++++-- 5 files changed, 75 insertions(+), 6 deletions(-) create mode 100644 openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-aa6b887-2fa8801/build-record.json create mode 100644 openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-aa6b887-2fa8801/image-manifest.json diff --git a/docker/miles-ports/Dockerfile b/docker/miles-ports/Dockerfile index 8392ea724..f6656eff5 100644 --- a/docker/miles-ports/Dockerfile +++ b/docker/miles-ports/Dockerfile @@ -19,7 +19,7 @@ FROM docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c ARG MILES_REPOSITORY=https://github.com/michaellchung/miles -ARG MILES_COMMIT=b13c170ce7611104aee28e38ed01c188e95417a2 +ARG MILES_COMMIT=aa6b887d8387d5c575f7485ad6008c0f503423a8 ARG SGLANG_REPOSITORY=https://github.com/michaellchung/sglang ARG SGLANG_COMMIT=2fa880182eefdb31e64d8ea70317be27f9f32121 # setuptools-scm version of SGLANG_COMMIT (describe v0.5.20-69-g4e4148f1b); diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-aa6b887-2fa8801/build-record.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-aa6b887-2fa8801/build-record.json new file mode 100644 index 000000000..3ce5b56f8 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-aa6b887-2fa8801/build-record.json @@ -0,0 +1,12 @@ +{ + "tag": "ghcr.io/michaellchung/yeto-miles-ports:aa6b887-2fa8801", + "digest": "sha256:2c4bd1dbea4d4f72ff2365a968dc32929e7372d2613d68498e64ed00c82eaa41", + "layer_tar_sha256": "a41bc783fde069b3d16bdcc2814554b2808b0c9fedbf62847ca13dd41a6ad499", + "base": "docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "base_index": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "miles_commit": "aa6b887d8387d5c575f7485ad6008c0f503423a8", + "sglang_commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "sglang_version": "0.5.21.dev70+g2fa8801", + "peft_version": "0.20.0", + "peft_wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-aa6b887-2fa8801/image-manifest.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-aa6b887-2fa8801/image-manifest.json new file mode 100644 index 000000000..5de233f3f --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-aa6b887-2fa8801/image-manifest.json @@ -0,0 +1,50 @@ +{ + "base": { + "index_digest": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "manifest_digest": "sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "miles_commit": "9e4260de047a704208535c0e90c531929879ab40", + "platform": "linux/amd64", + "repository": "docker.io/radixark/miles", + "sglang_commit": "e5bba1f9c433da5352e521a8bb826bab960e7836" + }, + "builder": "scripts/build_miles_ports_image.sh (crane append); docker/miles-ports/Dockerfile equivalent", + "image": "ghcr.io/michaellchung/yeto-miles-ports", + "megatron_bridge": { + "commit": "8cd3466d14d2337c8492827b3712482c2b3e4866", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_bridge_init.py)" + ] + }, + "megatron_lm": { + "commit": "a84b105473eca9c9c6e49dbf1c8aa460c5cf796e", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_lm_lazy_modelopt.patch)" + ], + "path": "/root/Megatron-LM", + "repository": "https://github.com/radixark/Megatron-LM" + }, + "miles": { + "commit": "aa6b887d8387d5c575f7485ad6008c0f503423a8", + "install": "base's editable install (pip install -e /root/miles --no-deps) kept (same path); /root is on the base's sys.path, so without PYTHONPATH import miles is a namespace package over the /root/miles checkout -- yeto runs with PYTHONPATH=$HOME/miles (regular package)", + "path": "/root/miles", + "repository": "https://github.com/michaellchung/miles", + "upstream_commit": "9e4260de047a704208535c0e90c531929879ab40" + }, + "peft": { + "install": "wheel unpacked into the base's site-packages (pip install --no-deps equivalent)", + "version": "0.20.0", + "wheel": "peft-0.20.0-py3-none-any.whl", + "wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" + }, + "pure_python_overlay": true, + "schema": 1, + "sglang": { + "commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "install": "pip install -e /sgl-workspace/sglang/python[all] --no-deps (base's editable install, same path)", + "path": "/sgl-workspace/sglang", + "repository": "https://github.com/michaellchung/sglang", + "upstream_commit": "571212b636baca45e10fa3b4da11a289123f3235", + "version": "0.5.21.dev70+g2fa8801" + }, + "tag": "aa6b887-2fa8801" +} \ No newline at end of file diff --git a/tests/test_rl_ports_image.py b/tests/test_rl_ports_image.py index a23990181..770f8aa5f 100644 --- a/tests/test_rl_ports_image.py +++ b/tests/test_rl_ports_image.py @@ -70,7 +70,7 @@ def test_build_inputs_match_the_pins(): def test_commits_and_digest_are_pinned_together_with_the_build_record(): """Commit pins, tag, Dockerfile ARGs and the image digest move together: the latest build record must name exactly the pinned commits and digest.""" - record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801" + record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-aa6b887-2fa8801" record = json.loads((record_dir / "build-record.json").read_text()) assert record["miles_commit"] == rl.MILES_NEXT_COMMIT assert record["sglang_commit"] == rl.SGLANG_NEXT_COMMIT diff --git a/yeto/rl/adapters/miles/pins.py b/yeto/rl/adapters/miles/pins.py index 9be42278d..00332cd0c 100644 --- a/yeto/rl/adapters/miles/pins.py +++ b/yeto/rl/adapters/miles/pins.py @@ -58,8 +58,10 @@ class MilesRevisionPins(NamedTuple): # Adam proof leaves the hybrid optimizer's master_param mirror out (S27 IMG2); # a477bcbe0 = + one parameter may span several fragments via flat ranges (S27); # 0b64b40c1 = + rebase_full_parameter_scheduler_anchor for peer-snapshot recovery; -# b13c170ce = + a chunked apply skips fragments the masters already hold (S27 G-X). -MILES_NEXT_COMMIT = "b13c170ce7611104aee28e38ed01c188e95417a2" +# b13c170ce = + a chunked apply skips fragments the masters already hold (S27 G-X); +# aa6b887d8 = + deferred export with on-demand materialize, faster finite/hash +# checks, hybrid R3 layer mapping (S27 G-AA). +MILES_NEXT_COMMIT = "aa6b887d8387d5c575f7485ad6008c0f503423a8" # sgl-project/sglang ``sglang-miles`` head when radixark/miles@9e4260d was # committed (upstream's Dockerfile follows that branch unpinned). SGLANG_NEXT_REPOSITORY = "https://github.com/michaellchung/sglang" @@ -133,6 +135,11 @@ class MilesRevisionPins(NamedTuple): # b13c170ce = + per-fragment full-parameter apply: a chunked apply skips # fragments the masters already hold (S27 G-X). "b13c170ce7611104aee28e38ed01c188e95417a2", + # 23d4de975 = + deferred full-parameter export with on-demand materialize, + # faster finite/hash checks, hybrid R3 layer mapping (S27 G-AA). + "23d4de975b18f76e35f2f3d1093ad774a3a68086", + # aa6b887d8 = + push checks small chunks inline (S27 G-AA); no loss path. + "aa6b887d8387d5c575f7485ad6008c0f503423a8", ) # First Miles commit whose full-parameter sync reads distributed-optimizer @@ -177,7 +184,7 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # every SHA. PUBLIC on ghcr.io/michaellchung since 2026-10-07 (anonymous # pull); a private image needs SKYPILOT_DOCKER_{USERNAME,PASSWORD,SERVER} or # --rl-image-private (yeto.launcher.registry_login_for; read:packages token). -# Tag b13c170-2fa8801; (S27 IMG4: Miles b13c170ce = 0b64b40c1 + per-fragment full-parameter apply (G-X); previous 0b64b40-2fa8801 @sha256:2b8c36f6...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/). Earlier: tag 0b64b40-2fa8801; (S27 IMG3: Miles 0b64b40c1 = cc693c693 + flat-range fragment split + scheduler anchor rebase; previous cc693c6-2fa8801 @sha256:2011cbf9... (a477bcb-2fa8801 built but voided); build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/). Earlier: tag cc693c6-2fa8801; (S27 IMG2: Miles cc693c693 = 56190a5a8 (full-parameter sync with distributed optimizer + CPU offload, G-B) + Adam proof without the hybrid master mirror; previous 56190a5-2fa8801 @sha256:5254e34e... (failed on GPU: global apply changed local Adam state), ed5f64f-2fa8801 @sha256:a4769254...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/). Earlier: tag ed5f64f-2fa8801; (S27 IMG: Miles ed5f64fd7, host-Adam round cuts (B8); previous 44d41ef-2fa8801 @sha256:5b9dc2a9...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-ed5f64f-2fa8801/). Earlier: tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt +# Tag aa6b887-2fa8801; (S27 IMG5: Miles aa6b887d8 = b13c170ce + deferred full-parameter export, on-demand materialize, faster checks, hybrid R3 layers (G-AA); previous b13c170-2fa8801 @sha256:814a7e27...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-aa6b887-2fa8801/). Earlier: tag b13c170-2fa8801; (S27 IMG4: Miles b13c170ce = 0b64b40c1 + per-fragment full-parameter apply (G-X); previous 0b64b40-2fa8801 @sha256:2b8c36f6...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-23-b13c170-2fa8801/). Earlier: tag 0b64b40-2fa8801; (S27 IMG3: Miles 0b64b40c1 = cc693c693 + flat-range fragment split + scheduler anchor rebase; previous cc693c6-2fa8801 @sha256:2011cbf9... (a477bcb-2fa8801 built but voided); build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-0b64b40-2fa8801/). Earlier: tag cc693c6-2fa8801; (S27 IMG2: Miles cc693c693 = 56190a5a8 (full-parameter sync with distributed optimizer + CPU offload, G-B) + Adam proof without the hybrid master mirror; previous 56190a5-2fa8801 @sha256:5254e34e... (failed on GPU: global apply changed local Adam state), ed5f64f-2fa8801 @sha256:a4769254...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-cc693c6-2fa8801/). Earlier: tag ed5f64f-2fa8801; (S27 IMG: Miles ed5f64fd7, host-Adam round cuts (B8); previous 44d41ef-2fa8801 @sha256:5b9dc2a9...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-22-ed5f64f-2fa8801/). Earlier: tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt # Megatron-LM patch in the layer; previous 17aeb5c-2fa8801 @sha256:b304f857...; # build record openspec/changes/rl-infra-spec/evidence/ports-image/ # 2026-10-14-44d41ef-2fa8801/). Earlier: tag 17aeb5c-2fa8801; (S24 L2.1-L2.4: Miles s24-miles-import, peft 0.20.0 and the @@ -186,7 +193,7 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # ports-image/2026-10-14-17aeb5c-2fa8801/). Earlier: tag 386ae3b-2fa8801; (S23 P0-1.11 GPU run: Miles s23-sao-streaming-roles; previous 64b591a-2fa8801 @sha256:fa2413be...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-386ae3b-2fa8801/). Earlier: (S19 #1: critic family + algo-supplement metrics baked in; previous ddce209-2fa8801 @sha256:9c252c38...) (S18 ARU-3: Miles s18-agentic-suspend = 2f7871fb2 + --agentic-suspend-between-turns; previous 2f7871f-2fa8801 @sha256:62b4f164..., efbbc63-2fa8801 @sha256:a7990076..., 8bc5223-4e4148f @sha256:4aeafd77...) ghcr.io/michaellchung (public package); linux/amd64 only. MILES_NEXT_IMAGE = ( "docker:ghcr.io/michaellchung/yeto-miles-ports@sha256:" - "814a7e276d160b5febdf0f64f526c2a6426dd47115e1abdec3db183fa18e05c0" + "2c4bd1dbea4d4f72ff2365a968dc32929e7372d2613d68498e64ed00c82eaa41" ) MILES_NEXT_IMAGE_MANIFEST = "/opt/yeto/image-manifest.json" # Nebius VM images whose /var/lib/docker already holds a docker image's From 46fde5ea484b94cc746acce153f922ccc4df2586 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:16:47 +0000 Subject: [PATCH 182/233] rl: numeric final-answer reward for one-number labels (DAPO-Math rows) Add yeto.rl.numeric_answer_reward:score (registered as "numeric_answer"). It grades the text after the last : the last \boxed{...} with balanced braces, else the last "Answer: ..." line. The answer must be one number after LaTeX clean-up (thousands separators such as 1\,007 and 1{,}440, \text{}, $, units, "x =", degrees, percent), or a simple \frac{a}{b} / a^{b}. There is no "last number in the last 200 characters" fallback. On 153,210 recorded S27 DAPO-Math responses (offline cut evaluations), the old gsm8k_grade credits 962 wrong answers (\boxed{\frac{25}{7}} read as 25, cut-off text ending in a matching number) and misses 200 right ones (\boxed{\frac{3625}{29}} = 125, \boxed{210^2} = 44100). gsm8k_grade and yeto.rl.gsm8k_reward are unchanged (golden samples and the bare-arm copy in tools/eval/efficiency/eff_reward.py still match). Co-Authored-By: Claude Opus 5.5 --- tests/test_import_boundaries.py | 1 + tests/test_rl_numeric_answer_reward.py | 140 +++++++++++++++++++++++++ yeto/rl/numeric_answer_reward.py | 29 +++++ yeto/rl/rewards/builtin.py | 106 +++++++++++++++++++ yeto/rl/rewards/registry.py | 1 + 5 files changed, 277 insertions(+) create mode 100644 tests/test_rl_numeric_answer_reward.py create mode 100644 yeto/rl/numeric_answer_reward.py diff --git a/tests/test_import_boundaries.py b/tests/test_import_boundaries.py index 7994a6840..bac4b5d66 100644 --- a/tests/test_import_boundaries.py +++ b/tests/test_import_boundaries.py @@ -54,6 +54,7 @@ CORE_FILES = ( # rewards / filters (audit R1-R3) "yeto/rl/math_reward.py", "yeto/rl/gsm8k_reward.py", "yeto/rl/length_reward.py", + "yeto/rl/numeric_answer_reward.py", "yeto/rl/filters.py", # algorithm extensions living beside the package (audit A1, A3, A5) "yeto/rl/critic_warmup.py", "yeto/rl/teacher_forcing.py", "yeto/rl/grad_audit.py", diff --git a/tests/test_rl_numeric_answer_reward.py b/tests/test_rl_numeric_answer_reward.py new file mode 100644 index 000000000..4287cae82 --- /dev/null +++ b/tests/test_rl_numeric_answer_reward.py @@ -0,0 +1,140 @@ +"""CPU tests of the numeric final-answer reward (yeto.rl.numeric_answer_reward). + +REAL_TAILS are the last 260 characters of real responses to DAPO-Math-17k rows +(the S27 task prompt "... Answer: \\boxed{$Answer} ...") from the S27 offline +cut evaluations (infra-drafts/SESSION27/logs//eval-island*/cuts/*.json, +labels from s1-runs/s27-x/eval/heldout-0.6b.jsonl). Each row stores the score of +the old ``gsm8k_grade`` and of the new grader; the old scores reproduce the +recorded rewards. NEMOTRON_HEADS are the response heads Miles printed for +Nemotron-3 Nano in s27-gr-g2/launch.stream.log ("First rollout sample" / +"Finish rollout", 512-character previews). As text cut off before an answer, +they must score 0. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from yeto.rl.rewards import registry +from yeto.rl.rewards.builtin import ( + answer_number, + final_answer, + gsm8k_grade, + last_boxed, + numeric_answer_grade, +) +from yeto.rl.rewards.types import Trajectory + +# (response tail, label, old gsm8k_grade, new numeric_answer_grade) +REAL_TAILS = [ + # GK/s27-xk-x1-k1go20-b h0 + ('\n\nFrom earlier, we have $P = 2Q$. Substitute this into the equation above:\n\n$$\nQ - 12(2Q) = -416\n\\Rightarrow Q - 24Q = -416\n\\Rightarrow -23Q = -416 \\Rightarrow Q = \\frac{416}{23} = 18\n$$\n\nThen $P = 2Q = 36$\n\nSo the value of $xy + yz + zx$ is:\n\n$$\n\\boxed{36}\n$$', + '672', 0.0, 0.0), + # GK/s27-xk-x1-k1go20-b h46 + ("m/s}$\n- Cate's speed = $6 \\, \\text{m/s}$\n\nSince they are moving towards each other, the time it takes for Dan to reach Cate is:\n\n$$\nt = \\frac{50 \\, \\text{m}}{8 + 6} = \\frac{50}{14} = \\frac{25}{7} \\approx 3.571 \\, \\text{seconds}\n$$\n\nAnswer: \\boxed{\\frac{25}{7}}", + '25', 1.0, 0.0), + # GK/s27-xk-x1-k1go20-b h56 + ("\n\\frac{6050}{3}\n$$\n\nPerforming the division:\n\n$$\n\\frac{6050}{3} = 2016.666...\n$$\n\nBut since we can't have fractional triangles, the **maximum number of non-overlapping non-collinear triangles** is thus:\n\n$$\n2016\n$$\n\n---\n\n### Final answer:\n\nAnswer: \\boxed{2016}", + '2016', 1.0, 1.0), + # GK/s27-xk-x1-k1go20-b h94 + ('^{(0)}|} = \\frac{1}{3^7}\n$$\n\nSo:\n\n$$\n\\frac{|OA_{2014}^{(7)}|}{|OA_{2014}^{(0)}|} = \\frac{1}{3^7}\n$$\n\nWe can now write:\n\n$$\n\\boxed{\\frac{1}{3^7} = \\frac{1}{2187}}\n$$\n\nSo, $ p = 1 $, $ q = 2187 $, and $ p+q = 2188 $.\n\n---\n\n### ✅ Final Answer:\n\n$$\n\\boxed{2188}\n$$', + '2188', 1.0, 1.0), + # GK/s27-xk-x1-k1go20-b h115 + ('solve for $ x $:\n\n$$\n12x - 225 = 17(-x + 200).\n$$\n\nExpand the right-hand side:\n\n$$\n12x - 225 = -17x + 3400.\n$$\n\nCombine like terms:\n\n$$\n12x + 17x = 3400 + 225,\n$$\n\n$$\n29x = 3625.\n$$\n\nSolve for $ x $:\n\n$$\nx = \\frac{3625}{29}.\n$$\n\nAnswer: \\boxed{\\frac{3625}{29}}', + '125', 0.0, 1.0), + # GK/s27-xk-x1-k1go20-b h152 + ('\nSo assume that $ x = 1 $, $ y = 2 $ and check:\n\nFirst equation: $ 1 - 3(1)(4) = 1 - 12 = -11 \\ne 3 $\n\nTry $ x = 1 $, $ y = 1.5 \\Rightarrow -1.5 $ each time.\n\nEventually, desired $ x = a^2 $ and $ y = b^2 $ such that:\n\n$$\nx^3 - 3xy^2 = 3, \\quad y^3 - 3x^2y = 3', + '3', 1.0, 0.0), + # GK/s27-xk-x1-k1go20-b h224 + ('mes 25 \\times 36 \\times 49 \\times 64 \\times 81) = (1 \\times 2^2 \\times 3^2 \\times 5^2 \\times 6^2 \\times 7^2)\n$$\n\nNow this is a square of:\n\n$$\n(2 \\times 3 \\times 5 \\times 7)^2 = 210^2\n$$\n\n---\n\n### Step 3: Conclude the result\n\n$$\n\\text{Answer: } \\boxed{210^2}\n$$', + '44100', 0.0, 1.0), + # GK/s27-xk-x1-k1go20-b h253 + (' between two options based on the remaining value.\n\nThis behavior ensures that the player who starts (Euhan) will **always be the one forced to say it**, so **Minjune has a winning strategy** for any $N > 2019$.\n\n---\n\n### Final Answer:\n\n$$\nAnswer: \\boxed{2020}', + '2020', 1.0, 1.0), + # GK/s27-xk-x1-k1go20-b h2 + ('ble value** of $a$ is:\n\n$$\n\\boxed{9997}\n$$\n\n---\n\n### Step 4: Find the sum of the **smallest** and **largest** values\n\n$$\n\\text{Smallest } a = 3 \\\\\n\\text{Largest } a = 9997 \\\\\n\\text{Sum} = 3 + 9997 = \\boxed{9997 + 3 = 10000}\n$$\n\n---\n\n**Answer: $\\boxed{10000}$**', + '10000', 1.0, 1.0), + # GK/s27-xk-x1-k1go20-b h52 + (' = 4 $\n- $ xy = 2 $, $ x + y + (x, y) = 4 $ → So this is **not valid**\n\nHence, we discard this pair.\n\n---\n\nSo the **only valid pair** is:\n\n- $ x = 2 $, $ y = 3 $\n\n---\n\n**Answer: $ \\frac{1}{\\text{Once}} $**\n\nBut when writing Answer:\n\n---\n\nAnswer: \\boxed{(2, 3)}', + '3', 1.0, 0.0), + # X1/s27-x1-k1-a h251 + ('{3}{2} \\right)^2 } = \\sqrt{ \\left( \\frac{1}{2} \\right)^2 + \\left( -1 \\right)^2 } = \\frac{1}{\\sqrt{2}}\n$$\n\nSo the area of square $PQRS$ is:\n\n$$\n\\text{Area} = s^2 = \\left( \\frac{1}{\\sqrt{2}} \\right)^2 = \\frac{1}{2}\n$$\n\n### Final Answer:\n$$\n\\boxed{\\frac{1}{2}}\n$$', + '1', 1.0, 0.0), + # X1/s27-x1-strict20-a h50 + (' A_{\\text{large}}\n$$\n$$\n14 = \\frac{4}{9} \\times A\n$$\n$$\nA = \\frac{14 \\times 9}{4} = \\frac{126}{4} = 31.5\n$$\n\nSo, the area of the large square is **31.5**, or in **fractional form** as $ \\frac{63}{2} $.\n\n---\n\n### Final Answer\n\n$$\nAnswer: \\boxed{\\frac{63}{2}}\n$$', + '63', 1.0, 0.0), +] + +# (response head as printed by Miles, label, recorded reward of the full response) +NEMOTRON_HEADS = [ + ("The cubic is \n", "5", 0.0), + ("We are given the system \n\n\\[\n\\begin{cases}\\", "18", 0.0), + ("We are given \n\n\\[\nx+", "10", 1.0), + ("The cubic can be written as \n\n\\[\n3x^{3}-3x^{2}+m=0\\qquad ( \\text{or } x", "0", 1.0), + ("Let \n\n\\[\nf_a(x)=x^{4}-5x+a ,\\qquad a\\in \\math", "0", 1.0), +] + + +@pytest.mark.parametrize("tail,label,old,new", REAL_TAILS) +def test_real_dapo_tails(tail, label, old, new): + assert gsm8k_grade(tail, label) == old + assert numeric_answer_grade(tail, label) == new + + +@pytest.mark.parametrize("head,label,_recorded", NEMOTRON_HEADS) +def test_nemotron_cut_off_heads_score_zero(head, label, _recorded): + assert numeric_answer_grade(head, label) == 0.0 + + +def test_old_fallback_credits_a_cut_off_response(): + # the "last number in the last 200 characters" fallback reads the "0" of "m=0" + head, label, _ = NEMOTRON_HEADS[3] + assert gsm8k_grade(head, label) == 1.0 + assert numeric_answer_grade(head, label) == 0.0 + + +@pytest.mark.parametrize("answer,value", [ + ("34", 34.0), ("-7", -7.0), ("1\\,007", 1007.0), ("1{,}440", 1440.0), ("270,901", 270901.0), + ("1\\!000\\!000", 1e6), ("1 000", 1000.0), ("\\$7", 7.0), ("\\text{\\textbf{2020}}", 2020.0), + ("{\\displaystyle \\textbf{2020}}", 2020.0), ("\\answer{2016}", 2016.0), ("x = 5", 5.0), + ("3 + 7 = 10", 10.0), ("45^\\circ", 45.0), ("45^{\\circ}", 45.0), ("25\\%", 25.0), + ("12 \\text{ cm}", 12.0), ("**34**", 34.0), ("\\(34\\)", 34.0), ("\\frac{3625}{29}", 125.0), + ("\\dfrac{1}{2}", 0.5), ("210^2", 44100.0), ("2^{10}", 1024.0), ("8.", 8.0), + # not one number: no credit for the last number inside an expression + ("2\\sqrt{3}", None), ("1, 2, 3", None), ("(2, 3)", None), ("3^{2188} + 1", None), + ("512 \\text{ and } 125", None), ("", None), (None, None), +]) +def test_answer_number(answer, value): + assert answer_number(answer) == value + + +def test_final_answer_rules(): + assert last_boxed("a \\boxed{\\frac{1}{2}} b") == "\\frac{1}{2}" + assert last_boxed("\\boxed{1} then \\boxed{2") == "1" # unclosed last box: cut off + assert last_boxed("no box") is None + # only the text after the last is graded + assert final_answer("\\boxed{3}Answer: \\boxed{4}") == "4" + assert final_answer("\\boxed{3}no answer") is None + # "Answer:" line when there is no box + assert final_answer("work\n**Answer:** 34") == "** 34" + assert numeric_answer_grade("work\n**Answer:** 34", "34") == 1.0 + assert numeric_answer_grade("so the total is 34", "34") == 0.0 # no last-number fallback + assert numeric_answer_grade("Answer: \\boxed{34}", "#### 34") == 1.0 # GSM8K label form + assert numeric_answer_grade("Answer: \\boxed{34}", None) == 0.0 + + +def test_entry_point_and_registry(): + from yeto.rl.numeric_answer_reward import score + + s = SimpleNamespace(response="Answer: \\boxed{1{,}440}", label="1440", metadata=None) + assert asyncio.run(score(None, s)) == 1.0 + assert s.metadata == {"success": True} + s = SimpleNamespace(response="Answer: \\boxed{\\frac{25}{7}}", label="25", metadata={}) + assert asyncio.run(score(None, s)) == 0.0 + assert s.metadata == {"success": False} + assert "numeric_answer" in registry.registered_names() + plugin = registry.load_reward("numeric_answer") + assert plugin.fn(Trajectory(response="\\boxed{7}", label="7")).value == 1.0 diff --git a/yeto/rl/numeric_answer_reward.py b/yeto/rl/numeric_answer_reward.py new file mode 100644 index 000000000..c2d84bc39 --- /dev/null +++ b/yeto/rl/numeric_answer_reward.py @@ -0,0 +1,29 @@ +"""Numeric final-answer reward for one-number labels (DAPO-Math-17k rows), in-package. + +The prompt asks for "Answer: \\boxed{$Answer}" on the last line and the label is +one number (an optional GSM8K-style "... #### N" prefix is ignored). The grader, +``yeto.rl.rewards.builtin.numeric_answer_grade``, reads the last balanced +``\\boxed{...}`` after any ```` (else the last "Answer: ..." line), +cleans LaTeX number formatting, and needs one number equal to the label. +Reward 1.0 / 0.0; ``sample.metadata["success"]`` records the same flag as +``yeto.rl.gsm8k_reward``. + +Launch as ``--reward-function yeto.rl.numeric_answer_reward:score``, or copy +this file next to the launcher as ``numeric_answer_reward.py`` and use +``--reward-function numeric_answer_reward:score`` (the s27-gr run.sh pattern). +``yeto.rl.gsm8k_reward`` is unchanged. +""" + +from __future__ import annotations + +from yeto.rl.math_reward import set_success +from yeto.rl.rewards.builtin import numeric_answer_grade as grade, numeric_answer_reward +from yeto.rl.rewards.types import Trajectory + +__all__ = ["grade", "score"] + + +async def score(args, sample, **kwargs): + result = numeric_answer_reward(Trajectory(response=sample.response or "", label=sample.label)) + set_success(sample, result.metadata["success"]) + return result.value diff --git a/yeto/rl/rewards/builtin.py b/yeto/rl/rewards/builtin.py index 22749bafe..ccb827471 100644 --- a/yeto/rl/rewards/builtin.py +++ b/yeto/rl/rewards/builtin.py @@ -54,6 +54,112 @@ def gsm8k_reward(traj: Trajectory) -> RewardResult: return RewardResult(value=value, metadata={SUCCESS_KEY: value == 1.0}) +# -- numeric final answer (DAPO-Math-17k style "Answer: \boxed{N}") ------------ +# +# ``gsm8k_grade`` above stays byte-identical (golden samples, bare-arm copy in +# tools/eval/efficiency/eff_reward.py). This grader is for rows whose label is +# one number and whose prompt asks for "Answer: \boxed{...}" on the last line. +# Differences from gsm8k_grade: +# * grades only the text after the last ````; +# * reads the last ``\boxed{...}`` with balanced braces (``\boxed{\frac{1}{2}}``); +# * if there is no ``\boxed``, reads the last "Answer: ..." line; there is no +# "last number in the last 200 characters" fallback, so a response cut off +# mid-derivation scores 0 instead of matching a stray number; +# * the answer must be one number after LaTeX clean-up (thousands separators +# ``1\,007`` / ``1{,}440`` / ``1,440``, ``\text{}``, ``$``, units, ``x = ``, +# ``^\circ``, ``\%``), or a simple ``\frac{a}{b}`` / ``a^{b}``; it never takes +# the last number of an expression (``\boxed{2\sqrt{3}}`` is not 3). + +_BOXED = re.compile(r"\\(?:boxed|fbox)\s*\{") +_ANSWER_LINE = re.compile(r"answer\s*[:：]\s*(.+?)\s*$", re.IGNORECASE | re.MULTILINE) +_NUMBER = r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)" +_TEXT_CMD = re.compile(r"\\(?:text|textbf|textrm|textit|mathrm|mathbf|mathit|bf|mbox|operatorname|answer)" + r"\s*\{([^{}]*)\}") +_FRAC = re.compile(rf"^({_NUMBER})?\\[dt]?frac\{{({_NUMBER})\}}\{{({_NUMBER})\}}$") +_POWER = re.compile(rf"^({_NUMBER})\^\{{?({_NUMBER})\}}?$") +_UNIT_TAIL = re.compile(rf"^({_NUMBER})[a-zA-Z]*\.?$") +_THOUSANDS = re.compile(r"(?<=\d)(?:\{,\}|,|\\,|\\!|\\ |~|\s)(?=\d{3}(?!\d))") +# spacing / delimiter / markup tokens removed before reading the number ("\\$" before "$") +_DROP_TOKENS = ("\\displaystyle", "\\left", "\\right", "\\!", "\\,", "\\;", "\\:", "\\ ", "~", + "\\(", "\\)", "\\[", "\\]", "*", "\\$", "$", "\\%", "%") + + +def last_boxed(text: str) -> str | None: + """Content of the last ``\\boxed{...}`` / ``\\fbox{...}`` with balanced braces; + None if there is none or its braces never close (a cut-off response).""" + starts = [m.end() for m in _BOXED.finditer(text)] + for start in reversed(starts): + depth = 1 + for i in range(start, len(text)): + if text[i] == "{": + depth += 1 + elif text[i] == "}": + depth -= 1 + if depth == 0: + return text[start:i] + return None + + +def answer_number(answer: str | None) -> float | None: + """The number an extracted final answer states, or None (see section notes).""" + if answer is None: + return None + s = str(answer) + for _ in range(3): # nested \text{\textbf{..}} + s = _TEXT_CMD.sub(r"\1", s) + # thousands separators only between digit groups: "1,2,3" stays a list (no match) + s = _THOUSANDS.sub("", s) + for token in _DROP_TOKENS: + s = s.replace(token, "") + s = re.sub(r"\^\s*\{?\\circ\}?|°", "", s) + s = re.sub(r"\s+", "", s) + s = s.split("=")[-1] # "x=5", "n = 12" + s = s.rstrip(".") + while s.startswith("{") and s.endswith("}"): + s = s[1:-1] + if re.fullmatch(_NUMBER, s): + return float(s) + m = _UNIT_TAIL.match(s) + if m: + return float(m.group(1)) + m = _FRAC.match(s) + if m and float(m.group(3)) != 0: + whole = float(m.group(1)) if m.group(1) else 0.0 + frac = float(m.group(2)) / float(m.group(3)) + return whole - frac if whole < 0 else whole + frac + m = _POWER.match(s) + if m: + try: + return float(m.group(1)) ** float(m.group(2)) + except OverflowError: + return None + return None + + +def final_answer(response: str) -> str | None: + """The final answer text of a response: last ``\\boxed{}`` after the last + ````, else the last "Answer: ..." line, else None.""" + text = (response or "").split("")[-1] + boxed = last_boxed(text) + if boxed is not None: + return boxed + lines = _ANSWER_LINE.findall(text) + return lines[-1] if lines else None + + +def numeric_answer_grade(response: str, label) -> float: + if label is None: + return 0.0 + gold = answer_number(str(label).split("####")[-1]) + pred = answer_number(final_answer(response)) + return 1.0 if (gold is not None and pred is not None and abs(pred - gold) < 1e-6) else 0.0 + + +def numeric_answer_reward(traj: Trajectory) -> RewardResult: + value = numeric_answer_grade(traj.response or "", traj.label) + return RewardResult(value=value, metadata={SUCCESS_KEY: value == 1.0}) + + # -- dense length / termination test reward ----------------------------------- def length_score(response: str, finished: float) -> float: diff --git a/yeto/rl/rewards/registry.py b/yeto/rl/rewards/registry.py index 5978cfe8f..190fd87e8 100644 --- a/yeto/rl/rewards/registry.py +++ b/yeto/rl/rewards/registry.py @@ -113,6 +113,7 @@ def _ensure_builtins() -> None: from yeto.rl.rewards import builtin for name, fn in (("math", builtin.math_reward), ("gsm8k", builtin.gsm8k_reward), + ("numeric_answer", builtin.numeric_answer_reward), ("length", builtin.length_reward), ("gdpo", builtin.gdpo_reward), ("gdpo_correctness", builtin.gdpo_correctness_reward)): register_reward(name, fn) From ca8a80be9c279702b955587861de83de790c962d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:29:00 +0000 Subject: [PATCH 183/233] rl run config: word the CP divisibility check without a backend name The context-parallel check that s27-gac-long added to the neutral core said 'Miles actor world ...'. test_rl_backend_registry requires that Miles-only wording lives in adapters/miles/run_config_rules.py only. Co-Authored-By: Claude Opus 5.5 --- yeto/rl/engine/run_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index a2604ceab..97ac98091 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -719,7 +719,7 @@ def resolve_rl_run_config( rules = _rules() rules.check_trainer_parallel(actor_gpus, tensor_parallel, pipeline_parallel) if context_parallel < 1 or actor_gpus % model_parallel: - raise ValueError("Miles actor world must be divisible by TP*PP*CP") + raise ValueError("trainer world must be divisible by TP*PP*CP") is_moe = getattr(provider, "num_moe_experts", None) is not None expert_parallel = getattr(args, "expert_parallel", None) or ( actor_gpus if is_moe else 1 From 73874f5e29aa0a1c0c97b65711aad2513ae9ef0e Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:00:15 +0000 Subject: [PATCH 184/233] rl: island lead cap, sync diagnostics and pushes computed ahead Lead cap (--rl-sync-max-local-rounds-ahead K, island --sweep-max-local-rounds-ahead; default off, old behaviour unchanged): in the asynchronous full-parameter rounds an island waits at a boundary instead of training a K-th local round on top of an unmerged push of its own. While it waits it keeps pushing open PULLs, applying broadcasts and serving peer snapshots; it stops waiting when the terminal phase starts and gives up (trains on) once the oldest blocking push has waited half of the wait timeout, before that pair could time out. With a fixed roster and K=1 a fast island stays at most one local round ahead of the slowest island, so both islands' c_steps (and tokens^2/steps weights) stay equal. The option lives on the island profile only: the layout, syncer profile, session contract and training contract hashes do not change. Diagnostics: every Miles trainer-group call is printed with its method name and elapsed time (TimedGroup); every push prints its delta, error feedback, encode and enqueue seconds, and the round event carries their per-boundary sums; each island prints host facts (CPU model and count, cgroup limits, memory, NUMA, OMP/MKL/torch threads) for the driver and for every trainer rank. Pushes computed ahead: with a deterministic codec (int4 without stochastic rounding, int8, q4) and no error feedback, up to YETO_FULL_PARAMETER_PUSH_WORKERS (default 3) worker threads compute the next pushes' deltas and wire encodings while the current push is enqueued in plan order. The frames are byte-identical to the streamed pushes. Co-Authored-By: Claude Opus 5.5 --- tests/test_s27_diagnostics.py | 105 ++++ tests/test_s27_lead_cap.py | 512 ++++++++++++++++++ tests/test_s27_push_prefetch.py | 231 ++++++++ yeto/launch/rl_args.py | 13 +- yeto/launch/rl_island_task.py | 2 + yeto/protocol.py | 138 ++++- .../miles/full_parameter/role_stream.py | 134 ++++- .../adapters/miles/full_parameter/session.py | 12 +- yeto/rl/adapters/miles/island_entry.py | 17 +- .../miles/models/chunked_full_parameter.py | 83 +++ yeto/rl/adapters/miles/run_flags.py | 7 + yeto/rl/engine/full_parameter_sync.py | 232 +++++++- yeto/rl/engine/run_flags.py | 15 + yeto/rl/host_info.py | 184 +++++++ yeto/syncer_shards.py | 9 + 15 files changed, 1654 insertions(+), 40 deletions(-) create mode 100644 tests/test_s27_diagnostics.py create mode 100644 tests/test_s27_lead_cap.py create mode 100644 tests/test_s27_push_prefetch.py create mode 100644 yeto/rl/host_info.py diff --git a/tests/test_s27_diagnostics.py b/tests/test_s27_diagnostics.py new file mode 100644 index 000000000..121ba9c9c --- /dev/null +++ b/tests/test_s27_diagnostics.py @@ -0,0 +1,105 @@ +"""S27 G-AG diagnostics: host facts, named Miles group calls, push time split.""" + +from __future__ import annotations + +import asyncio +import json + +import pytest + +from yeto.rl import host_info +from yeto.rl.adapters.miles.models import chunked_full_parameter as adapter_module +from yeto.rl.adapters.miles.models.chunked_full_parameter import ( + TimedGroup, + skip_resident_supported, + timed_group, + transport_supports, +) + + +def test_the_host_report_is_json_and_names_the_facts(): + report = host_info.host_report() + json.dumps(report) + for key in ("cpu_model", "cpu_count", "cpus_usable", "mem_total_gib", "cgroup_cpu_limit", + "cgroup_memory_limit_gib", "numa_nodes", "OMP_NUM_THREADS", "MKL_NUM_THREADS", + "torch_threads", "torch_interop_threads"): + assert key in report + assert report["cpu_count"] and report["mem_total_gib"] + + +def test_the_rank_report_never_raises(monkeypatch): + monkeypatch.setattr(host_info, "host_report", lambda: 1 / 0) + assert "ZeroDivisionError" in host_info.rank_host_report(object())["error"] + + +def test_rank_reports_print_per_rank_and_survive_a_failing_group(capsys): + class Group: + async def run_plugin(self, path, kwargs): + assert path == host_info.RANK_HOST_REPORT and kwargs == {} + return [host_info.rank_host_report(None), {"error": "x"}] + + class Broken: + async def run_plugin(self, path, kwargs): + raise RuntimeError("cell down") + + asyncio.run(host_info.report_rank_hosts(Group(), "actor")) + asyncio.run(host_info.report_rank_hosts(Broken(), "actor")) + asyncio.run(host_info.report_rank_hosts(object(), "actor")) # no run_plugin: nothing + err = capsys.readouterr().err + assert "[yeto-host] actor rank 0 {" in err and "[yeto-host] actor rank 1" in err + assert "rank report failed: RuntimeError: cell down" in err + + +class _Group: + def __init__(self): + self.calls = [] + self.flag = 7 + + async def apply_full_parameter_chunked_cut(self, cut, *, skip_resident_fragments=False): + self.calls.append(("apply", cut, skip_resident_fragments)) + return (1,) + + async def full_parameter_optimizer_states(self): + raise ValueError("proof failed") + + def sync_method(self, value): + return value + 1 + + +def test_the_timed_group_logs_each_call_and_keeps_signatures(capsys): + inner = _Group() + group = timed_group(inner, "actor") + assert isinstance(group, TimedGroup) and timed_group(group, "actor") is group + assert group.flag == 7 and group.sync_method(1) == 2 + assert asyncio.run(group.apply_full_parameter_chunked_cut("cut", skip_resident_fragments=True)) == (1,) + with pytest.raises(ValueError, match="proof failed"): + asyncio.run(group.full_parameter_optimizer_states()) + err = capsys.readouterr().err + assert "miles op=apply_full_parameter_chunked_cut role=actor elapsed_s=" in err and "ok=true" in err + assert "miles op=full_parameter_optimizer_states role=actor" in err and "ok=false" in err + assert inner.calls == [("apply", "cut", True)] + # Feature detection reads the wrapped signature. + assert skip_resident_supported(group) is skip_resident_supported(inner) is True + assert transport_supports(group, "apply_full_parameter_chunked_cut", "skip_resident_fragments", "x") + assert getattr(group, "materialize_full_parameter_fragments", None) is None + + +def test_the_timed_group_can_be_switched_off(monkeypatch): + monkeypatch.setenv(adapter_module.OP_TIMING_ENV, "0") + inner = _Group() + assert timed_group(inner, "actor") is inner + + +def test_the_round_event_reports_each_boundarys_push_time_split(): + from types import SimpleNamespace + + from yeto.rl.engine.full_parameter_sync import FullParameterSync + + actor = SimpleNamespace(push_timing_totals={"count": 2, "delta_s": 3.0, "total_s": 5.0}) + sync = FullParameterSync(SimpleNamespace(streams={"actor": actor}), groups={"actor": None}, + run=None, trajectories=None) + assert sync._push_timing_fields() == {"sync/push_count": 2, "sync/push_delta_s": 3.0, "sync/push_total_s": 5.0} + actor.push_timing_totals.update(count=5, delta_s=4.5, total_s=9.0) + assert sync._push_timing_fields() == {"sync/push_count": 3, "sync/push_delta_s": 1.5, "sync/push_total_s": 4.0} + assert FullParameterSync(SimpleNamespace(streams={"actor": SimpleNamespace()}), groups={"actor": None}, + run=None, trajectories=None)._push_timing_fields() == {} diff --git a/tests/test_s27_lead_cap.py b/tests/test_s27_lead_cap.py new file mode 100644 index 000000000..3a0cf2484 --- /dev/null +++ b/tests/test_s27_lead_cap.py @@ -0,0 +1,512 @@ +"""S27 G-AG: the local-round lead cap of the asynchronous full-parameter rounds. + +``--rl-sync-max-local-rounds-ahead K`` (island ``--sweep-max-local-rounds-ahead``) +makes an island wait at a boundary instead of training a K-th local round on +top of an unmerged push of its own. With a fixed roster (quorum 2 of 2) every +merge needs every island, so with K = 1 a fast island stays at most one local +round ahead of a slow one instead of training rounds whose progress the next +broadcast overwrites and whose c_steps inflate its tokens^2/steps merge weight. + +Covered here: the pairing rule on fakes (who blocks, the half-timeout escape, +the terminal phase, peer snapshots during the wait), the profile and the +launcher/island wiring (default off, hashes unchanged), and two islands of +the real pairing and session (:class:`FullParameterSync`) against the real +Rust syncer: a fast island and a slow island, with and without the cap. +""" + +from __future__ import annotations + +import asyncio +import threading +import time +from dataclasses import replace +from types import SimpleNamespace + +import numpy as np +import pytest +import test_rl_full_parameter_async as harness +from test_rl_full_parameter_async import ( # noqa: F401 - pytest fixtures + ASYNC, + _flag, + _hmac_key, + _plan, + needs_cargo, + syncer_binary, +) +from test_rl_full_parameter_sync import _PairedRole + +from yeto.protocol import PullRequest +from yeto.rl.adapters.miles.full_parameter.role_stream import RoleSubmissionCandidate +from yeto.rl.engine.full_parameter_sync import ( + FullParameterSync, + FullParameterSyncProfile, + RolePairing, +) + +# -- the pairing rule on fakes ---------------------------------------------------- + + +def _profile(**overrides): + values = {"roles": ("actor",), "total_fragment_steps": 40, "fragments": 2, "local_horizon": 1, + "pipeline_depth": 2, "wait_timeout": 900.0, "poll_seconds": 0.001, "partial_quorum": True} + values.update(overrides) + return FullParameterSyncProfile(**values) + + +def _candidate(global_step: int, local_step: int) -> RoleSubmissionCandidate: + return RoleSubmissionCandidate( + role="actor", permit=PullRequest((global_step - 1) % 2, global_step, 1), + base_version=max(0, global_step - 2), local_step=local_step, c_steps=1, c_tokens=11, + payload_bytes=8, submission_cut=object(), anchor=object()) + + +def _pairing(role, **profile): + role.fragment_versions = (0, 0) + return RolePairing({"actor": role}, _profile(**profile)) + + +def test_the_lead_cap_is_off_by_default_and_needs_the_streaming_profile(): + assert _profile().max_local_rounds_ahead is None + assert _profile(max_local_rounds_ahead=1).max_local_rounds_ahead == 1 + with pytest.raises(ValueError, match="integer >= 1"): + _profile(max_local_rounds_ahead=0) + with pytest.raises(ValueError, match="streaming profile"): + _profile(sweep=True, partial_quorum=False, fragments=2, pipeline_depth=2, + local_horizon=1, max_local_rounds_ahead=1) + + +def test_k_minus_one_rounds_may_train_on_top_of_an_unmerged_push(): + role = _PairedRole("actor", [_candidate(1, 1), _candidate(2, 1)], [], release_broadcasts=False) + pairing = _pairing(role, max_local_rounds_ahead=1) + pairing.local_round = 1 + pairing.submit_ready() + assert sorted(pair.local_round for pair in pairing.inflight.values()) == [1, 1] + assert len(pairing.lead_blockers()) == 2 # K = 1: no round on top of an unmerged push + pairing.profile = replace(pairing.profile, max_local_rounds_ahead=2) + assert pairing.lead_blockers() == [] # K = 2: one round on top is allowed + pairing.local_round = 2 + assert len(pairing.lead_blockers()) == 2 # ... but not a second one + off = _pairing(_PairedRole("actor", [_candidate(1, 1)], [], release_broadcasts=False)) + off.local_round = 1 + off.submit_ready() + off.local_round = 9 + assert off.lead_blockers() == [] + + +@pytest.mark.asyncio +async def test_the_wait_ends_when_the_blocking_round_merges_and_serves_peers_meanwhile(): + role = _PairedRole("actor", [_candidate(1, 1), _candidate(2, 1)], [], release_broadcasts=False) + pairing = _pairing(role, max_local_rounds_ahead=1) + pairing.local_round = 1 + pairing.submit_ready() + served = [] + + async def serve(): + served.append(time.monotonic()) + if len(served) == 3: + role.release_broadcasts = True # the peer pushed: both rounds merge + return 0 + + waited = await pairing.wait_for_lead(serve) + assert waited > 0 and pairing.inflight == {} + assert role.fragment_versions == (1, 2) + assert len(served) >= 3 + assert (pairing.lead_waits, pairing.lead_wait_escapes) == (1, 0) + + +@pytest.mark.asyncio +async def test_the_wait_gives_up_at_half_the_wait_timeout_and_never_times_out_itself(): + role = _PairedRole("actor", [_candidate(1, 1)], [], release_broadcasts=False) + pairing = _pairing(role, max_local_rounds_ahead=1, wait_timeout=0.4) + pairing.local_round = 1 + pairing.submit_ready() + started = time.monotonic() + await pairing.wait_for_lead() + assert 0.15 <= time.monotonic() - started < 0.4 # half of 0.4 s, before the pair's timeout + assert (pairing.lead_waits, pairing.lead_wait_escapes) == (1, 1) + assert len(pairing.inflight) == 1 and pairing.failure is None + + +@pytest.mark.asyncio +async def test_no_wait_in_the_terminal_phase_or_without_blockers(): + role = _PairedRole("actor", [_candidate(1, 1)], [], release_broadcasts=False) + pairing = _pairing(role, max_local_rounds_ahead=2) + pairing.local_round = 1 + pairing.submit_ready() + assert await pairing.wait_for_lead() == 0.0 # K = 2: one round on top is allowed + pairing.local_round = 2 + pairing.terminal_seen = True # the terminal barrier waits instead + assert await pairing.wait_for_lead() == 0.0 + assert pairing.lead_waits == 0 + + +# -- launcher and island wiring ------------------------------------------------------ + + +def test_the_launcher_forwards_the_cap_to_every_island_and_no_hash_moves(): + base = ASYNC + ("--rl-sync-fragment-pipeline", "2") + _, plain = _plan(base) + _, capped = _plan(base + ("--rl-sync-max-local-rounds-ahead", "1")) + for before, after in zip(plain["island_requests"], capped["island_requests"], strict=True): + assert "--sweep-max-local-rounds-ahead" not in before["learner_command"] + assert _flag(after["learner_command"], "--sweep-max-local-rounds-ahead") == "1" + for flag in ("--syncer-profile-sha256", "--parameter-layout-sha256"): + assert _flag(after["learner_command"], flag) == _flag(before["learner_command"], flag) + assert [s["command"] for s in plain["role_syncers"]] == [s["command"] for s in capped["role_syncers"]] + + +@pytest.mark.parametrize(("extra", "message"), [ + (("--rl-sync-max-local-rounds-ahead", "1"), "async"), # the strict sweep + (("--rl-sync-full-parameter-rounds", "async", "--rl-sync-max-local-rounds-ahead", "0"), "1..1000"), +]) +def test_the_launcher_refuses_a_cap_outside_the_async_rounds(extra, message): + with pytest.raises((ValueError, SystemExit), match=message): + _plan(harness.FULL + extra) + + +def test_the_island_builds_the_profile_with_the_cap(): + from yeto.rl.adapters.miles.island_entry import _max_local_rounds_ahead + + assert _max_local_rounds_ahead(SimpleNamespace(), True) is None + assert _max_local_rounds_ahead(SimpleNamespace(sweep_max_local_rounds_ahead=1), True) == 1 + with pytest.raises(ValueError, match="async"): + _max_local_rounds_ahead(SimpleNamespace(sweep_max_local_rounds_ahead=1), False) + with pytest.raises(ValueError, match="1..1000"): + _max_local_rounds_ahead(SimpleNamespace(sweep_max_local_rounds_ahead=0), True) + + +def test_the_cap_stays_out_of_the_session_contract(): + """The cap lives on FullParameterSyncConfig, not on RoleStreamConfig, so + the session contract hash (HELLO) of a capped island is unchanged.""" + from test_rl_full_parameter_sync import _role_stream_config + + from yeto.rl.adapters.miles.full_parameter.session import FullParameterSyncConfig + + config = _role_stream_config("actor", 30100) + config = replace(config, partial_quorum=True) + plain = FullParameterSyncConfig((config,)) + capped = FullParameterSyncConfig((config,), max_local_rounds_ahead=1) + assert plain.roles == capped.roles + assert capped.profile().max_local_rounds_ahead == 1 and plain.profile().max_local_rounds_ahead is None + + +# -- two islands on the real Rust syncer --------------------------------------------- + +FRAGMENTS = 4 +PIPELINE = 2 +OUTER_ROUNDS = 3 # every fragment merges 3 times +NUMEL = 3 +TOKENS = 5 # tokens per local round: tokens^2/steps = c_steps * 25 + + +def _initial(fid: int) -> np.ndarray: + return np.asarray([float(fid) + 0.25 * i for i in range(NUMEL)], dtype=np.float32) + + +def _step(learner: int, local_round: int, fid: int) -> np.ndarray: + """The local training update of one round (deterministic, learner-specific).""" + return np.asarray([0.01 * (learner + 1) + 0.001 * local_round + 0.0001 * fid + 0.00001 * i + for i in range(NUMEL)], dtype=np.float32) + + +class _WireRole: + """A RoleStream over a real SyncerClient: f32 pushes of ``cut - anchor``, + broadcasts overwrite the fragment (alpha = 0), FINAL ends the run.""" + + monotonic_broadcasts = True + role = "actor" + + def __init__(self, learner: int, client, record) -> None: + self.learner = learner + self.client = client + self.record = record + self.params = {f: _initial(f) for f in range(FRAGMENTS)} + self.anchor = {f: _initial(f) for f in range(FRAGMENTS)} + self.versions = [0] * FRAGMENTS + self.rounds = 0 + self.rounds_at_anchor = [0] * FRAGMENTS + self.cut: dict | None = None + self.cut_versions: tuple[int, ...] | None = None + self.cut_rounds = 0 + self.pending: dict[int, PullRequest] = {} + self.staged: dict[tuple[int, int], np.ndarray] = {} + self.terminal_submitted = False + self.final_acknowledged = False + self.manifest = None + self.final: dict[int, np.ndarray] | None = None + + # -- RoleStream --------------------------------------------------------- + @property + def fragment_versions(self): + return tuple(self.versions) + + @property + def content_hash(self): + return "x" + + @property + def finalizing(self): + return self.client.finalizing.is_set() + + async def initialize(self, group): + """Like the Miles stream, which waits for the syncer's initial cut: + training starts once the syncer opened its first rounds.""" + del group + deadline = time.monotonic() + 20 + while not self.pending: + assert time.monotonic() < deadline, "the syncer opened no round" + self._drain() + await asyncio.sleep(0.005) + + async def complete_local_round(self, *, rollout_id, trajectories): + del trajectories + self.rounds += 1 + for f in range(FRAGMENTS): + self.params[f] = self.params[f] + _step(self.learner, self.rounds, f) + self.cut = {f: value.copy() for f, value in self.params.items()} + self.cut_versions = tuple(self.versions) + self.cut_rounds = self.rounds + self.record("round", self.learner, rollout_id + 1) + + def resident_state(self, policy_version): + return SimpleNamespace(policy_version=policy_version, policy_tensor_hash=lambda: "h") + + def _drain(self): + self.client.check_health() + for permit in self.client.drain_pulls(): + known = self.pending.get(permit.fragment_id) + if known is None or (permit.global_step, permit.round_attempt) > ( + known.global_step, known.round_attempt): + self.pending[permit.fragment_id] = permit + + def ready_submission_candidates(self, include_deferred=False): + del include_deferred + if self.cut is None or self.final is not None: + return () + self._drain() + out = [] + for f, permit in sorted(self.pending.items()): + base = max(0, permit.global_step - FRAGMENTS) + c_steps = self.cut_rounds - self.rounds_at_anchor[f] + if self.versions[f] != base or self.cut_versions[f] != base or c_steps < 1: + continue + out.append(RoleSubmissionCandidate( + role="actor", permit=permit, base_version=base, local_step=self.cut_rounds, + c_steps=c_steps, c_tokens=c_steps * TOKENS, payload_bytes=NUMEL * 4, + submission_cut=self.cut, anchor=self.anchor[f])) + return tuple(out) + + def retain_candidate(self, candidate): + del candidate + + def release_candidate(self, candidate): + del candidate + + def push_candidate(self, candidate): + f = candidate.permit.fragment_id + if self.pending.get(f) != candidate.permit: + return None + delta = candidate.submission_cut[f] - candidate.anchor + if self.client.push_fragment_parts( + f, candidate.permit.global_step, candidate.permit.round_attempt, candidate.base_version, + candidate.local_step, candidate.c_steps, candidate.c_tokens, + [delta.astype(np.float32).tobytes()]) is not True: + return None + del self.pending[f] + self.terminal_submitted |= candidate.permit.global_step == FRAGMENTS * OUTER_ROUNDS + self.record("push", self.learner, (candidate.permit.global_step, candidate.c_steps, + candidate.local_step)) + return SimpleNamespace(role="actor", fragment_id=f, global_step=candidate.permit.global_step, + round_attempt=candidate.permit.round_attempt, + base_version=candidate.base_version, c_steps=candidate.c_steps, + c_tokens=candidate.c_tokens, payload_bytes=NUMEL * 4, + pull_to_push_seconds=0.0) + + async def stage_available_broadcasts(self): + self.client.check_health() + for update in self.client.drain_updates(): + if update.version > self.versions[update.fragment_id]: + self.staged[(update.fragment_id, update.version)] = np.frombuffer( + bytes(update.data), dtype=np.float32)[-NUMEL:].copy() + return tuple(self.staged) + + @property + def staged_fragment_versions(self): + return tuple(self.staged) + + def has_staged_broadcast(self, fragment_id, version): + return (fragment_id, version) in self.staged + + async def apply_staged_broadcast(self, fragment_id, version, *, catch_up=False): + del catch_up + value = self.staged.pop((fragment_id, version)) + self.params[fragment_id] = value.copy() + self.anchor[fragment_id] = value.copy() + self.versions[fragment_id] = version + self.rounds_at_anchor[fragment_id] = self.rounds + permit = self.pending.get(fragment_id) + if permit is not None and permit.global_step <= version: + del self.pending[fragment_id] + return (fragment_id,) + + async def wait_and_apply_final(self): + manifest, fragments = await asyncio.to_thread(self.client.wait_for_final_fragments, 30) + self.final = {f.fragment_id: np.frombuffer(bytes(f.data), dtype=np.float32)[-NUMEL:].copy() + for f in fragments} + self.params = {f: value.copy() for f, value in self.final.items()} + self.versions = list(manifest.versions) + self.manifest = manifest + + def acknowledge_final(self): + self.client.acknowledge_finalization(self.manifest) + self.final_acknowledged = True + + def discard_detached_submission_cut(self): + pass + + def close(self): + self.client.close() + + +def _island(learner, client, profile, train_s, record, errors): + loop = asyncio.new_event_loop() + role = _WireRole(learner, client, record) + sync = FullParameterSync(RolePairing({"actor": role}, profile), groups={"actor": None}, + run=loop.run_until_complete, trajectories=lambda rollout_id: None) + driver = SimpleNamespace(emit=lambda name, **fields: record("event", learner, (name, fields)), + outer_version=None) + try: + sync.start(driver) + for rollout_id in range(1000): + record("start", learner, rollout_id + 1) + time.sleep(train_s) + if sync.boundary(driver, rollout_id=rollout_id, stats=None).stop: + sync.published(driver, rollout_id=rollout_id, policy_hash="h") + sync.finish(driver) + break + except BaseException as error: # noqa: BLE001 - re-raised by the test + errors.append(error) + finally: + sync.close() + loop.close() + return role + + +def _run_two_islands(binary, tmp_path, monkeypatch, *, cap, fast_s=0.005, slow_s=0.25): + monkeypatch.setattr(harness, "P", FRAGMENTS) + monkeypatch.setattr(harness, "NUMEL", NUMEL) + extra = ("--rl-sync-fragment-pipeline", str(PIPELINE), "--rl-sync-quorum", "2", + "--rl-sync-grace-ms", "0", "--rl-sync-max-base-lag", "0", + "--rl-sync-learner-weight", "tokens2-over-steps", + "--rl-sync-outer-lr", "1.0", "--rl-sync-outer-momentum", "0") + tmp_path.mkdir(parents=True, exist_ok=True) + log_file = tmp_path / "syncer.log" + process, port, profile_hash = harness._start_async_syncer( + binary, tmp_path, *extra, total_steps=OUTER_ROUNDS, log_file=log_file) + profile = _profile(total_fragment_steps=FRAGMENTS * OUTER_ROUNDS, fragments=FRAGMENTS, + pipeline_depth=PIPELINE, wait_timeout=60.0, poll_seconds=0.002, + max_local_rounds_ahead=cap) + lock = threading.Lock() + log: list[tuple[float, str, int, object]] = [] + + def record(kind, learner, value): + with lock: + log.append((time.monotonic(), kind, learner, value)) + + clients = [harness._client(learner, port, profile_hash) for learner in (0, 1)] + errors: list[BaseException] = [] + roles: dict[int, _WireRole] = {} + try: + for client in clients: + client.start() + for fid in range(FRAGMENTS): + clients[0].send_init(fid, _initial(fid).tobytes()) + threads = [threading.Thread( + target=lambda learner=learner, train_s=train_s: roles.__setitem__( + learner, _island(learner, clients[learner], profile, train_s, record, errors))) + for learner, train_s in ((0, fast_s), (1, slow_s))] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=120) + assert not any(thread.is_alive() for thread in threads), "islands did not finish" + if errors: + raise errors[0] + assert process.wait(timeout=15) == 0 + finally: + for client in clients: + client.close() + if process.poll() is None: + process.terminate() + process.wait(timeout=10) + return SimpleNamespace(roles=roles, log=log, syncer_log=log_file.read_text()) + + +def _max_lead(log, fast=0, slow=1) -> int: + """Largest (rounds the fast island completed) - (rounds the slow island + completed) at the moment the fast island starts training a round.""" + completed = {fast: 0, slow: 0} + lead = 0 + for _, kind, learner, value in sorted(log, key=lambda item: item[0]): + if kind == "round": + completed[learner] = value + elif kind == "start" and learner == fast: + lead = max(lead, (value - 1) - completed[slow]) + return lead + + +def _pushes(log, learner): + return {step: (c_steps, local) for _, kind, who, value in log if kind == "push" and who == learner + for step, c_steps, local in [value]} + + +def _check_merges(run): + steps = range(1, FRAGMENTS * OUTER_ROUNDS + 1) + for learner in (0, 1): + assert sorted(_pushes(run.log, learner)) == list(steps), learner # every merge has both islands + assert "push rejected" not in run.syncer_log and "OutOfRound" not in run.syncer_log + finals = [run.roles[learner].final for learner in (0, 1)] + for fid in range(FRAGMENTS): + assert np.array_equal(finals[0][fid], finals[1][fid]) + assert np.all(np.isfinite(finals[0][fid])) + return finals[0] + + +@needs_cargo +@pytest.mark.usefixtures("_hmac_key") +def test_real_syncer_fast_island_stays_within_the_cap_and_both_islands_join_every_merge( + syncer_binary, tmp_path, monkeypatch): # noqa: F811 - the imported fixture + capped = _run_two_islands(syncer_binary, tmp_path / "capped", monkeypatch, cap=1) + finals = _check_merges(capped) + assert _max_lead(capped.log) <= 1 + fast = _pushes(capped.log, 0) + slow = _pushes(capped.log, 1) + # Bounded local progress per push: the fast island trains at most K + 1 + # rounds on a fragment before it pushes it, so its merge weight + # (c_steps * T^2) stays within (K + 1)x of the slow island's. + assert max(c for c, _ in fast.values()) <= 2 + assert max(c for c, _ in slow.values()) <= 2 + waits = [fields for _, kind, who, (name, fields) in + ((t, k, w, v) for t, k, w, v in capped.log if k == "event") + if who == 0 and name == "rl_full_parameter_round"] + assert waits[-1]["sync/max_local_rounds_ahead"] == 1 + assert waits[-1]["sync/lead_waits"] >= 1 and waits[-1]["sync/lead_wait_escapes"] == 0 + # Same inputs, same outcome: the capped run is reproducible bit for bit. + again = _run_two_islands(syncer_binary, tmp_path / "again", monkeypatch, cap=1) + finals_again = _check_merges(again) + for fid in range(FRAGMENTS): + assert finals[fid].view(np.int32).tolist() == finals_again[fid].view(np.int32).tolist() + assert _pushes(again.log, 0) == fast and _pushes(again.log, 1) == slow + + +@needs_cargo +@pytest.mark.usefixtures("_hmac_key") +def test_real_syncer_without_the_cap_the_fast_island_runs_ahead(syncer_binary, tmp_path, monkeypatch): # noqa: F811 + """The control (old behaviour): same islands, no cap. Every merge still + has both islands and the result is finite, but the fast island runs many + rounds ahead and pushes inflated c_steps.""" + free = _run_two_islands(syncer_binary, tmp_path, monkeypatch, cap=None) + _check_merges(free) + assert _max_lead(free.log) > 1 + assert max(c for c, _ in _pushes(free.log, 0).values()) > 2 + assert not any(kind == "event" and "sync/lead_waits" in value[1] for _, kind, _, value in free.log) diff --git a/tests/test_s27_push_prefetch.py b/tests/test_s27_push_prefetch.py new file mode 100644 index 000000000..a8cf0ce11 --- /dev/null +++ b/tests/test_s27_push_prefetch.py @@ -0,0 +1,231 @@ +"""S27 G-AG: pushes computed ahead on worker threads are byte-identical. + +``YETO_FULL_PARAMETER_PUSH_WORKERS`` (default 3) lets the pairing compute the +next pushes' deltas and wire encodings (``prepare_push``) on worker threads +while the pushing thread enqueues the current one, in plan order. Covered: +the client's encode-ahead frame equals the streamed frame for every +deterministic codec (with and without tensor segments); the role stream's +prepared push sends the same bytes and keeps the same bookkeeping; the +pairing pushes in plan order, prefetches at most ``workers`` pushes, and +falls back to the streamed push when a role cannot encode ahead (error +feedback, stochastic rounding) or the variable is 1; and the per-push time +split reaches the round event. +""" + +from __future__ import annotations + +import threading +import time +from dataclasses import replace +from types import SimpleNamespace + +import numpy as np +import pytest +from test_rl_full_parameter_sync import _PairedRole, _role_stream_config + +from yeto.fragments import MERGE_AVG, Fragment, FragmentLayout +from yeto.protocol import ( + DTYPE_F32_PUSH_I4, + DTYPE_F32_PUSH_I8, + DTYPE_F32_PUSH_Q4, + DTYPE_FLAG_TENSOR_BLOCKS, + PullRequest, + SyncerClient, +) +from yeto.rl.adapters.miles.full_parameter.role_stream import ( + MilesFullParameterRoleStream, + RoleSubmissionCandidate, +) +from yeto.rl.engine import full_parameter_sync as sync_module +from yeto.rl.engine.full_parameter_sync import FullParameterSyncProfile, RolePairing + +SIZES = (700, 1300, 37) + + +def _layout(): + return FragmentLayout([Fragment(MERGE_AVG, [(f"w{f}.{i}", n) for i, n in enumerate(SIZES)]) + for f in range(2)]) + + +def _client(dtype: int, **kwargs) -> SyncerClient: + client = SyncerClient(("127.0.0.1", 9), 0, _layout(), dtype=dtype, num_streams=1, connect_timeout=1, + max_reconnects=0, **kwargs) + client.sent = [] + + def capture(msg_type, prefix, parts, size, *, label, before_last_enqueue=None): + data = b"".join(bytes(part) for part in parts) + assert len(data) == size, label + client.sent.append((msg_type, bytes(prefix), data)) + return True + + client._send_large_parts = capture + return client + + +def _delta(seed: int) -> np.ndarray: + return (np.random.default_rng(seed).standard_normal(sum(SIZES)) * 0.01).astype(np.float32) + + +def _parts(delta: np.ndarray): + raw = delta.tobytes() + return [raw[:1000], raw[1000:5001], raw[5001:]] # unaligned to the codec blocks + + +@pytest.mark.parametrize("dtype", [DTYPE_F32_PUSH_I4, DTYPE_F32_PUSH_I8, DTYPE_F32_PUSH_Q4, + DTYPE_F32_PUSH_I4 | DTYPE_FLAG_TENSOR_BLOCKS]) +def test_an_encoded_ahead_push_frame_equals_the_streamed_frame(dtype): + client = _client(dtype) + assert client.can_encode_ahead + delta = _delta(1) + timings: dict[str, float] = {} + assert client.push_fragment_parts(1, 3, 1, 1, 2, 2, 9, _parts(delta), timings=timings) is True + assert {"produce_s", "push_s"} <= set(timings) + encoded = client.encode_push_parts(1, _parts(delta), timings={}) + assert client.push_fragment_parts(1, 3, 1, 1, 2, 2, 9, (), encoded=encoded) is True + assert client.sent[0] == client.sent[1] + with pytest.raises(ValueError, match="another fragment"): + client.push_fragment_parts(0, 2, 1, 0, 2, 2, 9, (), encoded=encoded) + + +def test_stochastic_rounding_and_f32_are_not_encoded_ahead(): + assert not _client(DTYPE_F32_PUSH_I4, push_i4_stochastic_seed=3).can_encode_ahead + from yeto.protocol import DTYPE_F32 + + plain = _client(DTYPE_F32) + assert not plain.can_encode_ahead + with pytest.raises(ValueError, match="cannot be encoded ahead"): + plain.encode_push_parts(0, _parts(_delta(0))) + + +def _stream(dtype_name="int4", **config): + """A Miles role stream with a fake adapter over a real (unconnected) client.""" + config = replace(_role_stream_config("actor", 30100), push_dtype=dtype_name, + expected_fragments=2, total_fragment_steps=8, **config) + stream = MilesFullParameterRoleStream(config) + stream.client = _client(config.session_dtype) + deltas = {0: _delta(10), 1: _delta(11)} + threads = [] + + def delta_parts(anchor, cut, fragment_id, **_kwargs): + def parts(): + threads.append(threading.current_thread().name) + return iter(_parts(deltas[fragment_id] + np.float32(cut.shift))) + return parts + + stream.adapter = SimpleNamespace(layout=SimpleNamespace(fragments=_layout()), + delta_parts_from_authoritative=delta_parts) + stream.threads = threads + return stream + + +def _candidate(stream, step: int, cut, anchor): + permit = PullRequest((step - 1) % 2, step, 1) + stream._pending[permit.fragment_id] = permit + stream._pinned_cuts[id(cut)] = (cut, 1) + stream._pinned_anchors[id(anchor)] = (anchor, 1) + return RoleSubmissionCandidate("actor", permit, max(0, step - 2), 1, 1, 7, sum(SIZES) * 4, cut, anchor) + + +def test_the_role_stream_sends_the_same_bytes_with_a_prepared_push(): + streamed, prepared = _stream(), _stream() + cut, anchor = SimpleNamespace(shift=0.5), object() + assert prepared.encodes_ahead + first = streamed.push_candidate(_candidate(streamed, 1, cut, anchor)) + candidate = _candidate(prepared, 1, cut, anchor) + ahead = prepared.prepare_push(candidate) + second = prepared.push_candidate(candidate, ahead) + assert streamed.client.sent == prepared.client.sent + assert replace(first, pull_to_push_seconds=0) == replace(second, pull_to_push_seconds=0) + assert prepared._pending == streamed._pending == {} + for stream in (streamed, prepared): + assert stream.push_timing_totals["count"] == 1 + assert {"delta_s", "encode_s", "enqueue_s", "total_s"} <= set(stream.push_timing_totals) + with pytest.raises(RuntimeError, match="another candidate"): + other = _candidate(prepared, 2, SimpleNamespace(shift=0.0), anchor) + prepared.push_candidate(other, ahead) + + +def test_error_feedback_keeps_the_streamed_push(): + stream = _stream(push_int4_error_feedback="bf16") + assert stream.error_feedback is not None and not stream.encodes_ahead + + +# -- the pairing --------------------------------------------------------------------- + + +class _PreparingRole(_PairedRole): + """Records the thread and the order of prepares and pushes.""" + + encodes_ahead = True + + def __init__(self, *args, delay=0.0, **kwargs): + super().__init__(*args, **kwargs) + self.delay = delay + self.prepared_on: dict[int, str] = {} + self.pushed: list[tuple[int, object]] = [] + self.active = 0 + self.max_active = 0 + self.lock = threading.Lock() + + def prepare_push(self, candidate): + with self.lock: + self.active += 1 + self.max_active = max(self.max_active, self.active) + time.sleep(self.delay) + with self.lock: + self.active -= 1 + self.prepared_on[candidate.permit.global_step] = threading.current_thread().name + return ("bytes", candidate.permit.global_step) + + def push_candidate(self, candidate, prepared=None): + self.pushed.append((candidate.permit.global_step, prepared)) + return super().push_candidate(candidate) + + +def _candidates(steps): + return [RoleSubmissionCandidate("actor", PullRequest((s - 1) % 8, s, 1), 0, 1, 1, 7, 8, object(), object()) + for s in steps] + + +def _pairing(role, depth=8): + role.fragment_versions = (0,) * 8 + profile = FullParameterSyncProfile(roles=("actor",), total_fragment_steps=64, fragments=8, local_horizon=1, + pipeline_depth=depth, wait_timeout=900.0, poll_seconds=0.001, + partial_quorum=True) + return RolePairing({"actor": role}, profile) + + +def test_the_pairing_pushes_prepared_bytes_in_plan_order_with_bounded_lookahead(monkeypatch): + monkeypatch.setenv(sync_module.PUSH_WORKERS_ENV, "3") + role = _PreparingRole("actor", _candidates(range(1, 9)), [], release_broadcasts=False, delay=0.02) + progress = _pairing(role).submit_ready() + assert [value.global_step for value in progress["actor"].submissions] == list(range(1, 9)) + assert role.pushed == [(step, ("bytes", step)) for step in range(1, 9)] + assert all(name.startswith("yeto-push-prepare") for name in role.prepared_on.values()) + assert 1 < role.max_active <= 3 + + +@pytest.mark.parametrize("workers", ["1", "bad"]) +def test_one_worker_or_a_bad_value_keeps_the_streamed_push(monkeypatch, workers): + monkeypatch.setenv(sync_module.PUSH_WORKERS_ENV, workers) + role = _PreparingRole("actor", _candidates(range(1, 4)), [], release_broadcasts=False) + _pairing(role).submit_ready() + assert [prepared for _, prepared in role.pushed] == [None, None, None] and role.prepared_on == {} + + +def test_a_failed_prepare_fail_stops_and_leaves_no_worker_running(monkeypatch): + monkeypatch.setenv(sync_module.PUSH_WORKERS_ENV, "3") + + class _Failing(_PreparingRole): + def prepare_push(self, candidate): + if candidate.permit.global_step == 2: + raise ValueError("encode failed") + return super().prepare_push(candidate) + + role = _Failing("actor", _candidates(range(1, 7)), [], release_broadcasts=False, delay=0.01) + pairing = _pairing(role) + with pytest.raises(ValueError, match="encode failed"): + pairing.submit_ready() + assert pairing.failure is not None and role.closed + assert [step for step, _ in role.pushed] == [1] + assert role.active == 0 diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 3d8e177d1..f7dbb5c95 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -9,7 +9,11 @@ from pathlib import Path from ..gpu_spec import parse_gpu_spec -from ..rl.engine.run_flags import FULL_PARAMETER_ASYNC_DEFAULT_PIPELINE, FULL_PARAMETER_ASYNC_MAX_PIPELINE +from ..rl.engine.run_flags import ( + FULL_PARAMETER_ASYNC_DEFAULT_PIPELINE, + FULL_PARAMETER_ASYNC_MAX_PIPELINE, + FULL_PARAMETER_MAX_LOCAL_ROUNDS_AHEAD, +) from ..rl.engine.selection import ( ASYNC_GRACE_CAP_MS, ASYNC_NESTEROV_MIN_H, @@ -174,6 +178,7 @@ def sao_contract_mounts(args, learner_id: int) -> dict[str, str]: ("rl_sync_expert_token_weights", False, "--rl-sync-expert-token-weights"), ("rl_sync_init_mode", "auto", "--rl-sync-init-mode"), ("rl_sync_wait_timeout_s", 900.0, "--rl-sync-wait-timeout-s"), + ("rl_sync_max_local_rounds_ahead", None, "--rl-sync-max-local-rounds-ahead"), ) @@ -247,6 +252,12 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: wait_timeout = float(getattr(args, "rl_sync_wait_timeout_s", 900.0)) if not 0 < wait_timeout <= 7 * 86400 or wait_timeout != int(wait_timeout): problems.append("--rl-sync-wait-timeout-s in whole seconds within 1..604800") + lead = getattr(args, "rl_sync_max_local_rounds_ahead", None) + if lead is not None: + if not async_rounds: + problems.append("--rl-sync-max-local-rounds-ahead only with --rl-sync-full-parameter-rounds async") + elif not 1 <= lead <= FULL_PARAMETER_MAX_LOCAL_ROUNDS_AHEAD: + problems.append(f"--rl-sync-max-local-rounds-ahead within 1..{FULL_PARAMETER_MAX_LOCAL_ROUNDS_AHEAD}") checkpoint_every = int(getattr(args, "rl_sync_checkpoint_every", 1)) if async_rounds: # The async checkpoint writer and the sync/train overlap are diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 984c8f298..6ed3b90be 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -121,6 +121,8 @@ def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: sweep += " --sweep-overlap" if getattr(args, "rl_sync_push_no_tensor_blocks", False): sweep += " --sweep-push-no-tensor-blocks" + if getattr(args, "rl_sync_max_local_rounds_ahead", None) is not None: + sweep += f" --sweep-max-local-rounds-ahead {int(args.rl_sync_max_local_rounds_ahead)}" if float(getattr(args, "rl_sync_wait_timeout_s", 900.0)) != 900.0: sweep += f" --sweep-wait-timeout-s {float(args.rl_sync_wait_timeout_s):g}" if getattr(args, "rl_sync_fragment_max_bytes", 2 << 30) != 2 << 30: diff --git a/yeto/protocol.py b/yeto/protocol.py index 4447c2a75..f178b5070 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -36,7 +36,7 @@ import time from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from .downlink import HEADER_LEN as DOWNLINK_HEADER_LEN from .downlink import DownlinkAnchors, DownlinkChunkSink @@ -915,6 +915,32 @@ def _connect_log(message: str) -> None: print(f"[yeto-syncer-client] {message}", file=sys.stderr, flush=True) +@dataclass(frozen=True) +class EncodedPush: + """Wire parts of one PUSH_FRAGMENT payload (``SyncerClient.encode_push_parts``).""" + + fragment_id: int + push_base: int + nbytes: int + parts: tuple[bytes, ...] + + +def timed_parts(parts: Iterable[Any], timings: dict[str, float], key: str) -> Iterator[Any]: + """Yield ``parts`` unchanged; add the seconds spent producing each one + (inside the upstream iterator) to ``timings[key]``.""" + iterator = iter(parts) + timings.setdefault(key, 0.0) + while True: + started = time.monotonic() + try: + part = next(iterator) + except StopIteration: + timings[key] += time.monotonic() - started + return + timings[key] += time.monotonic() - started + yield part + + class SyncerClient: """Non-blocking striped syncer connection owned by one learner process. @@ -923,6 +949,9 @@ class SyncerClient: forever); once exhausted, check_health raises. """ + # push_fragment_parts takes ``timings`` (S27 G-AG push time split). + push_timings_supported = True + # Transport defaults (S26 E3b); __init__ overrides them per client. stream_schedule: str = "rr" stream_stall_timeout: float | None = None @@ -1545,9 +1574,20 @@ def push_fragment_parts( *, before_last_enqueue: Callable[[], None] | None = None, expert_tags: Sequence[tuple[int, int, int]] | None = None, + timings: dict[str, float] | None = None, + encoded: EncodedPush | None = None, ) -> bool: """Stream one PUSH_FRAGMENT from bounded contiguous byte parts. + ``timings`` (S27 G-AG diagnostics): when given, ``produce_s`` gets the + seconds spent producing the wire parts (the caller's ``tensor_parts`` + plus the wire encoding) and ``push_s`` the whole call; the rest of + ``push_s`` is framing and enqueueing. The bytes do not change. + + ``encoded`` (S27 G-AG): wire parts :meth:`encode_push_parts` already + produced for this fragment (``tensor_parts`` must then be empty); the + frame is byte-identical to streaming the same deltas. + The syncer observes the same logical inner frame as ``push_fragment``. Only one CHUNK-sized envelope is assembled at a time; neither the complete tensor nor the complete inner frame is materialized here. @@ -1560,6 +1600,50 @@ def push_fragment_parts( raise ValueError("round_attempt must be positive") if c_steps < 1: raise ValueError("c_steps must be positive") + if encoded is None: + expected, tensor_parts = self._push_wire_parts(fragment_id, tensor_parts) + else: + if encoded.fragment_id != fragment_id or encoded.push_base != self._push_base: + raise ValueError("encoded PUSH_FRAGMENT parts belong to another fragment or codec") + expected, tensor_parts = encoded.nbytes, encoded.parts + head = struct.pack( + " tuple[int, Iterable[bytes | bytearray | memoryview]]: + """The wire size and wire parts of one PUSH_FRAGMENT's f32 deltas.""" numel = self.layout.fragments[fragment_id].numel segments = self.push_segments(fragment_id) if self._push_base == DTYPE_Q4: @@ -1577,26 +1661,38 @@ def push_fragment_parts( expected = _tensor_nbytes(push_wire_dtype(self.dtype), numel) if self._push_base == DTYPE_F32_PUSH_BF16: tensor_parts = _f32_parts_to_bf16(tensor_parts) - head = struct.pack( - " bool: + """Whether :meth:`encode_push_parts` may run off the sending thread + and ahead of the push: a deterministic block codec (int8, q4, int4 + without stochastic rounding) whose output parts are fresh bytes.""" + return self._push_base in (DTYPE_F32_PUSH_I8, DTYPE_F32_PUSH_Q4) or ( + self._push_base == DTYPE_F32_PUSH_I4 and self._push_i4_rng is None) + + def encode_push_parts( + self, + fragment_id: int, + tensor_parts: Iterable[bytes | bytearray | memoryview], + *, + timings: dict[str, float] | None = None, + ) -> EncodedPush: + """Encode one fragment's f32 deltas to wire parts now (S27 G-AG). + + Pure: reads only the layout and codec settings, so another thread may + encode the next fragments while this one sends. ``timings`` gets + ``produce_s``, as in :meth:`push_fragment_parts`.""" + if not self.can_encode_ahead: + raise ValueError("this push codec cannot be encoded ahead of its push") + self._check_owned("PUSH_FRAGMENT", fragment_id) + expected, parts = self._push_wire_parts(fragment_id, tensor_parts) + if timings is not None: + parts = timed_parts(parts, timings, "produce_s") + materialized = tuple(bytes(part) if not isinstance(part, bytes) else part for part in parts) + if sum(len(part) for part in materialized) != expected: + raise ValueError(f"PUSH_FRAGMENT fragment {fragment_id} encoded to the wrong size") + return EncodedPush(fragment_id, self._push_base, expected, materialized) def _expert_tag_block( self, fragment_id: int, expert_tags: Sequence[tuple[int, int, int]] | None diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 35b8b66d1..9e1988999 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -33,6 +33,7 @@ FinalManifest, PullRequest, SyncerClient, + timed_parts, ) from yeto.push_error_feedback import RESIDUAL_DTYPES, PushErrorFeedback from yeto.syncer_shards import ShardedSyncerClient @@ -44,7 +45,9 @@ StoredAuthoritativeFragment, defer_capture_supported, optimizer_proof_without_model_check_supported, + timed_group, ) +from yeto.rl.host_info import print_host_report, report_rank_hosts from yeto.rl.engine.policy_digest import ( PolicyDigest, TrainerResidentState, @@ -340,6 +343,16 @@ class RoleFragmentSubmission: pull_to_push_seconds: float +@dataclass(frozen=True) +class PreparedPush: + """One candidate's wire bytes computed ahead (``prepare_push``).""" + + submission_cut: Any + anchor: Any + encoded: Any # yeto.protocol.EncodedPush + timings: dict[str, float] + + @dataclass(frozen=True) class RoleSubmissionCandidate: """One immutable answer to a role syncer's retained PULL permit.""" @@ -426,6 +439,9 @@ def __init__(self, config: RoleStreamConfig, *, ray_module=None) -> None: self.group = None self.adapter: MilesChunkedFullParameterAdapter | None = None self.client = None + # S27 G-AG: summed push time split (_note_push_timing), read by the + # session's round event. + self.push_timing_totals: dict[str, float] = {} self.heartbeat: LeaseHeartbeat | None = None self.sink = None self.current: ReferencedPolicyCut | None = None @@ -546,6 +562,11 @@ def finalizing(self) -> bool: async def initialize(self, group) -> None: if self.group is not None: raise RuntimeError(f"{self.config.algorithm} {self.config.role} stream initialized twice") + # S27 G-AG diagnostics: host facts of the driver and of every trainer + # rank, then one timed line per Miles group call. + print_host_report(f"{self.config.role} driver learner {self.config.learner_id}") + await report_rank_hosts(group, f"{self.config.role} learner {self.config.learner_id}") + group = timed_group(group, self.config.role) adapter, initial = await MilesChunkedFullParameterAdapter.initialize( group, policy_version=0, @@ -1457,15 +1478,45 @@ def retained_replay_candidate( raise RuntimeError("role-stream reconnect replay changed a retained candidate") return replace(candidate, permit=permit) + @property + def encodes_ahead(self) -> bool: + """Whether :meth:`prepare_push` may compute a push's wire bytes ahead + of it and off the event-loop thread (S27 G-AG): no int4 error + feedback (its residual state is committed per push, in push order) + and a deterministic codec (``SyncerClient.can_encode_ahead``).""" + return (self.error_feedback is None and self.client is not None and self.adapter is not None + and bool(getattr(self.client, "can_encode_ahead", False)) + and callable(getattr(self.client, "encode_push_parts", None))) + + def prepare_push(self, candidate: RoleSubmissionCandidate) -> PreparedPush: + """Compute one candidate's delta and wire encoding now (thread-safe: + reads the retained cut and anchor, writes nothing shared). The bytes + equal what :meth:`push_candidate` streams for the same candidate.""" + + if not self.encodes_ahead: + raise RuntimeError("role stream cannot encode a push ahead") + fragment_id = candidate.permit.fragment_id + timings: dict[str, float] = {} + started = time.monotonic() + parts = timed_parts(self.adapter.delta_parts_from_authoritative( + candidate.anchor, candidate.submission_cut, fragment_id, ray_module=self.ray_module, + )(), timings, "delta_s") + encoded = self.client.encode_push_parts(fragment_id, parts, timings=timings) + timings["prepare_s"] = time.monotonic() - started + return PreparedPush(candidate.submission_cut, candidate.anchor, encoded, timings) + def push_candidate( self, candidate: RoleSubmissionCandidate, + prepared: PreparedPush | None = None, ) -> RoleFragmentSubmission | None: """Release one coordinator-approved candidate on its exact permit. A zero-chunk outage returns ``None``. A partial-generation loss is allowed to escape to the paired coordinator, which retains both cuts and waits for this exact logical PULL to be replayed after reconnect. + ``prepared`` (:meth:`prepare_push` of this candidate's cut and anchor) + sends bytes computed ahead instead of streaming them now. """ if ( @@ -1481,21 +1532,29 @@ def push_candidate( if permit != candidate.permit: return None fragment_id = candidate.permit.fragment_id - parts = self.adapter.delta_parts_from_authoritative( + if prepared is not None: + return self._push_prepared(candidate, prepared) + # S27 G-AG: split this push's time (the parts are yielded unchanged). + timings: dict[str, float] = {} + started = time.monotonic() + parts = timed_parts(self.adapter.delta_parts_from_authoritative( candidate.anchor, candidate.submission_cut, fragment_id, ray_module=self.ray_module, - )() + )(), timings, "delta_s") if self.error_feedback is not None: - parts = self.error_feedback.parts( + parts = timed_parts(self.error_feedback.parts( fragment_id, candidate.permit.global_step, self.config.expected_fragments, self.adapter.layout.fragments.fragments[fragment_id].numel, parts, self.client.push_segments(fragment_id), - ) + ), timings, "feedback_incl_s") + extra: dict[str, Any] = {} if candidate.expert_tags is None else {"expert_tags": candidate.expert_tags} + if getattr(self.client, "push_timings_supported", False): + extra["timings"] = timings try: queued = self.client.push_fragment_parts( fragment_id, @@ -1506,12 +1565,14 @@ def push_candidate( candidate.c_steps, candidate.c_tokens, parts, - **({} if candidate.expert_tags is None else {"expert_tags": candidate.expert_tags}), + **extra, ) except BaseException: if self.error_feedback is not None: self.error_feedback.discard(fragment_id) raise + finally: + self._note_push_timing(candidate, timings, time.monotonic() - started) if self.error_feedback is not None: if queued is True: self.error_feedback.commit(fragment_id, candidate.permit.global_step) @@ -1519,6 +1580,37 @@ def push_candidate( self.error_feedback.discard(fragment_id) if queued is not True: return None + return self._pushed(candidate) + + def _push_prepared(self, candidate: RoleSubmissionCandidate, + prepared: PreparedPush) -> RoleFragmentSubmission | None: + if (prepared.submission_cut is not candidate.submission_cut or prepared.anchor is not candidate.anchor + or prepared.encoded.fragment_id != candidate.permit.fragment_id): + raise RuntimeError("prepared push belongs to another candidate") + extra: dict[str, Any] = {} if candidate.expert_tags is None else {"expert_tags": candidate.expert_tags} + timings = dict(prepared.timings) + started = time.monotonic() + try: + queued = self.client.push_fragment_parts( + candidate.permit.fragment_id, + candidate.permit.global_step, + candidate.permit.round_attempt, + candidate.base_version, + candidate.local_step, + candidate.c_steps, + candidate.c_tokens, + (), + encoded=prepared.encoded, + **extra, + ) + finally: + timings["push_s"] = time.monotonic() - started + self._note_push_timing(candidate, timings, timings["push_s"], prepared=True) + if queued is not True: + return None + return self._pushed(candidate) + + def _pushed(self, candidate: RoleSubmissionCandidate) -> RoleFragmentSubmission: self._answered[ (candidate.permit.global_step, candidate.permit.round_attempt) ] = self._link_generation(candidate.permit.fragment_id) @@ -1537,6 +1629,38 @@ def push_candidate( max(0.0, time.monotonic() - candidate.permit.received_at), ) + def _note_push_timing(self, candidate: RoleSubmissionCandidate, timings: dict[str, float], + total: float, *, prepared: bool = False) -> None: + """Print one push's time split and add it to :attr:`push_timing_totals`. + + ``delta_s``: the delta from the cut and the anchor (Ray get, checks, + subtraction); ``feedback_s``: int4 error feedback; ``encode_s``: the + wire encoding; ``enqueue_s``: framing and enqueueing (it waits only + when the send queue is full; the upload itself is asynchronous); + ``total_s``: time on the pushing thread. A prepared push + (``prepared=1``) computed delta and encoding ahead on a worker + thread (``prepare_s``), so its ``total_s`` is the enqueue only.""" + + delta = timings.get("delta_s", 0.0) + feedback = max(0.0, timings.get("feedback_incl_s", delta) - delta) + produced = timings.get("produce_s") + encode = None if produced is None else max(0.0, produced - delta - feedback) + if prepared: + enqueue = timings.get("push_s", total) + else: + enqueue = None if produced is None else max(0.0, timings.get("push_s", total) - produced) + split = {"delta_s": delta, "feedback_s": feedback, "encode_s": encode, "enqueue_s": enqueue, + "prepare_s": timings.get("prepare_s") if prepared else None, "total_s": total} + totals = self.push_timing_totals + totals["count"] = totals.get("count", 0) + 1 + for name, value in split.items(): + if value is not None: + totals[name] = totals.get(name, 0.0) + value + text = " ".join(f"{name}={value:.3f}" for name, value in split.items() if value is not None) + print(f"[yeto-full-parameter] push role={self.config.role} fragment={candidate.permit.fragment_id} " + f"step={candidate.permit.global_step} c_steps={candidate.c_steps} prepared={int(prepared)} {text}", + file=sys.stderr, flush=True) + def _queue_pulls(self) -> None: # A redial drops the permits and answers of that link only: the # syncer resends PULLs of open rounds on the new group. With a diff --git a/yeto/rl/adapters/miles/full_parameter/session.py b/yeto/rl/adapters/miles/full_parameter/session.py index d333be204..edd89e1f1 100644 --- a/yeto/rl/adapters/miles/full_parameter/session.py +++ b/yeto/rl/adapters/miles/full_parameter/session.py @@ -37,10 +37,16 @@ class FullParameterSyncConfig: roles: tuple[RoleStreamConfig, ...] sweep: bool = False + # S27 G-AG lead cap (``--sweep-max-local-rounds-ahead``); None = off. + # Island-local: not part of the session contract or the syncer profile. + max_local_rounds_ahead: int | None = None def __post_init__(self) -> None: - names = tuple(config.role for config in self.roles) - FullParameterSyncProfile.from_configs(dict(zip(names, self.roles, strict=True)), sweep=self.sweep) + self.profile() + + def profile(self) -> FullParameterSyncProfile: + return FullParameterSyncProfile.from_configs( + self.by_role(), sweep=self.sweep, max_local_rounds_ahead=self.max_local_rounds_ahead) def by_role(self) -> dict[str, RoleStreamConfig]: return {config.role: config for config in self.roles} @@ -69,7 +75,7 @@ def build_full_parameter_sync( raise TypeError("Miles full-parameter sync configuration is missing") configs = config.by_role() streams = {role: MilesFullParameterRoleStream(stream) for role, stream in configs.items()} - pairing = RolePairing(streams, FullParameterSyncProfile.from_configs(configs, sweep=config.sweep)) + pairing = RolePairing(streams, config.profile()) return FullParameterSync( pairing, groups=groups, diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index dd4586520..ebc6a2b45 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -21,6 +21,7 @@ from yeto.rl.engine.resume import DEFAULT_CUT_EVERY from yeto.rl.engine.run_flags import ( FULL_PARAMETER_ASYNC_MAX_PIPELINE, + FULL_PARAMETER_MAX_LOCAL_ROUNDS_AHEAD as MAX_LOCAL_ROUNDS_AHEAD, ISLAND, add_rl_flags, rl_option_pairs, @@ -954,7 +955,21 @@ def _dense_full_sync_config( # an island pool sets the epoch (and lease); else peer recovery may **{"syncer_epoch": _island_syncer_epoch(args), **_island_pool_membership(args)}, ) - return FullParameterSyncConfig((actor,), sweep=not async_rounds) + return FullParameterSyncConfig((actor,), sweep=not async_rounds, + max_local_rounds_ahead=_max_local_rounds_ahead(args, async_rounds)) + + +def _max_local_rounds_ahead(args, async_rounds: bool) -> int | None: + """``--sweep-max-local-rounds-ahead`` (launcher ``--rl-sync-max-local-rounds-ahead``): + the S27 G-AG lead cap of the asynchronous rounds; None (off) by default.""" + value = getattr(args, "sweep_max_local_rounds_ahead", None) + if value is None: + return None + if not async_rounds: + raise ValueError("--sweep-max-local-rounds-ahead needs --full-parameter-rounds async") + if type(value) is not int or not 1 <= value <= MAX_LOCAL_ROUNDS_AHEAD: + raise ValueError(f"--sweep-max-local-rounds-ahead must be an integer within 1..{MAX_LOCAL_ROUNDS_AHEAD}") + return value def _sweep_wait_timeout_s(args) -> float: diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index 2cd701261..50ad5aeb0 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -1257,3 +1257,86 @@ def _metadata_only_reference(reference: ReferencedPolicyCut) -> ReferencedPolicy content_hash=reference.content_hash, transport_cut=metadata, ) + + +# -- S27 G-AG diagnostics: one line per Miles trainer-group call --------------- + +OP_TIMING_ENV = "YETO_FULL_PARAMETER_OP_TIMING" + + +class TimedGroup: + """Forward every attribute of a Miles trainer group; log each method call. + + The production group is a ``RayWorkerHandle`` around the TrainerController + actor, whose ``run_plugin`` log lines carry no operation name. Each call + through this wrapper prints ``[yeto-full-parameter] miles op= + role= elapsed_s= ok=`` on the driver once it returns, so + the controller's ``run_plugin`` lines of that interval belong to it (one + group method may run several plugins, e.g. apply: topology, prepare, + commit, finalize). ``functools.wraps`` keeps each method's signature, so + the feature checks above see the wrapped method unchanged.""" + + def __init__(self, group: Any, label: str) -> None: + object.__setattr__(self, "_timed_group", group) + object.__setattr__(self, "_timed_label", label) + + def __getattr__(self, name: str) -> Any: + value = getattr(self._timed_group, name) + if name.startswith("_") or not callable(value): + return value + import functools + import inspect + import time + + label = self._timed_label + + @functools.wraps(value) + def call(*args: Any, **kwargs: Any) -> Any: + started = time.monotonic() + try: + result = value(*args, **kwargs) + except BaseException: + _log_group_call(name, label, started, ok=False) + raise + if not inspect.isawaitable(result): + _log_group_call(name, label, started, ok=True) + return result + return _timed_await(result, name, label, started) + + return call + + def __setattr__(self, name: str, value: Any) -> None: + setattr(self._timed_group, name, value) + + def __repr__(self) -> str: + return f"TimedGroup({self._timed_group!r})" + + +async def _timed_await(awaitable: Any, name: str, label: str, started: float) -> Any: + try: + result = await awaitable + except BaseException: + _log_group_call(name, label, started, ok=False) + raise + _log_group_call(name, label, started, ok=True) + return result + + +def _log_group_call(name: str, label: str, started: float, *, ok: bool) -> None: + import sys + import time + + print(f"[yeto-full-parameter] miles op={name} role={label} " + f"elapsed_s={time.monotonic() - started:.3f} ok={'true' if ok else 'false'}", + file=sys.stderr, flush=True) + + +def timed_group(group: Any, label: str) -> Any: + """``group`` wrapped in :class:`TimedGroup` (``YETO_FULL_PARAMETER_OP_TIMING=0``: + unchanged).""" + + import os + + if group is None or isinstance(group, TimedGroup) or os.environ.get(OP_TIMING_ENV, "1") == "0": + return group + return TimedGroup(group, label) diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 4f93307fb..a606efb07 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -357,6 +357,13 @@ help="full-parameter: f32 byte bound of one sync fragment; a larger " "parameter spans several fragments as dimension-0 row slices", ), + island_flag( + "--sweep-max-local-rounds-ahead", + type=int, + default=None, + help="full-parameter async rounds: lead cap K >= 1, wait at a boundary instead of " + "training a K-th local round on top of an unmerged push (default off)", + ), island_flag( "--sweep-wait-timeout-s", type=float, diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index ba139fcb8..9835ddb72 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -59,10 +59,11 @@ import asyncio import os +import sys import time from collections.abc import Mapping, Sequence from dataclasses import dataclass, field, replace -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol, Self from yeto.protocol import PartialMessageGenerationLost from yeto.rl.pool_client import LeaveReason @@ -167,6 +168,10 @@ class FullParameterSyncProfile: # S27 P2: the syncer closes rounds with quorum K <= islands, so a fragment # broadcast may arrive without this island's submission (overwrite it). partial_quorum: bool = False + # S27 G-AG lead cap K (streaming only; None = off, the old behaviour): an + # island trains at most K - 1 local rounds on top of an unmerged push of + # its own; K = 1 starts no round while one is unmerged (RolePairing.wait_for_lead). + max_local_rounds_ahead: int | None = None def __post_init__(self) -> None: if not self.roles or len(set(self.roles)) != len(self.roles): @@ -179,6 +184,11 @@ def __post_init__(self) -> None: raise ValueError("sync/train overlap needs the dense strict sweep") if self.partial_quorum and self.sweep: raise ValueError("a partial quorum needs the streaming profile, not the strict sweep") + if self.max_local_rounds_ahead is not None: + if type(self.max_local_rounds_ahead) is not int or self.max_local_rounds_ahead < 1: + raise ValueError("max_local_rounds_ahead must be an integer >= 1") + if self.sweep or self.overlap: + raise ValueError("a local-round lead cap needs the streaming profile (no sweep, no overlap)") @property def total_rounds(self) -> int: @@ -193,7 +203,8 @@ def effective_pipeline_depth(self) -> int: @classmethod def from_configs( - cls, configs: Mapping[str, Any], *, sweep: bool = False + cls, configs: Mapping[str, Any], *, sweep: bool = False, + max_local_rounds_ahead: int | None = None, ) -> FullParameterSyncProfile: """Profile of role-stream configs that pass :func:`require_lockstep`.""" @@ -210,6 +221,7 @@ def from_configs( sweep=sweep, overlap=bool(getattr(first, "overlap", False)), partial_quorum=bool(getattr(first, "partial_quorum", False)), + max_local_rounds_ahead=max_local_rounds_ahead, ) @@ -233,6 +245,9 @@ class PairedSubmission: candidates: dict[str, Any] submissions: dict[str, Any] = field(default_factory=dict) retained_at: float = field(default_factory=time.monotonic) + # The island's local round (RolePairing.local_round) when it retained the + # pair: the round whose cut the pair pushes (lead cap). + local_round: int = 0 @property def logical_key(self) -> LogicalKey: @@ -280,6 +295,12 @@ def __init__(self, streams: Mapping[str, RoleStream], profile: FullParameterSync self.caught_up_fragments = 0 # catch-up jumps (one per fragment and jump) # Terminal phase: last rounds closed without a trainer apply (_retires). self.retired_rounds = 0 + # Lead cap: local rounds this island completed (set by the session at + # each boundary) and what its waits did. + self.local_round = 0 + self.lead_waits = 0 + self.lead_wait_seconds = 0.0 + self.lead_wait_escapes = 0 # -- fail-stop --------------------------------------------------------- @@ -358,11 +379,14 @@ def submit_ready(self) -> dict[str, RoleProgress]: # A round whose broadcast is staged or applied closed without # this island: pushing would only upload a refused delta. shared = [key for key in shared if not self._round_closed(key[1], key[0])] - for key in shared[: max(0, capacity)]: - pair = self._new_pair({role: by_role[role].pop(key) for role in self.streams}) - self._retain(pair) - self.inflight[key] = pair - self._push_once(pair, dict(pair.candidates), produced) + planned = [(key, {role: by_role[role].pop(key) for role in self.streams}) + for key in shared[: max(0, capacity)]] + with _PushPrefetch(self.streams, planned) as prefetch: + for key, candidates in planned: + pair = self._new_pair(candidates) + self._retain(pair) + self.inflight[key] = pair + self._push_once(pair, dict(pair.candidates), produced, prepared=prefetch.take(key)) self._check_timeouts(by_role, now) except BaseException as error: self.fail_stop(error) @@ -435,7 +459,7 @@ def _new_pair(self, candidates: dict[str, Any]) -> PairedSubmission: or candidate.c_steps < self.profile.local_horizon ): raise RuntimeError("role retained cuts are not one H-qualified round") - return PairedSubmission(candidates) + return PairedSubmission(candidates, local_round=self.local_round) def _retain(self, pair: PairedSubmission) -> None: retained: list[str] = [] @@ -457,6 +481,7 @@ def _push_once( pair: PairedSubmission, observed: Mapping[str, Any | None], produced: dict[str, list[Any]], + prepared: Mapping[str, Any] | None = None, ) -> None: for role, stream in self.streams.items(): replayed = _retained_replay(pair.candidates[role], observed.get(role)) @@ -466,8 +491,10 @@ def _push_once( # A PULL visible after a successful enqueue proves that generation # did not commit it: replay the exact retained bytes. pair.submissions.pop(role, None) + ahead = None if prepared is None else prepared.get(role) try: - submission = stream.push_candidate(replayed) + submission = (stream.push_candidate(replayed) if ahead is None + else stream.push_candidate(replayed, ahead)) except PartialMessageGenerationLost: submission = None if submission is not None: @@ -724,6 +751,79 @@ def _journal(self, pair: PairedSubmission) -> None: self.journal.append(dict(pair.submissions)) self._release(pair) + # -- local-round lead cap (S27 G-AG) ----------------------------------------- + + def lead_blockers(self) -> list[PairedSubmission]: + """In-flight pairs that stop the next local round under the lead cap K. + + At a boundary after local round ``r`` the next round is ``r + 1``. A + pair retained at local round ``p`` (its cut is round ``p``) blocks when + ``p <= r + 1 - K``: training round ``r + 1`` would put ``K`` or more + rounds on top of that unmerged push. With a fixed roster a merge needs + every island's push, so with K = 1 an island starts a round only after + every island pushed everything it pushed: it stays at most one local + round ahead of the slowest island. A larger K allows K - 1 rounds on + top of an unmerged push; that bounds the island's c_steps per push, + not its absolute lead (pushes made one boundary late can repeat).""" + + limit = self.profile.max_local_rounds_ahead + if limit is None: + return [] + return [pair for pair in self.inflight.values() if pair.local_round <= self.local_round + 1 - limit] + + async def wait_for_lead(self, serve: Any = None) -> float: + """Lead cap: block at this boundary until :meth:`lead_blockers` is empty. + + The wait keeps the island's sync duties: every poll materializes and + pushes PULLs it has candidates for (from this boundary's cut), applies + arrived broadcasts, and runs ``serve`` (peer snapshots). It stops at + once when the terminal phase starts (``any_finalizing``): the terminal + barrier then waits instead. + + No new timer: each pair's wait timeout keeps running inside + :meth:`submit_ready`, so a broadcast that never comes still fails + there. The lead wait gives up (and the island trains on, the old + behaviour) once the oldest blocking pair has waited half of + ``wait_timeout`` since its last push, before that pair could time out. + This also breaks a theoretical deadlock where each island waits for a + push the other can make only after training (possible only when two + broadcasts reach the two islands in different orders). + + Returns the seconds this boundary waited.""" + + blockers = self.lead_blockers() + if not blockers or self.any_finalizing or self.failure is not None: + return 0.0 + started = time.monotonic() + poll = max(self.profile.poll_seconds, LEAD_WAIT_POLL_S) + escaped = False + oldest = min(pair.local_round for pair in blockers) + print(f"[yeto-full-parameter] lead wait: local round {self.local_round}, " + f"{len(blockers)} pair(s) pushed at local round {oldest}.. are unmerged " + f"(lead cap {self.profile.max_local_rounds_ahead})", file=sys.stderr, flush=True) + while True: + await asyncio.sleep(poll) + await self.prepare_submissions() + self.submit_ready() + await self.commit_staged_broadcasts() + if serve is not None: + await serve() + blockers = self.lead_blockers() + if not blockers or self.any_finalizing: + break + give_up = min(pair.retained_at for pair in blockers) + self.profile.wait_timeout / 2 + if time.monotonic() >= give_up: + escaped = True + break + waited = time.monotonic() - started + self.lead_waits += 1 + self.lead_wait_seconds += waited + self.lead_wait_escapes += int(escaped) + print(f"[yeto-full-parameter] lead wait: local round {self.local_round} waited {waited:.1f}s" + + (" and gave up at half the wait timeout (trains on)" if escaped else ""), + file=sys.stderr, flush=True) + return waited + # -- terminal barrier ------------------------------------------------------ async def finish_without_more_training(self) -> dict[str, RoleProgress]: @@ -786,6 +886,90 @@ def _candidate_map(candidates: Sequence[Any], *, role: str) -> dict[LogicalKey, BATCH_APPLY_ENV = "YETO_FULL_PARAMETER_BATCH_APPLY" +# S27 G-AG: threads that compute the next pushes' deltas and wire encodings +# while the current one is enqueued (1 = compute each push in turn, as before). +PUSH_WORKERS_ENV = "YETO_FULL_PARAMETER_PUSH_WORKERS" +DEFAULT_PUSH_WORKERS = 3 +_PUSH_POOL: list[Any] = [] # (workers, ThreadPoolExecutor), created on first use + + +def push_workers() -> int: + try: + return max(1, int(os.environ.get(PUSH_WORKERS_ENV, DEFAULT_PUSH_WORKERS))) + except ValueError: + return 1 + + +def _push_pool(workers: int) -> Any: + if not _PUSH_POOL or _PUSH_POOL[0][0] != workers: + from concurrent.futures import ThreadPoolExecutor + + _PUSH_POOL[:] = [(workers, ThreadPoolExecutor(max_workers=workers, + thread_name_prefix="yeto-push-prepare"))] + return _PUSH_POOL[0][1] + + +class _PushPrefetch: + """Compute the deltas and wire encodings of the next planned pushes on + worker threads (S27 G-AG) while the pushing thread enqueues them in plan + order. + + Only when every role can (``encodes_ahead``) and + ``YETO_FULL_PARAMETER_PUSH_WORKERS`` > 1; otherwise :meth:`take` returns + None and each push computes its bytes itself, as before. The bytes are + the same either way: a prepared push reads the same retained cut and + anchor, and the codec is deterministic (no stochastic rounding, no error + feedback). At most ``workers`` pushes are computed ahead; on exit every + unfinished computation is cancelled or awaited, so no worker outlives the + submit call.""" + + def __init__(self, streams: Mapping[str, Any], planned: Sequence[tuple[LogicalKey, Mapping[str, Any]]]): + self.streams = streams + self.workers = push_workers() + self.enabled = len(planned) > 1 and self.workers > 1 and all( + getattr(stream, "encodes_ahead", False) and callable(getattr(stream, "prepare_push", None)) + for stream in streams.values()) + self.queue = list(planned) + self.futures: dict[LogicalKey, dict[str, Any]] = {} + + def __enter__(self) -> Self: + if self.enabled: + self._fill() + return self + + def _fill(self) -> None: + pool = _push_pool(self.workers) + while self.queue and len(self.futures) < self.workers: + key, candidates = self.queue.pop(0) + self.futures[key] = {role: pool.submit(self.streams[role].prepare_push, candidate) + for role, candidate in candidates.items()} + + def take(self, key: LogicalKey) -> dict[str, Any] | None: + if not self.enabled: + return None + futures = self.futures.pop(key, None) + if futures is None: + return None + try: + return {role: future.result() for role, future in futures.items()} + finally: + self._fill() + + def __exit__(self, *exc: object) -> None: + for futures in self.futures.values(): + for future in futures.values(): + future.cancel() + for futures in self.futures.values(): + for future in futures.values(): + if not future.cancelled(): + try: + future.result() + except BaseException: # noqa: BLE001, S110 - the push that needed it already failed + pass + self.futures.clear() + self.queue.clear() +# Lead cap: poll interval of the boundary wait (the wait lasts a peer's round). +LEAD_WAIT_POLL_S = 0.05 @dataclass(frozen=True) @@ -859,6 +1043,8 @@ def __init__( self._upload_before: tuple | None = None # Partial quorum: overwritten rounds at the previous boundary. self._overwritten_before = 0 + self._lead_wait_s = 0.0 # lead cap: this boundary's wait + self._push_timing_before: dict[str, float] | None = None @property def outer_version(self) -> int | None: @@ -895,6 +1081,8 @@ def _publish_outer_version(self, driver: IslandDriver) -> None: def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> SyncBoundary: del stats self._timing = SyncTiming(None, None) + self._lead_wait_s = 0.0 + self.pairing.local_round = rollout_id + 1 self.terminal = self.run(self._boundary(rollout_id)) versions = self.pairing.fragment_versions() state = self.actor.resident_state(rollout_id + 1) @@ -908,6 +1096,8 @@ def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> Sync "sync/global_policy_hash": state.policy_tensor_hash(), "sync/inflight_pairs": len(self.pairing.inflight), **(self._partial_quorum_fields(versions) if self.pairing.profile.partial_quorum else {}), + **self._lead_fields(), + **self._push_timing_fields(), **self._timing.fields(), **self._upload_fields(), }, @@ -927,6 +1117,28 @@ def _partial_quorum_fields(self, versions: tuple[int, ...]) -> dict[str, int]: "sync/outer_version": merged_version(versions), } + def _lead_fields(self) -> dict[str, object]: + pairing = self.pairing + if pairing.profile.max_local_rounds_ahead is None: + return {} + return { + "sync/max_local_rounds_ahead": pairing.profile.max_local_rounds_ahead, + "sync/lead_wait_s": round(self._lead_wait_s, 3), + "sync/lead_waits": pairing.lead_waits, + "sync/lead_wait_escapes": pairing.lead_wait_escapes, + } + + def _push_timing_fields(self) -> dict[str, object]: + """The actor's push time split of this boundary (S27 G-AG): delta + computation, error feedback, wire encoding and enqueue, summed over + the fragments it pushed; streams without the counters add nothing.""" + totals = getattr(self.actor, "push_timing_totals", None) + if not isinstance(totals, dict): + return {} + before, self._push_timing_before = self._push_timing_before or {}, dict(totals) + return {f"sync/push_{name}": round(value - before.get(name, 0.0), 3) if name != "count" + else int(value - before.get(name, 0)) for name, value in totals.items()} + def _stream_snapshots(self) -> tuple | None: snapshots = getattr(getattr(self.actor, "client", None), "stream_snapshots", None) return snapshots() if callable(snapshots) else None @@ -965,6 +1177,8 @@ async def _boundary(self, rollout_id: int) -> bool: submit_ts = unix_now() await pairing.commit_staged_broadcasts() await self._serve_peer_snapshots() + if pairing.profile.max_local_rounds_ahead is not None: + self._lead_wait_s = await pairing.wait_for_lead(self._serve_peer_snapshots) if pairing.profile.sweep: await self._wait_for_sweep(before) # Broadcasts are staged in the background, so the receive time diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index dff8b850e..4557f3f83 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -58,6 +58,9 @@ # The island's wait for a paired broadcast, the terminal barrier and the FINAL # cut (--rl-sync-wait-timeout-s / island --sweep-wait-timeout-s). FULL_PARAMETER_DEFAULT_WAIT_TIMEOUT_S = 900.0 +# S27 G-AG lead cap of the asynchronous rounds (--rl-sync-max-local-rounds-ahead +# / island --sweep-max-local-rounds-ahead): the largest accepted K. +FULL_PARAMETER_MAX_LOCAL_ROUNDS_AHEAD = 1000 ISLAND = "island" LAYERS = (CLI, ISLAND) @@ -1003,6 +1006,18 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "dimension-0 row slices, each merged as its own tensor (RDA per slice). " "A smaller bound changes the layout hash, so the probe must use the same value", ), + cli_flag( + "--rl-sync-max-local-rounds-ahead", + type=int, + default=None, + help="full-parameter async rounds only: lead cap K >= 1 (default off). At a " + "boundary an island waits, still pushing and applying, instead of training a " + "K-th local round on top of an unmerged push of its own. K=1 starts no round " + "while one of its pushes is unmerged: with a fixed roster a fast island stays " + "at most one local round ahead of the slowest island. The wait gives up (and " + "trains on) after half of --rl-sync-wait-timeout-s since that push. Island-local: " + "the layout, syncer profile and training contract hashes do not change", + ), cli_flag( "--rl-sync-wait-timeout-s", type=float, diff --git a/yeto/rl/host_info.py b/yeto/rl/host_info.py new file mode 100644 index 000000000..ead5b7ecb --- /dev/null +++ b/yeto/rl/host_info.py @@ -0,0 +1,184 @@ +"""Host facts an island prints at start (S27 G-AG diagnostics). + +Two islands with the same image and the same requested CPU and memory ran +memory-heavy boundary work 2-4x apart (s27-gr-g2), and no log said on which +hardware. :func:`host_report` collects the facts that decide that speed: +CPU model and count, the CPUs this process may use, cgroup CPU and memory +limits, host memory, NUMA nodes, and the thread settings of OpenMP, MKL, +OpenBLAS and torch. It reads only ``/proc``, ``/sys`` and the environment +(never a credential) and never raises: a fact it cannot read is ``None``. + +:func:`rank_host_report` is the same report as a trainer-rank plugin +(``TrainGroup.run_plugin``), because the full-parameter capture, apply and +materialize run in the rank processes, whose thread settings Ray may set +differently from the driver's. +""" + +from __future__ import annotations + +import json +import os +import platform +import socket +import sys +from pathlib import Path +from typing import Any + +THREAD_ENV = ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", + "NUMEXPR_NUM_THREADS", "MILES_FULL_PARAMETER_HASH_THREADS") +RANK_HOST_REPORT = f"{__name__}.rank_host_report" +RANK_HOST_REPORT_ENV = "YETO_RANK_HOST_REPORT" + + +def _read(path: str) -> str | None: + try: + return Path(path).read_text().strip() + except (OSError, UnicodeDecodeError): + return None + + +def _cpu_model() -> str | None: + text = _read("/proc/cpuinfo") or "" + for line in text.splitlines(): + key, _, value = line.partition(":") + if key.strip() in ("model name", "Model", "cpu model"): + return value.strip() + return platform.processor() or None + + +def _meminfo_gib() -> dict[str, float | None]: + values: dict[str, float | None] = {"mem_total_gib": None, "mem_available_gib": None} + for line in (_read("/proc/meminfo") or "").splitlines(): + key, _, rest = line.partition(":") + name = {"MemTotal": "mem_total_gib", "MemAvailable": "mem_available_gib"}.get(key) + if name is not None: + try: + values[name] = round(int(rest.split()[0]) / (1 << 20), 1) # kB -> GiB + except (IndexError, ValueError): + pass + return values + + +def _cgroup_cpu_limit() -> float | None: + """CPUs the cgroup quota allows (v2 ``cpu.max``, else v1 CFS); None = no limit.""" + text = _read("/sys/fs/cgroup/cpu.max") + if text: + quota, _, period = text.partition(" ") + if quota != "max": + try: + return round(int(quota) / int(period or 100000), 2) + except ValueError: + return None + return None + quota, period = _read("/sys/fs/cgroup/cpu/cpu.cfs_quota_us"), _read("/sys/fs/cgroup/cpu/cpu.cfs_period_us") + try: + if quota is not None and period is not None and int(quota) > 0: + return round(int(quota) / int(period), 2) + except ValueError: + pass + return None + + +def _cgroup_memory_limit_gib() -> float | None: + for path in ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes"): + text = _read(path) + if text is None: + continue + if text == "max": + return None + try: + value = int(text) + except ValueError: + return None + return None if value >= 1 << 60 else round(value / (1 << 30), 1) + return None + + +def _numa_nodes() -> int | None: + try: + return len([p for p in Path("/sys/devices/system/node").glob("node[0-9]*")]) or None + except OSError: + return None + + +def _torch_threads() -> dict[str, int | None]: + """torch's intra-op and inter-op threads, only when torch is already imported + (this report never imports torch itself).""" + torch = sys.modules.get("torch") + if torch is None: + return {"torch_threads": None, "torch_interop_threads": None} + out: dict[str, int | None] = {} + for key, getter in (("torch_threads", "get_num_threads"), ("torch_interop_threads", "get_num_interop_threads")): + try: + out[key] = int(getattr(torch, getter)()) + except Exception: # noqa: BLE001 - a report never raises + out[key] = None + return out + + +def host_report() -> dict[str, Any]: + """Facts of this host and process; every value is JSON-serializable.""" + report: dict[str, Any] = {} + for key, read in ( + ("hostname", socket.gethostname), + ("kernel", platform.release), + ("cpu_model", _cpu_model), + ("cpu_count", os.cpu_count), + ("cpus_usable", lambda: len(os.sched_getaffinity(0))), + ("cgroup_cpu_limit", _cgroup_cpu_limit), + ("cgroup_memory_limit_gib", _cgroup_memory_limit_gib), + ("numa_nodes", _numa_nodes), + ("pid", os.getpid), + ): + try: + report[key] = read() + except Exception: # noqa: BLE001 - a report never raises + report[key] = None + try: + report.update(_meminfo_gib()) + except Exception: # noqa: BLE001, S110 - a report never raises + pass + for name in THREAD_ENV: + report[name] = os.environ.get(name) + try: + report.update(_torch_threads()) + except Exception: # noqa: BLE001, S110 - a report never raises + pass + return report + + +def print_host_report(label: str) -> dict[str, Any]: + """Print one ``[yeto-host]