diff --git a/examples/experimental/openenv_swebench/README.md b/examples/experimental/openenv_swebench/README.md new file mode 100644 index 0000000000..7ec4055956 --- /dev/null +++ b/examples/experimental/openenv_swebench/README.md @@ -0,0 +1,142 @@ +# OpenEnv SWE-bench-style GRPO (GLM-4.7-Flash, single node) + +Train GLM-4.7-Flash with GRPO on **SWE-bench-style** tasks (SWE-Rebench-V2 +"donor" variants) through the HuggingFace [OpenEnv](https://github.com/huggingface/openenv) +env server. This is a sibling of [`../openenv`](../openenv) (Terminal-Bench-2): +same env server, same agentic loop, same reward-marker protocol — only the task +family and its grading differ. A miles-side adapter runs the multi-turn agentic +loop (`reset(task_id)` → { policy emits one shell command → `step(exec)` → feed +output back } → verify) against an **unmodified** OpenEnv env server; the reward +is the binary verifier result (1.0 if the task's `test_command` passes, 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. + +## How this differs from the Terminal-Bench-2 example + +The OpenEnv env server is generic: it **pulls** the image named in a task's +`task.toml [environment].docker_image`, copies the task dir into the container, +and exposes an `exec` step. Everything task-family-specific is in this adapter: + +| | Terminal-Bench-2 (`../openenv`) | SWE-bench-style (here) | +| --- | --- | --- | +| Agent working dir | fixed `/app` | **auto-detected repo root** (`/app`, `/testbed`, `/`, …) | +| Verifier | `tests/test.sh` writes `/logs/verifier/reward.txt` | `tests/test.sh` prints `RESULT: PASSED`/`FAILED`, exits 0/1 (no `reward.txt`) | +| Reward source | value in `reward.txt` | derived from the `RESULT:` verdict line | +| No verdict recovered | drop sample | drop sample | + +Both incompatibilities are handled in `swebench_agent_function.py` — the OpenEnv +repo is **not** patched or vendored. + +## 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 env client (the same `tbench2_env` client the TB2 example + uses; isolate it if its deps clash with the miles image): + + ```bash + pip install -e /envs/tbench2_env + ``` + +## 1. Make each task pullable (data prep) + +The env server only ever **pulls** `task.toml [environment].docker_image`; it +does not build a Dockerfile. SWE-Rebench-V2 donor tasks ship an +`environment/Dockerfile` (e.g. `FROM docker.io/swerebenchv2/:`) and no +`docker_image`. So, once per task, build the image and record it — this is +offline data prep, **no code change**: + +```bash +# For each task dir in your pool: +docker build -t /:latest /environment +docker push /:latest +# then add to /task.toml under [environment]: +# docker_image = "/:latest" +``` + +Many donor Dockerfiles only `FROM` a `swerebenchv2/...` base plus a `git reset`; +if the env host can pull that base directly you can point `docker_image` straight +at it, skipping the per-task build. + +## 2. Build the prompt data + +Emit one prompt row per `task_id` (the task dir name): + +```bash +python make_swebench_data.py --tasks_dir /root/swebench_pool --output /root/swebench_train.jsonl +# add --n 8 for a small smoke subset +``` + +## 3. Start the env server + +Run it in a separate shell (or off-node). Point `TB2_TASKS_DIR` at the pool from +step 1/2. Docker mode 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=/root/swebench_pool 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`. + +## 4. Launch training + +```bash +python run-openenv-swebench.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/swebench_train.jsonl` | Prompt set from step 2 | +| `--num-rollout` | `40` | Number of GRPO steps | +| `OPENENV_MAX_TURNS` | `30` | Max agent turns per episode | +| `OPENENV_TASK_WORKDIR` | (unset → auto-detect) | Force a fixed repo path instead of detecting it | +| `OPENENV_SWEBENCH_TESTS_SRC` | `/task/tests` | Where the env stages the task's tests in the container | +| `OPENENV_EVAL_CMD` | (built-in) | Override the whole grading command (must print `__SB_REWARD__:` last) | +| `OPENENV_MAX_ROLLOUT_TIME_SECONDS` | `3600` | Per-episode wall-clock cap; a straggler that exceeds it is terminated and the sample is dropped (the cap covers infra phases too, so a timeout can't produce a verdict — dropping avoids a false negative) | +| `OPENENV_SWEBENCH_TESTS_SNAPSHOT` | `/opt/.sb_pristine_tests` | Where the pristine grader tests are snapshotted at reset; grading reads this so mid-episode tampering with the staged tests can't reach the grader. Set `""` to disable | +| `OPENENV_SKIP_CLEANUP` | (unset) | Set `1` to skip the pre-run process reaper (use when another of your jobs is intentionally sharing the node) | +| `--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). If every task is always-fail + or always-pass, GRPO sees a flat signal — curate a variance band (or use a + stronger base) to see a learning climb. +- **Dropped samples.** If the verifier can't produce a `RESULT:` verdict (e.g. it + can't locate the repo because the agent trashed the checkout), the episode has + no recoverable reward and the sample is **dropped** rather than scored 0 — same + policy as the TB2 adapter, so an infra/harness failure never becomes a false + negative in training. A `OPENENV_MAX_ROLLOUT_TIME_SECONDS` timeout is treated + the same way: the cap covers capacity/reset/generation/eval, so a timeout may + fire before the verifier ran, which is a no-verdict case — dropped, not scored 0. +- **Verifier tampering.** The agent shares the task container with the grader + assets, so it *could* edit them before grading. The adapter snapshots the + pristine tests at reset and grades from that snapshot + (`OPENENV_SWEBENCH_TESTS_SNAPSHOT`), which defeats naive tampering, but SWE-bench + images run the agent as root so this is not a hard boundary — a fully robust + setup would grade in a separate container the agent never touches. +- **`_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, so an unclean disconnect + (trainer crash) can orphan containers. Sweep stale containers between runs. +- **Open-file limit.** 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`. Start the server with a raised limit (`ulimit -n 1048576`, as in step 3). diff --git a/examples/experimental/openenv_swebench/make_swebench_data.py b/examples/experimental/openenv_swebench/make_swebench_data.py new file mode 100644 index 0000000000..08a9e223f8 --- /dev/null +++ b/examples/experimental/openenv_swebench/make_swebench_data.py @@ -0,0 +1,84 @@ +"""Generate a prompt dataset for the OpenEnv SWE-bench-style run. + +The *tasks* are the SWE-Rebench-V2 "donor" variants: a directory tree (one per +task) that the env server serves at reset(). Each task dir contains a +``task.toml``, an ``instruction.md`` (served as the reset() instruction), an +``environment/`` (the image the env pulls), and a ``tests/`` verifier. This +builder only needs to emit, per task, the system prompt (how the agent should +behave) plus the ``task_id`` in metadata; ``swebench_agent_function`` drives the +multi-turn loop and reads metadata["task_id"]. + +task_ids are the top-level task directory names in the pool (they must equal the +dir name, since the env resolves a task by ``tasks_dir/``); a valid task +dir contains a ``task.toml``. + + # all tasks in the pool + python make_swebench_data.py --tasks_dir /root/swebench_pool \ + --output /root/swebench_train.jsonl + # a small smoke subset + python make_swebench_data.py --tasks_dir /root/swebench_pool \ + --output /root/swebench_smoke.jsonl --n 8 + # an explicit subset + python make_swebench_data.py --tasks_dir /root/swebench_pool \ + --tasks synthdonor_99designs__aws-vault_8a6_738a5936 +""" + +import json +from pathlib import Path + +from tap import Tap + +# The agent contract must match swebench_agent_function._multi_turn: one shell +# command per turn inside a single ```bash block; TASK_COMPLETE to stop. The +# working directory is the repository root (the adapter cds the agent there). +_SYSTEM = ( + "You are an autonomous software engineer working inside a git repository. You " + "will be given a task describing a change to make to the codebase, then " + "interact with a real Linux shell whose working directory is the repository " + "root. On each turn respond with EXACTLY ONE shell command inside a single " + "```bash code block and nothing else. Explore the repository, implement the " + "required change, and verify it (e.g. by building the project or running the " + "relevant tests). Do not edit the test files that grade the task. When you are " + "confident the task is fully complete, reply with TASK_COMPLETE (with no code " + "block)." +) + + +class Args(Tap): + tasks_dir: str = "/root/swebench_pool" # pool of SWE-Rebench-V2 donor task dirs + output: str = "/root/swebench_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)} SWE-bench-style tasks to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/examples/experimental/openenv_swebench/run-openenv-swebench.py b/examples/experimental/openenv_swebench/run-openenv-swebench.py new file mode 100644 index 0000000000..5ef7d52c96 --- /dev/null +++ b/examples/experimental/openenv_swebench/run-openenv-swebench.py @@ -0,0 +1,211 @@ +"""OpenEnv SWE-bench-style learning launcher (GLM-4.7-Flash). + +Drives the OpenEnv env via ``swebench_agent_function.run`` over the SWE-Rebench-V2 +"donor" task family. This is *multi-turn*: the adapter runs an agentic loop +(reset(task_id) -> {policy emits a shell command -> step(exec) -> feed output +back} -> verify) and the reward is the binary verifier result (1.0 if the task's +test_command passes, 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). This is + # the same generic OpenEnv tbench2 env server the TB2 example uses. + pip install -e /envs/tbench2_env + # 2. Make each task pullable. The env server only PULLS the image named in + # task.toml [environment].docker_image; SWE-Rebench-V2 donor tasks ship an + # environment/Dockerfile instead. Build each task's image, push it to a + # registry the env host can reach, and set docker_image in its task.toml. + # (Offline data prep -- no code change; see README for the loop.) + # 3. Build prompt-data (task_ids) from the pool. + python make_swebench_data.py --tasks_dir /root/swebench_pool --output /root/swebench_train.jsonl + # 4. Serve the env pointing at the pool dir. + TB2_MODE=docker TB2_TASKS_DIR=/root/swebench_pool MAX_CONCURRENT_ENVS=32 \ + python -m tbench2_env.server.app --port 8003 + + NOTE: the binary sparse reward needs a task subset where the base policy + *sometimes* succeeds (advantage variance), or GRPO sees a flat signal. + +Usage: + python run-openenv-swebench.py --openenv-env-url http://:8003 +""" + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import swebench_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_swebench/" + prompt_data: str = "/root/swebench_train.jsonl" + + # Training settings (small; multi-turn so responses run long) + num_rollout: int = 40 # number of GRPO steps + 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-swebench-learn") + wandb_team: str = os.environ.get("WANDB_TEAM", "") + wandb_run_name: str = "openenv-swebench-learn" + + # Prometheus settings + use_prometheus: bool = True + prometheus_port: int = 9090 + prometheus_run_name: str = "openenv-swebench-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) diff --git a/examples/experimental/openenv_swebench/swebench_agent_function.py b/examples/experimental/openenv_swebench/swebench_agent_function.py new file mode 100644 index 0000000000..9922935f9a --- /dev/null +++ b/examples/experimental/openenv_swebench/swebench_agent_function.py @@ -0,0 +1,571 @@ +"""OpenEnv SWE-bench-style <-> miles adapter. + +A sibling of ``examples/experimental/openenv`` (Terminal-Bench-2). Same env +server, same agentic loop, same reward-marker protocol -- only the task family +differs. SWE-bench-style tasks are the SWE-Rebench-V2 "donor" variants: a git +repository baked into a prebuilt image, an instruction describing a code change, +and a shell verifier (``tests/test.sh`` + ``tests/config.json``) that runs the +task's ``test_command`` and prints ``RESULT: PASSED`` / ``RESULT: FAILED``. + +miles selects the policy and calls ``run`` once per episode via +``--custom-agent-function-path swebench_agent_function.run``. The episode drives +the OpenEnv 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. + +Two things make this family distinct from TB2, and both live entirely in this +adapter (the OpenEnv env server stays UNMODIFIED upstream): + + 1. Working directory. TB2 tasks live at a fixed /app; SWE-Rebench-V2 donor + images put the repo at varying paths (/app, /testbed, /, ...). + We auto-detect the repo root once (the same probe tests/test.sh uses) and + cd the agent into it, so the agent edits the real checkout. Override with + OPENENV_TASK_WORKDIR to force a fixed dir. + 2. Scoring. The donor verifier does not write /logs/verifier/reward.txt; it + prints ``RESULT: PASSED`` / ``RESULT: FAILED`` and exits 0/1. We run it via + a plain ``exec`` step and derive the reward from that verdict line. + +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/verify 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 the sample is dropped (the cap spans infra + phases too, so a timeout may precede any verdict -- dropping, + rather than scoring 0, avoids injecting a false negative while + still bounding stragglers that would stall the 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 force a fixed container dir for every agent command + (default: empty -> auto-detect the repo root per task). + OPENENV_CONDA_ENV conda env to activate for every agent command so the agent's + dev/test loop matches the grader (default: "testbed", the env + SWE-bench images install the repo + pytest into). Set to "" + to disable; probe is a no-op on images without such an env. + OPENENV_SWEBENCH_TESTS_SRC where the upstream env stages the task's tests inside + the container (default: /task/tests); copied to /tests for test.sh. + OPENENV_SWEBENCH_TESTS_SNAPSHOT path the pristine tests are snapshotted to at reset + (default: /opt/.sb_pristine_tests); grading reads this snapshot so + mid-episode tampering with the staged tests can't reach the grader. + Set "" to disable and grade from OPENENV_SWEBENCH_TESTS_SRC directly. + OPENENV_EVAL_CMD override the whole grading command (must still print + "<_REWARD_MARKER>" as its last stdout line). +""" + +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 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) + +# Command extraction. The policy narrates in prose, often with illustrative +# ```python planning snippets, before emitting the ONE command it wants to run. +# So we must pick the *last* shell fence and never execute a ```python block. +# _CMD_FENCE_RE captures every fence's language tag (group 1) and body (group 2); +# _extract_command keeps only shell-flavored (or language-less) ones. +_CMD_FENCE_RE = re.compile(r"```([A-Za-z0-9_.+-]*)[ \t]*\r?\n?(.*?)```", re.DOTALL) +_SHELL_LANGS = {"", "bash", "sh", "shell", "console", "shell-session", "shellsession", "zsh"} +# GLM-4.7 sometimes ignores the prompt and emits its native tool-call format +# `bashCMD` as raw text (SGLang has no +# matching tool parser, so it lands in message.content). Recover the command from it. +_TOOLCALL_ARG_RE = re.compile(r"(.*?)", re.DOTALL | re.IGNORECASE) + +_FORMAT_NUDGE = ( + "Your last message contained no runnable shell command. Respond with EXACTLY ONE " + "shell command inside a single ```bash code block and nothing else, or reply " + "TASK_COMPLETE (with no code block) when the task is fully done." +) + +# Max chars of command output fed back to the policy per turn (keeps context bounded). +_OBS_CHAR_CAP = 4000 + +# --- Working directory: auto-detect the repo root ---------------------------- +# The env server execs commands in /task (a copy of the task *source*), but the +# code the agent must edit lives in the git repo baked into the image, whose path +# varies per SWE-Rebench-V2 donor (/app, /testbed, /, ...). We probe +# for it once (mirroring tests/test.sh: try /app and /testbed, then a shallow +# .git scan) and cache the answer, so every subsequent turn cds there cheaply. +# Setting OPENENV_TASK_WORKDIR forces a fixed dir and skips detection. +_TASK_WORKDIR = os.getenv("OPENENV_TASK_WORKDIR", "") +_REPO_ROOT_CACHE = "/tmp/.openenv_swebench_repo_root" +_DETECT_REPO_ROOT = ( + f"R=$(cat {_REPO_ROOT_CACHE} 2>/dev/null); " + 'if [ -z "$R" ]; then ' + 'for c in /app /testbed; do [ -d "$c/.git" ] && R="$c" && break; done; ' + '[ -z "$R" ] && R=$(find / -maxdepth 3 -type d -name .git 2>/dev/null ' + "| grep -v node_modules | head -1 | xargs -r dirname); " + "R=${R:-/}; " + f'echo "$R" > {_REPO_ROOT_CACHE}; fi; ' +) + +# --- Conda env: run agent commands in the graded environment ------------------ +# SWE-bench images install the repo (editable) plus pytest and every runtime dep +# into a conda env (default name "testbed"), NOT the base env. tests/test.sh +# activates it before grading, but the base env the container boots into lacks +# pytest and cannot even import the target package. Without activating it here, +# the agent runs blind -- `pytest` is not found and `import ` raises +# ModuleNotFoundError -- so it gets no test feedback and its edits are never the +# thing that gets graded. We activate the env (mirroring test.sh) for every agent +# command so the agent's dev/test loop matches the grader. The probe is a no-op +# on donor images that have no such env, so it stays safe for non-SWE-bench pools. +# Set OPENENV_CONDA_ENV="" to disable, or to another name to override "testbed". +_CONDA_ENV = os.getenv("OPENENV_CONDA_ENV", "testbed") +_ACTIVATE_ENV = ( + ( + 'for _p in /opt/miniconda3 /opt/conda "$HOME/miniconda3" "$HOME/anaconda3"; do ' + f'if [ -f "$_p/bin/activate" ] && [ -d "$_p/envs/{_CONDA_ENV}" ]; then ' + f'. "$_p/bin/activate" {_CONDA_ENV} 2>/dev/null || true; break; fi; done; ' + ) + if _CONDA_ENV + else "" +) + +# --- Scoring: SWE-Rebench-V2 donor verifier ---------------------------------- +# The donor verifier (tests/test.sh) self-locates the repo, applies the variant's +# test_patch, runs test_command from tests/config.json, writes output.json, and +# prints "RESULT: PASSED" (exit 0) or "RESULT: FAILED" (exit 1). It never writes +# /logs/verifier/reward.txt, so -- unlike the TB2 adapter -- we derive the reward +# from the verdict line: PASSED -> 1.0, FAILED -> 0.0, and neither present -> no +# recoverable reward (the verifier crashed before scoring, e.g. it could not find +# the repo), which the caller drops rather than scoring a false 0.0. +_TESTS_SRC = os.getenv("OPENENV_SWEBENCH_TESTS_SRC", "/task/tests") +# Verifier-tampering mitigation. The agent has a shell in the same container as +# the grader assets under _TESTS_SRC, so a reward-hacking policy could edit +# test.sh/config.json before grading. We snapshot the tests to _TESTS_SNAPSHOT +# once at reset -- before the agent gets a turn -- and grade from that snapshot, +# so tampering with _TESTS_SRC mid-episode no longer reaches the grader. NOTE: +# this defeats the obvious attack but is NOT a hard security boundary -- SWE-bench +# images run the agent as root, which can still reach the snapshot path; a fully +# robust design grades in a separate container the agent never touches (out of +# scope here: needs env-server support). Set OPENENV_SWEBENCH_TESTS_SNAPSHOT="" +# to disable and grade from _TESTS_SRC directly. +_TESTS_SNAPSHOT = os.getenv("OPENENV_SWEBENCH_TESTS_SNAPSHOT", "/opt/.sb_pristine_tests") +_GRADE_TESTS_SRC = _TESTS_SNAPSHOT or _TESTS_SRC +_REWARD_MARKER = "__SB_REWARD__:" +# test.sh's exit code, echoed on its own marker purely for diagnostics. +_TESTSH_RC_MARKER = "__SB_TESTSH_RC__:" +_CANONICAL_EVAL_CMD = os.getenv("OPENENV_EVAL_CMD") or ( + "mkdir -p /tests /logs/verifier && " + # test.sh reads /tests/config.json by absolute path, so stage the tests there. + # Grade from the pristine reset-time snapshot (_GRADE_TESTS_SRC) so mid-episode + # edits to _TESTS_SRC do not reach the grader; fall back to _TESTS_SRC if the + # snapshot was never taken (donor tasks / snapshot disabled). + f"SB_SRC={_GRADE_TESTS_SRC}; [ -d $SB_SRC ] || SB_SRC={_TESTS_SRC}; " + "cp -a $SB_SRC/. /tests/ 2>/dev/null || true; " + "bash /tests/test.sh > /tmp/sb_testsh.log 2>&1; rc=$?; " + f"echo {_TESTSH_RC_MARKER}$rc; " + # The tbench2 env wraps every command as bash -c '' and docker-py + # shlex-splits that string, so any single quote here collides with the + # wrapper's quotes and shreds the command. Keep this command + # single-quote-free: use double quotes for the grep pattern. + 'V=$(grep -aoE "RESULT: (PASSED|FAILED)" /tmp/sb_testsh.log | tail -1); ' + f'if echo "$V" | grep -q PASSED; then echo {_REWARD_MARKER}1.0; ' + f'elif echo "$V" | grep -q FAILED; then echo {_REWARD_MARKER}0.0; ' + # No verdict line: verifier crashed before scoring -> emit empty value so the + # caller recognizes an infra/harness failure and drops the sample. + f"else echo {_REWARD_MARKER}; fi" +) + + +def _apply_workdir(command: str) -> str: + """Prefix an agent command so it runs in the task's repo root.""" + if _TASK_WORKDIR: + return f"{_ACTIVATE_ENV}cd {_TASK_WORKDIR} && {command}" + return f'{_DETECT_REPO_ROOT}{_ACTIVATE_ENV}cd "$R" && {command}' + + +def _parse_reward_marker(output: str) -> float | None: + """Parse the reward value the eval exec echoed on its marker line. + + Returns None when no reward can be recovered -- no marker line, an empty + value (the verifier printed no RESULT 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 emits 0.0 and is returned so. + """ + 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 reset (container create), exec, and +# verify 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 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 _extract_command(reply: str) -> str | None: + """Extract the single shell command the policy wants to run this turn. + + Returns None when the reply carries no runnable command (pure reasoning, a + lone ```python planning snippet, or a TASK_COMPLETE sentinel). We prefer the + *last* shell fence (the command usually follows the reasoning), skip + ```python/```py blocks (planning, not commands), and fall back to GLM's + native format. We never execute the whole raw reply, which used + to feed prose/pseudo-code to bash and produce a syntax-error storm. + """ + shell_bodies = [ + body.strip() for lang, body in _CMD_FENCE_RE.findall(reply) if lang.lower() in _SHELL_LANGS and body.strip() + ] + if shell_bodies: + return shell_bodies[-1] + tool_args = [a.strip() for a in _TOOLCALL_ARG_RE.findall(reply) if a.strip()] + if tool_args: + return tool_args[-1] + return None + + +def _is_task_complete(reply: str) -> bool: + """True when the policy signals completion (TASK_COMPLETE after its reasoning).""" + tail = reply.rsplit("", 1)[-1] + tail = re.sub(r"<\|.*?\|>", "", tail).strip().upper() + return tail.startswith("TASK_COMPLETE") or tail.endswith("TASK_COMPLETE") + + +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. +# The OpenEnv env server is the generic tbench2 server (it pulls the task's +# docker_image, copies the task dir into the container, and exposes exec); the +# SWE-bench task family reuses it wholesale and differs only in the task dirs. +def _load_env() -> 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} -> verify. + + The policy emits one shell command per turn (a ```bash block or the bare + reply), executed in the task's repo root; the loop ends when the policy stops + emitting a command, says TASK_COMPLETE, or hits OPENENV_MAX_TURNS. Scoring + runs the donor verifier (tests/test.sh) via an ``exec`` step and derives the + binary reward from its RESULT verdict (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 verify() 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 + + # Snapshot the pristine grader assets before the agent gets a turn, so the + # final grade reads tests the agent could not have tampered with (see + # _TESTS_SNAPSHOT). Best-effort: donor tasks may have no _TESTS_SRC, in + # which case grading falls back to _TESTS_SRC at eval time. + if _TESTS_SNAPSHOT: + try: + await env.step( + action_cls( + action_type="exec", + command=( + f"rm -rf {_TESTS_SNAPSHOT} && mkdir -p {_TESTS_SNAPSHOT} && " + f"cp -a {_TESTS_SRC}/. {_TESTS_SNAPSHOT}/ 2>/dev/null || true" + ), + ) + ) + except Exception: + pass + + 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 = _extract_command(reply) + if command is None: + if _is_task_complete(reply): + break + # No runnable command and not done: nudge for the right format and + # retry rather than executing the raw reply (prose/planning code). + convo.append({"role": "user", "content": _FORMAT_NUDGE}) + continue + + 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 env server (TB2_OUTPUT_DIR=/tmp/tbench2_env_runs) leaves a + # per-episode trial dir under that path after every episode, which fills + # the disk and trips ENOSPC. One episode holds the sandbox at a time, so + # it is safe to purge them here. + # + # MODE-DEPENDENT: this runs *inside the execution sandbox* via env.step. + # - Daytona mode (what this run used): the sandbox IS the execution + # filesystem, so /tmp/tbench2_env_runs is reachable and the purge is both + # effective and necessary -- verified to fix ENOSPC on the 10 GiB Daytona + # snapshot. + # - Docker mode: TB2_OUTPUT_DIR lives on the env-server *host* and is not + # bind-mounted into the task container, so this purge is a no-op there. + # The proper fix for docker mode is server-side cleanup on session close; + # we deliberately do not patch the upstream env server, so docker-mode + # disk pressure must be managed on the host (e.g. a reaper on the env + # host) instead. + # + # BUT the same dir also holds repo_cache/ (TB2_CACHE_DIR defaults to + # output_dir/repo_cache) -- the shared checkout that reset() clones once + # and every later episode reads its task from. A blanket `rm -rf .../*` + # wiped repo_cache too, collapsing effective concurrency and exploding step + # time. Preserve repo_cache; delete only the ephemeral per-trial dirs. + 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 verify() 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 SWE-bench-style episode via the trained policy.""" + request_kwargs = request_kwargs or {} + metadata = metadata or {} + + classes = _load_env() + 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 (dropped below). + # 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: + # The wall-clock cap covers capacity wait + reset + generation + eval, so a + # timeout can fire before the verifier ever ran -- exactly the no-verdict + # case we drop elsewhere. Scoring it 0.0 would inject a false negative when + # the cause is infra latency, not a task the agent legitimately failed. So + # drop the sample (return None); wait_for still cancels the coroutine and + # frees the sandbox slot, so the straggler is bounded either way. + logger.warning( + f"OpenEnv swebench episode exceeded {_MAX_ROLLOUT_TIME_S:.0f}s; terminating and dropping sample" + ) + return None + except Exception as e: + logger.error(f"OpenEnv swebench episode failed: {e}", exc_info=True) + return None + finally: + await policy.close() + + # No recoverable reward means the verifier 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 swebench episode produced no verdict " + f"(test.sh exit code={agent_metrics.get('testsh_rc')}); " + "infra/harness failure, dropping sample" + ) + return None + + # eval_report is intentionally empty: the reward-marker protocol echoes back + # only the scalar reward. The verifier's detailed output.json is written + # inside the sandbox at /logs/verifier/output.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_swebench/swebench_generate.py b/examples/experimental/openenv_swebench/swebench_generate.py new file mode 100644 index 0000000000..89b17cabf9 --- /dev/null +++ b/examples/experimental/openenv_swebench/swebench_generate.py @@ -0,0 +1,16 @@ +"""Reward function for the OpenEnv SWE-bench-style run. + +Task-agnostic: the agent function (``swebench_agent_function.run``) stores the +verifier-computed binary reward in ``sample.metadata["reward"]``; this just reads +it back. Wired via ``--custom-rm-path swebench_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_swebench/swebench_launch_common.py b/examples/experimental/openenv_swebench/swebench_launch_common.py new file mode 100644 index 0000000000..e0da5cb54b --- /dev/null +++ b/examples/experimental/openenv_swebench/swebench_launch_common.py @@ -0,0 +1,190 @@ +"""Shared launch helpers for the OpenEnv SWE-bench-style learning launchers. + +``run-openenv-swebench.py`` (GLM-4.7-Flash) is the launcher in this example; +sibling per-model launchers 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 + +# Agent-side overrides that swebench_agent_function reads straight from os.getenv +# but that base_env_vars does not derive from a launcher flag. They must be +# forwarded explicitly: with MILES_SCRIPT_EXTERNAL_RAY=1 the rollout runs on a +# remote Ray cluster that does NOT inherit the submission shell's environment, so +# an override set only in the shell would silently revert to its default on the +# workers. Forward each one that is actually set (membership, not truthiness, so +# an intentional empty value like OPENENV_CONDA_ENV="" still propagates). +_AGENT_PASSTHROUGH_ENV_VARS = ( + "OPENENV_TASK_WORKDIR", + "OPENENV_CONDA_ENV", + "OPENENV_MESSAGE_TIMEOUT_S", + "OPENENV_SWEBENCH_TESTS_SRC", + "OPENENV_SWEBENCH_TESTS_SNAPSHOT", + "OPENENV_EVAL_CMD", +) + + +class LaunchArgs(Protocol): + """The config fields the shared helpers read (satisfied by each launcher's ScriptArgs).""" + + prompt_data: str + num_rollout: int + 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 stale training/rollout processes to free GPU resources before a run. + + ASSUMES a disposable, single-tenant pod: this example colocates training and + rollout on all 8 GPUs of one node, so any lingering ``sglang`` / ``train.py`` + / ``MegatronTrain`` process is a leftover from a previous crashed run on this + same box and must be reaped or it holds the GPUs. To avoid collateral damage + on a shared host we scope the kill to the current user's own processes + (``pgrep -u ``) and still exclude this launcher and its parent. Set + ``OPENENV_SKIP_CLEANUP=1`` to disable entirely (e.g. when another of your jobs + is intentionally running on the same node).""" + if os.getenv("OPENENV_SKIP_CLEANUP") == "1": + print("Cleanup skipped (OPENENV_SKIP_CLEANUP=1)") + return + my_pid = os.getpid() + ppid = os.getppid() + euid = os.geteuid() + print(f"Cleanup starting (pid={my_pid}, ppid={ppid}, euid={euid})") + targets = ["sglang", "train.py", "MegatronTrain"] + exclude = f"grep -v '^{my_pid}$' | grep -v '^{ppid}$'" + for t in targets: + subprocess.run( + f"pgrep -u {euid} -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 " + f"--num-rollout {args.num_rollout} " + 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 swebench_agent_function.run " + "--custom-rm-path swebench_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]: + env = { + "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, + } + # Forward advertised agent-side overrides so they survive to a remote Ray + # cluster (MILES_SCRIPT_EXTERNAL_RAY=1); see _AGENT_PASSTHROUGH_ENV_VARS. + for name in _AGENT_PASSTHROUGH_ENV_VARS: + if name in os.environ: + env[name] = os.environ[name] + return env + + +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