diff --git a/examples/experimental/openenv/README.md b/examples/experimental/openenv/README.md new file mode 100644 index 0000000000..f265cd482e --- /dev/null +++ b/examples/experimental/openenv/README.md @@ -0,0 +1,90 @@ +# OpenEnv Terminal-Bench-2 GRPO (GLM-4.7-Flash, single node) + +Train GLM-4.7-Flash with GRPO on the HuggingFace [OpenEnv](https://github.com/huggingface/openenv) +**Terminal-Bench-2 (tbench2)** environment. A miles-side adapter runs the multi-turn +agentic loop (`reset(task_id)` → { policy emits one shell command → `step(exec)` → +feed output back } → `evaluate`) against an unmodified OpenEnv env server; the reward +is the binary pytest result (1.0 = all tests pass, else 0.0). + +This guide targets a **single H200 node with 8 GPUs**. The run is colocated +(training + rollout on the same 8 GPUs): TP=4, EP=2, one SGLang engine per GPU. + +## Prerequisites + +- The node has Docker available (the env server launches one container per task). +- miles is installed and GLM-4.7-Flash weights are reachable (the launcher pulls + `zai-org/GLM-4.7-Flash` from HF and converts it to `torch_dist` on first run). +- Install the OpenEnv tbench2 env client (isolate it if its deps clash with the + miles image): + + ```bash + pip install -e /envs/tbench2_env + ``` + +## 1. Build the prompt data + +Clone the TB2 suite and emit one prompt row per `task_id`: + +```bash +git clone --depth 1 https://github.com/laude-institute/terminal-bench-2.git /workspace/terminal-bench-2 +python make_tbench2_data.py --tasks_dir /workspace/terminal-bench-2 --output /root/tbench2_train.jsonl +# add --n 8 for a small smoke subset +``` + +## 2. Start the env server + +Run it in a separate shell (or off-node — see note). Docker mode gives real TB2 +fidelity; it needs the Docker socket and pulls the per-task images on first use: + +```bash +# Raise the open-file limit first (see Notes): the WebSocket env server holds an +# FD per live session + Docker connection and leaks sockets on unclean +# disconnects, so the default 1024 soft limit is exhausted on a long run. +ulimit -n 1048576 +TB2_MODE=docker TB2_TASKS_DIR=/workspace/terminal-bench-2 MAX_CONCURRENT_ENVS=32 \ + python -m tbench2_env.server.app --port 8003 +``` + +`MAX_CONCURRENT_ENVS` caps live sandboxes; keep it at or below the rollout batch +concurrency. Per-task containers are heavy on disk — if you'd rather not colocate +them with the GPU workload, run the env server on a separate Docker host and point +the launcher at it via `--openenv-env-url http://:8003`. + +## 3. Launch training + +```bash +python run-openenv-tbench2.py --openenv-env-url http://localhost:8003 +``` + +Common overrides: + +| Flag / env var | Default | Purpose | +| --- | --- | --- | +| `--openenv-env-url` | `http://localhost:8003` | Env server URL | +| `--prompt-data` | `/root/tbench2_train.jsonl` | Prompt set from step 1 | +| `--num-rollout` | (launcher) | Number of GRPO steps | +| `OPENENV_MAX_TURNS` | `30` | Max agent turns per episode | +| `OPENENV_MAX_ROLLOUT_TIME_SECONDS` | `3600` | Per-episode wall-clock cap; a straggler that exceeds it is terminated and scored 0 | +| `--dump-details ` | off | Dump per-episode tokens/logprobs/masks/reward for inspection | +| `WANDB_KEY`, `--wandb-project`, `--wandb-team` | — | W&B logging | + +## Notes + +- **Reward signal.** The binary sparse reward needs a task subset where the base + policy *sometimes* succeeds (advantage variance). On the full TB2 suite, + GLM-4.7-Flash's low base solve-rate yields a near-flat GRPO signal — use a + variance-band subset (or a stronger base) to see a learning climb. +- **`_step` vs. rollout.** W&B `_step` is an internal log-call index that advances + several times per rollout; it is **not** the training step. Read the driver log's + `rollout N:` counter for true progress. +- **Sandbox leakage.** Upstream OpenEnv creates task containers with `remove=False` + and only tears them down on a clean session close (the idle reaper is off by + default), so an unclean disconnect (trainer crash) can orphan containers. Sweep + stale TB2 containers between runs, e.g. `docker rm -f` of any older than the + episode wall-cap. +- **Open-file limit.** The same unclean disconnects also leak socket FDs in the + env server process. On a long run under the default 1024 soft limit the accept + loop eventually fails every connection with `OSError: [Errno 24] Too many open + files`, silently throttling rollouts. Start the server with a raised limit + (`ulimit -n 1048576`, as in step 2); if a running server is already saturated, + restart it with the higher limit. diff --git a/examples/experimental/openenv/make_tbench2_data.py b/examples/experimental/openenv/make_tbench2_data.py new file mode 100644 index 0000000000..48d47da069 --- /dev/null +++ b/examples/experimental/openenv/make_tbench2_data.py @@ -0,0 +1,78 @@ +"""Generate a prompt dataset for the OpenEnv Terminal-Bench-2 (tbench2) run. + +The *tasks* are not hand-written: they are the +Terminal-Bench-2 suite shipped in the laude-institute/terminal-bench-2 repo. The +env serves the per-task instruction at reset(), so each prompt-data row only +needs the system prompt (how the agent should behave) plus the ``task_id`` in +metadata. ``openenv_agent_function`` drives the multi-turn loop and reads +metadata["task_id"]. + +task_ids are the top-level task directory names in the TB2 repo checkout +(e.g. "chess-best-move"); a valid task dir contains a ``task.toml``. + + # all 89 tasks + python make_tbench2_data.py --tasks_dir /workspace/terminal-bench-2 \ + --output /root/tbench2_train.jsonl + # a small smoke subset + python make_tbench2_data.py --tasks_dir /workspace/terminal-bench-2 \ + --output /root/tbench2_smoke.jsonl --n 8 + # an explicit subset + python make_tbench2_data.py --tasks_dir /workspace/terminal-bench-2 \ + --tasks chess-best-move,circuit-fibsqrt +""" + +import json +from pathlib import Path + +from tap import Tap + +# The agent contract must match openenv_agent_function._multi_turn: one shell +# command per turn inside a single ```bash block; TASK_COMPLETE to stop. +_SYSTEM = ( + "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)." +) + + +class Args(Tap): + tasks_dir: str = "/workspace/terminal-bench-2" # TB2 repo checkout + output: str = "/root/tbench2_train.jsonl" + n: int = 0 # 0 = all discovered tasks + tasks: str = "" # optional comma-separated explicit task_ids + + +def _discover_task_ids(tasks_dir: Path) -> list[str]: + return sorted(p.name for p in tasks_dir.iterdir() if (p / "task.toml").is_file()) + + +def main() -> None: + args = Args().parse_args() + tasks_dir = Path(args.tasks_dir).expanduser().resolve() + + if args.tasks: + task_ids = [t.strip() for t in args.tasks.split(",") if t.strip()] + else: + task_ids = _discover_task_ids(tasks_dir) + if args.n > 0: + task_ids = task_ids[: args.n] + + missing = [t for t in task_ids if not (tasks_dir / t / "task.toml").is_file()] + if missing: + raise SystemExit(f"task_ids without a task.toml in {tasks_dir}: {missing}") + + with open(args.output, "w") as f: + for tid in task_ids: + row = { + "prompt": [{"role": "system", "content": _SYSTEM}], + "metadata": {"task_id": tid}, + } + f.write(json.dumps(row) + "\n") + print(f"Wrote {len(task_ids)} TB2 tasks to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/examples/experimental/openenv/openenv_agent_function.py b/examples/experimental/openenv/openenv_agent_function.py new file mode 100644 index 0000000000..20335e1ded --- /dev/null +++ b/examples/experimental/openenv/openenv_agent_function.py @@ -0,0 +1,424 @@ +"""OpenEnv Terminal-Bench-2 <-> miles adapter. + +miles selects the policy and calls ``run`` once per episode via +``--custom-agent-function-path openenv_agent_function.run``. The episode drives +the OpenEnv tbench2 env through an agentic loop (reset -> {policy -> exec} -> +evaluate). + +The policy is always reached at ``base_url/v1`` through miles' session server, +so token ids + logprobs + loss masks are captured natively (no re-tokenization) +on every turn of the multi-turn episode. + +Env vars: + OPENENV_ENV_URL base_url of the env server (default: http://localhost:8003). + OPENENV_MAX_TURNS multi-turn cap (default: 30) + OPENENV_MESSAGE_TIMEOUT_S per-message WS recv timeout (default: 600; docker-mode + reset/exec/pytest routinely exceed the client default of 60) + OPENENV_MAX_ROLLOUT_TIME_SECONDS hard wall-clock cap for one episode (default: + 3600). An episode that does not return within the limit is + terminated and scored reward 0 (bounds long-trajectory + stragglers that would otherwise stall the whole rollout batch). + AGENT_MODEL_NAME model name sent to the policy (default: "model") + MILES_ROUTER_EXTERNAL_HOST optional host rewrite for off-cluster agents + OPENENV_TASK_WORKDIR container dir every agent command + eval runs in (default: + /app, the TB2 task image WORKDIR). Empty string disables the + prefix. Needed because upstream OpenEnv defaults to /task. + OPENENV_TB2_TESTS_SRC where the upstream env stages the task's tests inside the + container (default: /task/tests); copied to /tests for test.sh. +""" + +import asyncio +import logging +import os +import random +import re +import time +from collections.abc import Callable +from typing import Any +from urllib.parse import urlparse, urlunparse + +from openai import AsyncOpenAI + +logger = logging.getLogger(__name__) + +# Env slots are finite: hosted spaces admit one session at a time +# (CAPACITY_REACHED) and the docker-mode tbench2 server caps concurrent envs +# (MAX_CONCURRENT_ENVS), closing the WebSocket cleanly (ConnectionClosedOK) once +# full. Both are transient -- episodes hold a slot only for their rollout -- so +# jittered backoff + retry serializes the surplus rather than failing it. +# +# A rollout fans out more episodes than the server has slots (e.g. ~32 episodes +# vs a 16-session cap), so a queued episode must outwait a full episode ahead of +# it -- minutes, not seconds. The wait deadline is sized for that; the backoff +# ceiling is wide enough that 16 queued episodes don't hammer the server with +# reconnects while they wait. +_CAPACITY_MAX_WAIT_S = 1800.0 +_CAPACITY_BACKOFF_S = (1.0, 5.0) + +# Strip a single fenced block: ```python / ```bash / ``` ... ```. +_FENCE_RE = re.compile(r"```(?:python|py|bash|sh)?\s*\n?(.*?)```", re.DOTALL | re.IGNORECASE) + +# Max chars of command output fed back to the policy per turn (keeps context bounded). +_OBS_CHAR_CAP = 4000 + +# --- Adapter-driven Terminal-Bench-2 fidelity -------------------------------- +# Upstream OpenEnv's Tbench2DockerEnvironment runs the task container with workdir +# /task (a copy of the task *source*) and scores via bare `pytest tests/` there. +# Real TB2 tasks live at /app (the task image's WORKDIR) and are scored by the +# task's canonical tests/test.sh, which pins the pytest toolchain, copies test.py +# into /app, runs test_outputs.py, and writes the binary result to +# /logs/verifier/reward.txt. We reproduce that faithfully from the adapter -- +# without patching OpenEnv or vendoring it -- by (a) running every agent command +# in _TASK_WORKDIR and (b) driving the canonical harness through a plain `exec` +# step instead of the env's built-in (non-canonical) `evaluate` action. +# +# This assumes the env server is UNMODIFIED upstream, which copies the task dir +# (tests included) into the container at _TB2_TESTS_SRC. +_TASK_WORKDIR = os.getenv("OPENENV_TASK_WORKDIR", "/app") +_TB2_TESTS_SRC = os.getenv("OPENENV_TB2_TESTS_SRC", "/task/tests") + +# The eval exec echoes reward.txt on this marker so we can parse it out of stdout. +_REWARD_MARKER = "__TB2_REWARD__:" +# test.sh's exit code, echoed on its own marker purely for diagnostics: when a +# sample is dropped for having no recoverable reward, a nonzero rc points at a +# test.sh crash (infra/harness failure) vs. a clean run that wrote no verdict. +# It does NOT drive the drop decision -- a nonzero rc from merely-failing tests +# is a legitimate reward 0, not an infra error. +_TESTSH_RC_MARKER = "__TB2_TESTSH_RC__:" +# Honor an empty _TASK_WORKDIR (workdir prefix disabled) the same way +# _apply_workdir does, instead of silently forcing /app. +_EVAL_CD_CMD = f"cd {_TASK_WORKDIR} && " if _TASK_WORKDIR else "" +_CANONICAL_EVAL_CMD = ( + # rm the reward file first so a stale one can never be read back if test.sh + # fails to run (e.g. in a reused sandbox where /logs survives across episodes). + "mkdir -p /tests /logs/verifier && rm -f /logs/verifier/reward.txt && " + f"cp -a {_TB2_TESTS_SRC}/. /tests/ 2>/dev/null || true; " + f"{_EVAL_CD_CMD}bash /tests/test.sh > /tmp/tb2_testsh.log 2>&1; " + # $? here is test.sh's exit code (captured before any other command runs). + f"echo {_TESTSH_RC_MARKER}$?; " + f"echo {_REWARD_MARKER}$(cat /logs/verifier/reward.txt 2>/dev/null)" +) + + +def _apply_workdir(command: str) -> str: + """Prefix an agent command so it runs in the real task workdir (/app).""" + if not _TASK_WORKDIR: + return command + return f"cd {_TASK_WORKDIR} && {command}" + + +def _parse_reward_marker(output: str) -> float | None: + """Parse the reward.txt value the canonical-eval exec echoed on its marker line. + + Returns None when no reward can be recovered -- no marker line, an empty + value (reward.txt absent, i.e. test.sh never wrote a verdict), or a + non-numeric value. These are infra/harness failures, not a task the agent + legitimately failed, so the caller drops the sample rather than scoring a + false 0.0 that would pollute the training signal. A genuine failure writes + reward.txt = 0 and is returned as 0.0. + """ + for line in output.splitlines()[::-1]: + if _REWARD_MARKER in line: + raw = line.split(_REWARD_MARKER, 1)[1].strip() + if not raw: + return None + try: + return float(raw) + except ValueError: + return None + return None + + +def _parse_testsh_rc(output: str) -> int | None: + """Parse test.sh's exit code off its marker line (diagnostic only, may be absent).""" + for line in output.splitlines()[::-1]: + if _TESTSH_RC_MARKER in line: + raw = line.split(_TESTSH_RC_MARKER, 1)[1].strip() + try: + return int(raw) + except ValueError: + return None + return None + + +# Per-message WS recv timeout. Docker-mode tbench2 reset (container create), +# exec, and evaluate (pytest) each routinely exceed the EnvClient default of 60s. +_MESSAGE_TIMEOUT_S = float(os.getenv("OPENENV_MESSAGE_TIMEOUT_S", "600")) + +# Hard wall-clock cap for one episode. The per-message timeout above bounds a +# single env op, and OPENENV_MAX_TURNS bounds the turn count, but neither bounds +# total episode time: a long agentic trajectory can loop for turns * (long +# generation) and stall the whole rollout batch (a step finishes only when the +# slowest of all concurrent episodes returns). An episode exceeding this cap is +# terminated (its coroutine cancelled) and scored reward 0. +_MAX_ROLLOUT_TIME_S = float(os.getenv("OPENENV_MAX_ROLLOUT_TIME_SECONDS", "3600")) + + +def _is_retryable_env_error(e: BaseException) -> bool: + """True when an env op failed only because no env slot was free (transient). + + The docker-mode tbench2 server caps concurrent envs; over that cap it either + returns CAPACITY_REACHED or closes the WebSocket cleanly (ConnectionClosedOK). + Both mean "retry once a slot frees up", not a genuine episode failure. Match + the close exceptions by class name so the adapter need not import websockets. + """ + if "CAPACITY_REACHED" in str(e): + return True + return type(e).__name__ in {"ConnectionClosedOK", "ConnectionClosedError", "ConnectionClosed"} + + +def _resolve_session_url(base_url: str) -> str: + """Build the OpenAI-compatible policy URL, rewriting host for off-cluster agents.""" + session_url = f"{base_url}/v1" + external_host = os.getenv("MILES_ROUTER_EXTERNAL_HOST") + if external_host: + parsed = urlparse(session_url) + netloc = f"{external_host}:{parsed.port}" if parsed.port else external_host + session_url = urlunparse(parsed._replace(netloc=netloc)) + return session_url + + +def _extract_messages(prompt: Any) -> list[dict[str, str]]: + """Accept either a chat-message list or a raw string prompt.""" + if isinstance(prompt, list): + return list(prompt) + return [{"role": "user", "content": str(prompt)}] + + +def _strip_fence(text: str) -> str: + """Return the contents of a single fenced block, else the stripped text.""" + match = _FENCE_RE.search(text) + if match: + return match.group(1).strip() + return text.strip() + + +def _obs_field(result: Any, name: str) -> str: + """Read an Observation field off a StepResult, tolerating shape differences.""" + obs = getattr(result, "observation", result) + return str(getattr(obs, name, "") or "") + + +# Lazy import so the file loads without the env client present at import time. +def _load_tbench2() -> dict[str, Any]: + from tbench2_env import Tbench2Action, Tbench2Env + + return {"env": Tbench2Env, "action": Tbench2Action} + + +_DEFAULT_ENV_URL = "http://localhost:8003" + + +async def _with_env(env_cls: Any, env_url: str, body: Callable[[Any], Any]) -> Any: + """Open an env session and run ``body(env)``, retrying while a slot is busy.""" + deadline = asyncio.get_event_loop().time() + _CAPACITY_MAX_WAIT_S + while True: + try: + async with env_cls(base_url=env_url, message_timeout_s=_MESSAGE_TIMEOUT_S) as env: + return await body(env) + except Exception as e: + if _is_retryable_env_error(e) and asyncio.get_event_loop().time() < deadline: + await asyncio.sleep(random.uniform(*_CAPACITY_BACKOFF_S)) + continue + raise + + +async def _multi_turn( + classes: dict[str, Any], + env_url: str, + policy: AsyncOpenAI, + model_name: str, + messages: list[dict[str, str]], + request_kwargs: dict[str, Any], + metadata: dict[str, Any], +) -> tuple[float | None, dict[str, Any]]: + """Agentic loop: reset(task) -> {policy -> exec -> feed output back} -> evaluate (tbench2). + + The policy emits one shell command per turn (a ```bash block or the bare + reply), executed in the real task workdir (_TASK_WORKDIR, /app); the loop ends + when the policy stops emitting a command, says TASK_COMPLETE, or hits + OPENENV_MAX_TURNS. Scoring runs the task's canonical tests/test.sh via an + ``exec`` step and parses /logs/verifier/reward.txt for the binary reward + (faithful to Terminal-Bench-2, and needs no OpenEnv-side changes). + """ + action_cls = classes["action"] + task_id = metadata.get("task_id") or metadata.get("task_name") + max_turns = int(os.getenv("OPENENV_MAX_TURNS", "30")) + + async def body(env: Any) -> tuple[float | None, int, list[float], list[float], float, float, int | None]: + # Per-turn wall-clock timings. gen_times[i] is turn i's policy generation + # latency; tool_times[i] is turn i's env.step(exec) latency. reset_time and + # eval_time bracket the one-off reset() and the final evaluate() env steps. + gen_times: list[float] = [] + tool_times: list[float] = [] + + t0 = time.monotonic() + reset_result = await (env.reset(task_id=task_id) if task_id else env.reset()) + reset_time = time.monotonic() - t0 + instruction = _obs_field(reset_result, "instruction") + convo = list(messages) + if instruction: + convo.append({"role": "user", "content": instruction}) + + turns = 0 + while turns < max_turns: + turns += 1 + t0 = time.monotonic() + completion = await policy.chat.completions.create( + model=model_name, messages=convo, extra_body=request_kwargs + ) + gen_times.append(time.monotonic() - t0) + message = completion.choices[0].message + reply = message.content or "" + # Echo the assistant turn back verbatim. The session server stores the + # message exactly as SGLang emitted it -- content plus reasoning_content + # and tool_calls split out by the reasoning/tool-call parsers -- and + # matches each later request against that stored prefix. A thin + # {role, content} dict drops reasoning_content/tool_calls, diverges at + # the assistant turn, and trips "rollback failed: no assistant message + # in matched prefix". model_dump round-trips whatever the SDK parsed + # (extras like reasoning_content included). + convo.append(message.model_dump(exclude_none=True)) + + command = _strip_fence(reply) if "```" in reply else reply.strip() + if not command or command.upper().startswith("TASK_COMPLETE"): + break + + t0 = time.monotonic() + step_result = await env.step(action_cls(action_type="exec", command=_apply_workdir(command))) + tool_times.append(time.monotonic() - t0) + output = _obs_field(step_result, "output") + # Feed the command output back as a user turn, not a tool turn. GLM + # emits native tool_calls that we must echo verbatim (above) for the + # session server's prefix match; a role="tool" reply would then have + # to carry a matching tool_call_id and trips OpenAI tool-call + # validation. A plain user turn sidesteps the handshake -- the same + # text protocol the Harbor mini-swe-agent scaffold uses. + # + # Substitute a placeholder when a command produces no stdout: SGLang + # rejects an empty message content with "content cannot be empty". + content = output[:_OBS_CHAR_CAP] or "(no output)" + convo.append({"role": "user", "content": content}) + + t0 = time.monotonic() + eval_result = await env.step(action_cls(action_type="exec", command=_CANONICAL_EVAL_CMD)) + eval_time = time.monotonic() - t0 + eval_output = _obs_field(eval_result, "output") + reward = _parse_reward_marker(eval_output) + testsh_rc = _parse_testsh_rc(eval_output) + + # rm-hack: the tbench2 env server (TB2_OUTPUT_DIR=/tmp/tbench2_env_runs) + # leaves a per-episode trial dir under that path after every episode, which + # fills the sandbox overlay disk and trips ENOSPC. One episode holds the + # sandbox at a time, so it is safe to purge them here. + # + # BUT the same dir also holds repo_cache/ (TB2_CACHE_DIR defaults to + # output_dir/repo_cache) -- the shared terminal-bench-2 checkout that + # reset() clones once and every later episode reads its task from + # (repo_cache/terminal-bench-2-main/). A blanket `rm -rf .../*` wiped + # repo_cache too, so on a pooled/reused sandbox every episode after the first + # either re-cloned the whole repo (huge) or raced into "Task path not found", + # collapsing effective concurrency and exploding step time. Preserve + # repo_cache; delete only the ephemeral per-trial dirs beside it. + try: + await env.step( + action_cls( + action_type="exec", + command=( + "find /tmp/tbench2_env_runs -mindepth 1 -maxdepth 1 " + "! -name repo_cache -exec rm -rf {} + 2>/dev/null || true" + ), + ) + ) + except Exception: + pass + + return reward, turns, gen_times, tool_times, reset_time, eval_time, testsh_rc + + reward, turns, gen_times, tool_times, reset_time, eval_time, testsh_rc = await _with_env( + classes["env"], env_url, body + ) + total_gen_time = sum(gen_times) + # non_generation_time = everything the rollout spent outside policy generation: + # per-turn exec latency plus the one-off reset() and evaluate() env steps. Feeds + # Sample.non_generation_time so miles' throughput accounting subtracts env time. + total_tool_time = sum(tool_times) + reset_time + eval_time + return reward, { + "turns": turns, + "tool_calls": len(tool_times), + "gen_times": gen_times, + "tool_times": tool_times, + "reset_time": reset_time, + "eval_time": eval_time, + "total_gen_time": total_gen_time, + "total_tool_time": total_tool_time, + "testsh_rc": testsh_rc, + } + + +async def run( + base_url: str, + prompt: Any, + request_kwargs: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, + **kwargs, +) -> dict[str, Any] | None: + """Run one OpenEnv tbench2 episode via the trained policy.""" + request_kwargs = request_kwargs or {} + metadata = metadata or {} + + classes = _load_tbench2() + session_url = _resolve_session_url(base_url) + model_name = os.getenv("AGENT_MODEL_NAME", os.getenv("SWE_AGENT_MODEL_NAME", "model")) + env_url = os.getenv("OPENENV_ENV_URL", _DEFAULT_ENV_URL) + + policy = AsyncOpenAI(base_url=session_url, api_key="EMPTY") + messages = _extract_messages(prompt) + + try: + # Hard wall-clock cap: cancel the episode if it overruns and score it 0. + # wait_for cancels the coroutine, so any in-flight policy call / env.step + # is interrupted and the env session is closed by _with_env's async-with + # during cancellation cleanup. + reward, agent_metrics = await asyncio.wait_for( + _multi_turn(classes, env_url, policy, model_name, messages, request_kwargs, metadata), + timeout=_MAX_ROLLOUT_TIME_S, + ) + except asyncio.TimeoutError: + logger.warning(f"OpenEnv tbench2 episode exceeded {_MAX_ROLLOUT_TIME_S:.0f}s; " "terminating with reward 0") + # eval_report empty: the episode was cancelled before the canonical + # eval ever ran, so there is no pytest report to surface. + return { + "reward": 0.0, + "exit_status": "timeout", + "eval_report": {}, + "agent_metrics": {"timed_out": 1}, + } + except Exception as e: + logger.error(f"OpenEnv tbench2 episode failed: {e}", exc_info=True) + return None + finally: + await policy.close() + + # No recoverable reward means the canonical harness never produced a verdict + # (infra/harness failure, not a legitimate task failure). Drop the sample -- + # returning it as reward 0.0 would inject a false negative into training. + if reward is None: + logger.warning( + "OpenEnv tbench2 episode produced no canonical reward " + f"(test.sh exit code={agent_metrics.get('testsh_rc')}); " + "infra/harness failure, dropping sample" + ) + return None + + # eval_report is intentionally empty: the canonical-eval marker protocol + # (see _REWARD_MARKER) echoes back only the scalar reward. The detailed + # pytest CTRF report is written inside the sandbox at + # /logs/verifier/ctrf.json and is deliberately not captured back to the + # trainer, which consumes only `reward`. + return { + "reward": reward, + "exit_status": "completed", + "eval_report": {}, + "agent_metrics": agent_metrics, + } diff --git a/examples/experimental/openenv/openenv_generate.py b/examples/experimental/openenv/openenv_generate.py new file mode 100644 index 0000000000..abcd0ade9f --- /dev/null +++ b/examples/experimental/openenv/openenv_generate.py @@ -0,0 +1,16 @@ +"""Reward function for the OpenEnv Terminal-Bench-2 (tbench2) run. + +Task-agnostic: the agent function (``openenv_agent_function.run``) stores the +env-computed binary pytest reward in ``sample.metadata["reward"]``; this just +reads it back. Wired via ``--custom-rm-path openenv_generate.reward_func`` and +mirrors ``swe-agent-v2/generate.py:reward_func`` so it works for both the +single-sample (``async_rm``) and batched (``--custom-rm-path``) call paths. +""" + +from miles.utils.types import Sample + + +async def reward_func(args, samples: Sample | list[Sample], **kwargs) -> float | list[float]: + if isinstance(samples, list): + return [s.metadata.get("reward", 0.0) for s in samples] + return samples.metadata.get("reward", 0.0) diff --git a/examples/experimental/openenv/openenv_launch_common.py b/examples/experimental/openenv/openenv_launch_common.py new file mode 100644 index 0000000000..560ba6637d --- /dev/null +++ b/examples/experimental/openenv/openenv_launch_common.py @@ -0,0 +1,154 @@ +"""Shared launch helpers for the OpenEnv tbench2 learning launchers. + +``run-openenv-tbench2.py`` (GLM-4.7-Flash) is the launcher in this example; +sibling per-model launchers (e.g. a DeepSeek-V4-Flash variant) reuse the same +agentic adapter and differ only in the model-family serving/training profile. +The model-agnostic fragments (process cleanup, GRPO/optimizer/rollout/agent +flags, W&B + Prometheus wiring, and the OpenEnv env-var plumbing) live here so +those launchers cannot silently drift apart. Each launcher keeps only its own +perf/sglang/misc profile and its ``ScriptArgs`` defaults. +""" + +import os +import subprocess +import time +from typing import Protocol + + +class LaunchArgs(Protocol): + """The config fields the shared helpers read (satisfied by each launcher's ScriptArgs).""" + + prompt_data: str + rollout_batch_size: int + n_samples_per_prompt: int + max_seq_len: int + global_batch_size: int + + openenv_env_url: str + agent_model_name: str + openenv_max_turns: int + openenv_max_rollout_time_seconds: int + router_external_host: str + miles_host_ip: str + + wandb_key: str + wandb_project: str + wandb_team: str + wandb_run_name: str + + use_prometheus: bool + prometheus_port: int + prometheus_run_name: str + + +def cleanup() -> None: + """Kill old Ray jobs and stale processes to free GPU resources.""" + my_pid = os.getpid() + ppid = os.getppid() + print(f"Cleanup starting (pid={my_pid}, ppid={ppid})") + targets = ["sglang", "train.py", "MegatronTrain"] + exclude = f"grep -v '^{my_pid}$' | grep -v '^{ppid}$'" + for t in targets: + subprocess.run( + f"pgrep -f '{t}' | {exclude} | xargs -r kill 2>/dev/null || true", + shell=True, + ) + time.sleep(5) + print(f"Cleanup complete (pid={my_pid}) — old processes killed.") + + +def rollout_args(args: LaunchArgs) -> str: + return ( + f"--prompt-data {args.prompt_data} " + "--input-key prompt " + "--metadata-key metadata " + "--rollout-shuffle " + "--num-rollout 40 " + f"--rollout-batch-size {args.rollout_batch_size} " + f"--n-samples-per-prompt {args.n_samples_per_prompt} " + "--rollout-temperature 0.8 " + "--rollout-max-response-len 8192 " + f"--max-seq-len {args.max_seq_len} " + f"--global-batch-size {args.global_batch_size} " + "--balance-data " + ) + + +def grpo_args() -> str: + return ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.01 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.0 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + +def optimizer_args() -> str: + return ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + +def agent_args(tito_model: str) -> str: + """Agentic-rollout wiring. Only the TITO surface differs across models.""" + return ( + "--custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate " + "--custom-agent-function-path openenv_agent_function.run " + "--custom-rm-path openenv_generate.reward_func " + "--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_no_aborted " + f"--tito-model {tito_model} " + "--use-session-server " + "--session-server-port 30000 " + "--tito-allowed-append-roles user tool " + ) + + +def wandb_args(args: LaunchArgs) -> str: + if not args.wandb_key: + return "" + out = ( + "--use-wandb " + f"--wandb-project {args.wandb_project} " + f"--wandb-group {args.wandb_run_name} " + f"--wandb-key {args.wandb_key} " + ) + if args.wandb_team: + out += f"--wandb-team {args.wandb_team} " + return out + + +def prometheus_args(args: LaunchArgs) -> str: + if not args.use_prometheus: + return "" + return ( + "--use-prometheus " + f"--prometheus-port {args.prometheus_port} " + f"--prometheus-run-name {args.prometheus_run_name} " + ) + + +def base_env_vars(args: LaunchArgs, script_dir: str, megatron_path: str, miles_root: str) -> dict[str, str]: + return { + "PYTHONPATH": f"{megatron_path}:{script_dir}:{miles_root}", + "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1", + "OPENENV_ENV_URL": args.openenv_env_url, + "OPENENV_MAX_TURNS": str(args.openenv_max_turns), + "OPENENV_MAX_ROLLOUT_TIME_SECONDS": str(args.openenv_max_rollout_time_seconds), + "AGENT_MODEL_NAME": args.agent_model_name, + } + + +def apply_optional_env_vars(env: dict[str, str], args: LaunchArgs) -> None: + """Add host-rewrite env vars when the args request them.""" + if args.miles_host_ip: + env["MILES_HOST_IP"] = args.miles_host_ip + if args.router_external_host: + env["MILES_ROUTER_EXTERNAL_HOST"] = args.router_external_host diff --git a/examples/experimental/openenv/run-openenv-tbench2.py b/examples/experimental/openenv/run-openenv-tbench2.py new file mode 100644 index 0000000000..d3faf92a4d --- /dev/null +++ b/examples/experimental/openenv/run-openenv-tbench2.py @@ -0,0 +1,211 @@ +"""OpenEnv Terminal-Bench-2 (tbench2) learning launcher (GLM-4.7-Flash). + +Drives the OpenEnv tbench2 env via ``openenv_agent_function.run``. tbench2 is +*multi-turn*: the adapter runs an agentic loop (reset(task_id) -> {policy emits a +shell command -> step(exec) -> feed output back} -> evaluate) and the reward is +the binary pytest result (1.0 all tests pass, else 0.0). + +Prereqs: + # 1. Install the env client where the rollout runs (pulls camel-ai; isolate + # from the training env if its deps clash with the miles image). + pip install -e /envs/tbench2_env + # 2. Get the TB2 task suite + build prompt-data (task_ids). + git clone --depth 1 https://github.com/laude-institute/terminal-bench-2.git /workspace/terminal-bench-2 + python make_tbench2_data.py --tasks_dir /workspace/terminal-bench-2 --output /root/tbench2_train.jsonl + # 3. Serve the env. The tbench2 server supports concurrency natively via + # MAX_CONCURRENT_ENVS (no wrapper needed). Choose execution mode: + # TB2_MODE=docker -> real TB2 fidelity (needs docker.sock + image pulls) + # TB2_MODE=local -> runs in-process, ignores task Dockerfiles (degraded) + TB2_MODE=docker TB2_TASKS_DIR=/workspace/terminal-bench-2 MAX_CONCURRENT_ENVS=32 \ + python -m tbench2_env.server.app --port 8003 + + NOTE (open decisions before a real run): docker mode wants a Docker host with + disk + socket; colocating heavy per-task containers on the GPU pod is risky, + so the env server likely runs off-pod (use --openenv-env-url / the host + rewrite). The binary sparse reward also needs a task subset where the base + policy *sometimes* succeeds (advantage variance) -- e.g. the TB2 variance + band -- or GRPO sees a flat signal. + +Usage: + python run-openenv-tbench2.py --openenv-env-url http://:8003 +""" + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import openenv_launch_common as C +import typer + +import miles.utils.external_utils.command_utils as U + +SCRIPT_DIR = Path(__file__).resolve().parent + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + mode: Literal["normal", "debug_rollout_only"] = "normal" + run_id: str = U.create_run_id() + megatron_model_type: str = "glm4.7-flash" + num_gpus_per_node: int = 8 + megatron_path: str = "/root/Megatron-LM" + + # Paths + skip_prepare: bool = False + base_dir: str = "/root" + model_name: str = "GLM-4.7-Flash" + hf_checkpoint: str = "zai-org/GLM-4.7-Flash" + ref_load: str = "/root/GLM-4.7-Flash_torch_dist" + save_dir: str = "/workspace/GLM-4.7-Flash_openenv_tbench2/" + prompt_data: str = "/root/tbench2_train.jsonl" + + # Training settings (small; multi-turn so responses run long) + max_seq_len: int = 16384 + rollout_batch_size: int = 8 + n_samples_per_prompt: int = 8 + global_batch_size: int = 32 + + # OpenEnv settings + openenv_env_url: str = os.environ.get("OPENENV_ENV_URL", "http://localhost:8003") + agent_model_name: str = os.environ.get("AGENT_MODEL_NAME", "model") + openenv_max_turns: int = int(os.environ.get("OPENENV_MAX_TURNS", "30")) + # Hard wall-clock cap (seconds) per episode. An episode that does not return + # within the limit is terminated and scored reward 0, bounding long-trajectory + # stragglers that would otherwise stall the whole rollout batch. + openenv_max_rollout_time_seconds: int = int(os.environ.get("OPENENV_MAX_ROLLOUT_TIME_SECONDS", "3600")) + # When set, miles dumps full per-episode agent trajectories (tokens, logprobs, + # loss masks, reward, multi-turn messages) to /rollout_data/{rollout_id}.pt + # for post-hoc inspection via miles.utils.debug_utils.display_debug_rollout_data. + dump_details: str = os.environ.get("OPENENV_DUMP_DETAILS", "") + # Optional host rewrite for the policy URL (only needed if the in-process + # agent cannot reach the session server at its raw base_url host). + router_external_host: str = os.environ.get("MILES_ROUTER_EXTERNAL_HOST", "") + # Leave empty so miles resolves the numeric LAN IP itself. sgl-router's Rust + # binder rejects a hostname ("invalid socket address syntax"), and a numeric + # base_url host keeps the in-process policy client off hostname DNS too. + miles_host_ip: str = os.environ.get("MILES_HOST_IP", "") + + # W&B settings + wandb_key: str = os.environ.get("WANDB_KEY", os.environ.get("WANDB_API_KEY", "")) + wandb_project: str = os.environ.get("WANDB_PROJECT", "openenv-tbench2-learn") + wandb_team: str = os.environ.get("WANDB_TEAM", "") + wandb_run_name: str = "openenv-tbench2-learn" + + # Prometheus settings + use_prometheus: bool = True + prometheus_port: int = 9090 + prometheus_run_name: str = "openenv-tbench2-learn" + + +def prepare(args: ScriptArgs): + """Convert HF checkpoint to torch_dist format if not already done.""" + U.convert_checkpoint( + model_name=args.model_name, + megatron_model_type=args.megatron_model_type, + num_gpus_per_node=args.num_gpus_per_node, + dir_dst=args.base_dir, + hf_checkpoint=args.hf_checkpoint, + megatron_path=args.megatron_path, + ) + + +def execute(args: ScriptArgs): + ckpt_args = ( + f"--hf-checkpoint {args.hf_checkpoint} " + f"--ref-load {args.ref_load} " + f"--save {args.save_dir} " + "--save-interval 100 " + ) + + rollout_args = C.rollout_args(args) + + perf_args = ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 2 " # single 8-GPU node: TP=4 -> DP=2, so EP<=2 + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 16384 " + "--optimizer-cpu-offload " + "--overlap-cpu-optimizer-d2h-h2d " + "--use-precision-aware-optimizer " + ) + + grpo_args = C.grpo_args() + + optimizer_args = C.optimizer_args() + + sglang_args = ( + "--rollout-num-gpus-per-engine 1 " + "--sglang-mem-fraction-static 0.7 " + "--sglang-tool-call-parser glm47 " + "--sglang-reasoning-parser glm45 " + "--sglang-router-port 31000 " + ) + + agent_args = C.agent_args("glm47") + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--colocate " + f"--actor-num-nodes {args.num_nodes} " + f"--actor-num-gpus-per-node {args.num_gpus_per_node} " + f"--rollout-num-gpus {args.num_gpus_per_node} " + ) + + debug_args = "--debug-rollout-only " if args.mode == "debug_rollout_only" else "" + + dump_args = f"--dump-details {args.dump_details} " if args.dump_details else "" + + wandb_args = C.wandb_args(args) + + prometheus_args = C.prometheus_args(args) + + train_args = ( + f"{ckpt_args}" + f"{rollout_args}" + f"{optimizer_args}" + f"{grpo_args}" + f"{wandb_args}" + f"{prometheus_args}" + f"{perf_args}" + f"{sglang_args}" + f"{agent_args}" + f"{misc_args}" + f"{debug_args}" + f"{dump_args}" + ) + + extra_env_vars = C.base_env_vars(args, str(SCRIPT_DIR), args.megatron_path, U.repo_base_dir) + C.apply_optional_env_vars(extra_env_vars, args) + + U.execute_train( + train_args=train_args, + config=args, + num_gpus_per_node=args.num_gpus_per_node, + megatron_model_type=args.megatron_model_type, + megatron_path=args.megatron_path, + extra_env_vars=extra_env_vars, + ) + + +@U.dataclass_cli +def main(args: ScriptArgs): + C.cleanup() + if not args.skip_prepare: + prepare(args) + execute(args) + + +if __name__ == "__main__": + typer.run(main)