From d8a0e5d9e10fcf8dd054c006f270c010b3eafe98 Mon Sep 17 00:00:00 2001 From: Tao Lin Date: Tue, 28 Jul 2026 14:43:51 -0400 Subject: [PATCH] nemo-gym: add an upstream-native connector example on the agent-function TITO chain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New examples/experimental/nemo-gym recipe integrating NVIDIA-NeMo/Gym as an external environment ecosystem at the agent-function layer, the same shape as the Harbor and OpenEnv connectors: the session server records every chat-completions turn losslessly, and a thin agent function POSTs each sample to the sandbox-backed mini_swe_agent_2 agent with policy_base_url set to the session URL (the per-request override proposed in NVIDIA-NeMo/Gym#2166; the README points at that PR's branch until it merges). - nemogym_agent_function.py: one POST /run per sample; sampling kwargs mapped onto responses_create_params (temperature / top_p / max_output_tokens), the only channel mini_swe_agent_2 reads; no miles imports so it loads on CPU-only machines - nemogym_generate.py: reward hook reading the environment's grade - eval_nemogym_via_api.py + tests/: no-GPU validation — offline unit tests of the /run contract, a golden scan (gold patch through the sandbox + SWE-bench harness, no model), and an API-policy scan that drives real episodes through the same policy_base_url override - run.py: the validated GPU launcher (requires MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1, which gates the dynamic registration of the agentic flags) - user-guide docs: nemo-gym page, environments tables (agent-function connector row; NeMo-Gym's own sandbox provider API supports Daytona) Validated end-to-end: offline contract tests; a golden scan scoring 1.0 in the official SWE-bench container; a DeepSeek API-policy episode; and a 4x H200 GRPO training smoke (Qwen3-4B-Instruct-2507, SWE-bench Verified) whose episodes ran in real task containers with the official harness grading them and rewards flowing back into training. Known limitations are documented in the README: the official swebench package lacks eval specs for several SWE-Gym repos (train on SWE-bench-family instances until that gap closes), and Qwen3 templates trip the tito_session_mismatch soft diagnostic (empty skeleton on re-rendered assistant history; engine-recorded tokens stay lossless). This supersedes the fork-based NeMo-Gym integration removed in #1918: same ecosystem, now driven through upstream NVIDIA-NeMo/Gym with no submodules and lossless token recording. Co-Authored-By: Claude Fable 5 --- docs/docs.json | 3 +- docs/user-guide/environments.md | 3 +- docs/user-guide/nemo-gym.md | 57 ++++ examples/experimental/nemo-gym/README.md | 249 ++++++++++++++++++ .../nemo-gym/download_and_process_data.py | 92 +++++++ .../nemo-gym/eval_nemogym_via_api.py | 132 ++++++++++ .../nemo-gym/nemogym_agent_function.py | 151 +++++++++++ .../experimental/nemo-gym/nemogym_generate.py | 18 ++ examples/experimental/nemo-gym/run.py | 210 +++++++++++++++ .../tests/test_nemogym_agent_function.py | 152 +++++++++++ 10 files changed, 1065 insertions(+), 2 deletions(-) create mode 100644 docs/user-guide/nemo-gym.md create mode 100644 examples/experimental/nemo-gym/README.md create mode 100755 examples/experimental/nemo-gym/download_and_process_data.py create mode 100644 examples/experimental/nemo-gym/eval_nemogym_via_api.py create mode 100644 examples/experimental/nemo-gym/nemogym_agent_function.py create mode 100644 examples/experimental/nemo-gym/nemogym_generate.py create mode 100644 examples/experimental/nemo-gym/run.py create mode 100644 examples/experimental/nemo-gym/tests/test_nemogym_agent_function.py diff --git a/docs/docs.json b/docs/docs.json index 102b87c1a10..dfbc701e07e 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -180,7 +180,8 @@ "root": "user-guide/environments", "pages": [ "user-guide/harbor", - "user-guide/openenv" + "user-guide/openenv", + "user-guide/nemo-gym" ] } ] diff --git a/docs/user-guide/environments.md b/docs/user-guide/environments.md index b50bdb57236..878a27b58e3 100644 --- a/docs/user-guide/environments.md +++ b/docs/user-guide/environments.md @@ -20,6 +20,7 @@ where the environment itself comes from: |---|---| | [Harbor](/user-guide/harbor) | agent function | | [OpenEnv](/user-guide/openenv) | agent function | +| [NeMo-Gym](/user-guide/nemo-gym) | agent function | | [Strands Agents](https://github.com/radixark/miles/tree/main/examples/experimental/strands_sglang) | generate function | | [τ-bench](https://github.com/radixark/miles/tree/main/examples/experimental/tau-bench) | generate function | @@ -28,7 +29,7 @@ Sandbox providers are a different axis: they provision the task containers | Sandbox provider | Used within | |---|---| -| [Daytona](https://www.daytona.io/) | OpenEnv, Harbor | +| [Daytona](https://www.daytona.io/) | OpenEnv, Harbor, NeMo-Gym | All external ecosystem support is experimental. diff --git a/docs/user-guide/nemo-gym.md b/docs/user-guide/nemo-gym.md new file mode 100644 index 00000000000..7a7a6a87453 --- /dev/null +++ b/docs/user-guide/nemo-gym.md @@ -0,0 +1,57 @@ +--- +title: NeMo-Gym +description: Train on NVIDIA NeMo-Gym environments through the agent-function extension point. +--- + +[NeMo-Gym](https://github.com/NVIDIA-NeMo/Gym) is NVIDIA's RL environment +ecosystem: environments are HTTP *resources servers* (code execution, search, +SWE tasks, ...) paired with *agents* that drive an episode end-to-end and +grade it. Task containers run through NeMo-Gym's own sandbox provider API +(`nemo_gym.sandbox`) — Docker locally, or Daytona / Apptainer / ECS Fargate / +OpenSandbox — selected by config, no agent changes. + +Miles integrates NeMo-Gym as an +[agent-function integration](/user-guide/environments): per sample, the agent +function POSTs the task to a NeMo-Gym agent server's `/run` endpoint with +`policy_base_url` set to the session's OpenAI-compatible URL. NeMo-Gym runs +its agent harness (mini-swe-agent v2 in `mini_swe_agent_2`) against that URL, +so Miles' session server records every turn losslessly (token ids, logprobs, +loss masks — see [Rollout Endpoints](/user-guide/rollout-endpoints)); NeMo-Gym +grades the episode itself and the grade enters training through a custom +reward hook reading `sample.metadata["reward"]`. + +The per-request `policy_base_url` override is proposed upstream in +[NVIDIA-NeMo/Gym#2166](https://github.com/NVIDIA-NeMo/Gym/pull/2166); until it +merges, run the NeMo-Gym server from that PR's branch (upstream main plus one +small commit pair). + +## Try it + +The maintained recipe is **SWE-bench GRPO with mini-swe-agent** in +[`examples/experimental/nemo-gym`](https://github.com/radixark/miles/tree/main/examples/experimental/nemo-gym). +In short: + +1. **Environment side** — on any docker-capable host, clone NeMo-Gym (the + PR branch above until #2166 merges) and start the `mini_swe_agent_2` + responses-API agent server with the docker sandbox provider config. +2. **Data** — convert SWE-bench Verified to Miles prompt data with + `download_and_process_data.py`; the task instance rides in each sample's + `metadata`. +3. **Training side** — point `NEMO_GYM_URL` at the agent server and launch + `run.py` (requires `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`, which the + launcher sets), wiring the chain: + +```bash +--custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate +--custom-agent-function-path nemogym_agent_function.run +--custom-rm-path nemogym_generate.reward_func +--use-session-server +``` + +The recipe is validated end-to-end: golden and API-policy scans on a real +docker host, plus a 4-GPU GRPO training smoke whose episodes ran in real task +containers with the SWE-bench harness grading them. Follow the +[recipe README](https://github.com/radixark/miles/blob/main/examples/experimental/nemo-gym/README.md) +for the NeMo-Gym server setup, no-GPU validation (golden scan / API-policy +scan), the launch walkthrough, and known limitations (SWE-Gym eval specs, +Qwen3 template soft-mismatch diagnostics). diff --git a/examples/experimental/nemo-gym/README.md b/examples/experimental/nemo-gym/README.md new file mode 100644 index 00000000000..44e84d35c15 --- /dev/null +++ b/examples/experimental/nemo-gym/README.md @@ -0,0 +1,249 @@ +# SWE-agent training via NeMo-Gym + +## Introduction + +This example trains a SWE agent with Miles using NVIDIA's +[NeMo-Gym](https://github.com/NVIDIA-NeMo/Gym) as the environment ecosystem: +NeMo-Gym's sandbox-backed `mini_swe_agent_2` agent runs the +[mini-swe-agent](https://github.com/SWE-agent/mini-swe-agent) v2 harness inside +per-task SWE-bench containers (via the `nemo_gym.sandbox` provider API) and +grades every episode with the official SWE-bench harness; Miles owns training, +batch orchestration, and lossless token recording. + +NeMo-Gym plugs in at the **agent function** layer (see the +[Environments guide](../../../docs/user-guide/environments.md)), the same +shape as the Harbor and OpenEnv connectors: + +``` +Miles trainer ── session server (records every chat-completions turn: token + │ ids + logprobs + loss masks, no re-tokenization) + │ per-sample POST /run { task fields, policy_base_url = session URL } + ▼ +NeMo-Gym responses-API agent server (mini_swe_agent_2) + │ runs mini-swe-agent v2 against policy_base_url, + │ per-task container via a nemo_gym.sandbox provider (docker here; + │ daytona / apptainer / ecs_fargate / opensandbox also exist) + ▼ +reward (official SWE-bench harness) ──► sample.metadata ──► reward hook +``` + +- `nemogym_agent_function.py` — the connector: one `/run` POST per sample. +- `nemogym_generate.py` — reward hook (reads the NeMo-Gym grade from + `sample.metadata["reward"]`). +- `eval_nemogym_via_api.py`, `tests/` — no-GPU validation tooling (below). + +The per-request `policy_base_url` override this example relies on is proposed +upstream in [NVIDIA-NeMo/Gym#2166](https://github.com/NVIDIA-NeMo/Gym/pull/2166). +Until it merges, run the NeMo-Gym server from the PR branch +(`nblintao/Gym@mini-swe-agent-per-request-policy-url`, upstream main + that +one commit pair); afterwards, use upstream directly. + +## Validation status + +Validated end-to-end (2026-07-28) — the commands in this README are the exact +ones used: + +- offline contract tests: 7/7 pass; +- golden scan: gold patch through the official + `swebench/sweb.eval.x86_64.*` container scored **reward 1.0**; +- API-policy scan: DeepSeek drove a full episode through the + `policy_base_url` override — patch applied, FAIL_TO_PASS 4/5, a legitimate + reward 0.0; +- **GPU training smoke** (`run.py` defaults, 4x H200, + Qwen3-4B-Instruct-2507, SWE-bench Verified prompts): 3 synchronous GRPO + steps completed twice; every episode ran mini-swe-agent v2 in a real task + container on the NeMo-Gym host, the SWE-bench harness executed the task's + full test suite (e.g. 175/175 PASS_TO_PASS on an unresolved attempt), and + the grade flowed back into `rollout/raw_reward`. + +Two known limitations from the smoke run: + +- A 4B policy solves none of these tasks, so rewards were uniformly 0 — + GRPO then has zero advantage (`rollout/zero_std` fires). That's a model + capability floor, not a pipeline defect; expect the same until you use a + stronger policy or an easier task pool. +- `rollout/tito_session_mismatch_rate` reads 1.0 with this model: Qwen3 + chat templates insert an empty `` skeleton when re-rendering + assistant history, which the engine's actual output never contains. It is a + soft diagnostic — training tokens and loss masks come from the engine's + recorded token ids, which stay lossless — and is a property of the + model-family template, not of this connector. + +## Setting up the NeMo-Gym server + +Any docker-capable host works: a CPU box next to the cluster, or a container +beside the trainer (mount `/var/run/docker.sock` and share a docker network +with the trainer — that variant is not validated here). Set `NEMO_GYM_URL` to +wherever the server listens. + +```bash +# Until NVIDIA-NeMo/Gym#2166 merges; afterwards clone NVIDIA-NeMo/Gym instead. +git clone -b mini-swe-agent-per-request-policy-url https://github.com/nblintao/Gym.git +cd Gym + +curl -LsSf https://astral.sh/uv/install.sh | sh +source $HOME/.local/bin/env +uv venv --python 3.12 && source .venv/bin/activate +uv sync --extra dev +# Do NOT install the agent's requirements.txt into this venv. `gym env start` +# builds a per-server venv from it automatically; installing it here bumps +# shared pins (e.g. openai) past nemo-gym's caps, and the injected versions +# then make the child venv unresolvable. + +# Global config. The model-server entry must boot but receives no policy +# traffic in this setup — every /run carries its own policy_base_url override +# pointing at a Miles session URL. policy_model_name is the model name the +# harness sends on each request (any string Miles' router accepts). +echo "policy_base_url: http://localhost:9/v1 +policy_api_key: dummy +policy_model_name: model +default_host: 0.0.0.0" > env.yaml +``` + +Start the server, composing the agent config, the docker sandbox provider +config, and a model server config: + +```bash +gym env start \ + --config responses_api_agents/mini_swe_agent_2/configs/mini_swe_agent_2.yaml \ + --config nemo_gym/sandbox/providers/docker/configs/docker.yaml \ + --model-type vllm_model \ + '++mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.port=12000' \ + '++mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.concurrency=16' +``` + +The first spin-up resolves each server's own venv, which takes a few minutes; +the server is ready when the startup table lists `mini_swe_agent_2` on +port 12000 and uvicorn reports it running. + +## Preparing data + +On the trainer, download the task instances and convert them to Miles' prompt +data format. The validated smoke run uses **SWE-bench Verified**: + +```bash +cd miles/examples/experimental/nemo-gym +python download_and_process_data.py --input princeton-nlp/SWE-bench_Verified \ + --split test --subset verified --output /root/swe_verified.jsonl +``` + +Each row keeps the full SWE-bench-format instance (`instance_id`, `repo`, +`base_commit`, `problem_statement`, ...) in `metadata`, plus `subset` / +`split` — the agent function forwards all of it in the `/run` body, and +NeMo-Gym selects the per-task image from it. + +**SWE-Gym caveat**: `--input SWE-Gym/SWE-Gym --subset gym` produces the +training dataset this recipe ultimately targets (per-task images from +`docker.io/xingyaoww/...`), and episodes run fine — but the official +`swebench` package `mini_swe_agent_2` scores with does not carry eval specs +for several SWE-Gym repos (`KeyError: 'getmoto/moto'` at +`make_test_spec`), so those episodes error at grading and score 0. Until +SWE-Gym eval specs are available in that path (upstream gap), train on +SWE-bench-family instances or filter SWE-Gym to repos the `swebench` package +knows. + +## Wiring it into training + +The launcher is [`run.py`](run.py) (4 GPUs, smoke-scale defaults — scale up +--num-rollout / batch sizes for real training). Its prepare step downloads the +HF checkpoint and converts it to torch_dist on first run (`--skip-prepare` to +skip): + +```bash +export NEMO_GYM_URL="http://:12000" +# Only if the NeMo-Gym host cannot resolve the trainer's hostname (e.g. it +# reaches the trainer over a tailnet): +export MILES_ROUTER_EXTERNAL_HOST="" +python examples/experimental/nemo-gym/run.py +``` + +To wire the connector into a different launch script, the essential pieces +are this example's directory on `PYTHONPATH`, +`MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1` in the environment (it gates the +dynamic registration of the agentic flags below — without it train.py fails +with "unrecognized arguments"), and: + +```bash +--prompt-data /root/swe_verified.jsonl +--input-key prompt +--metadata-key metadata +--max-seq-len 16384 + +--custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate +--custom-agent-function-path nemogym_agent_function.run +--custom-rm-path nemogym_generate.reward_func +--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_no_aborted +--use-session-server +--session-server-ip 0.0.0.0 # listen on all interfaces for the dial-back +--tito-model qwen3 # match your policy model's TITO family +``` + +Per sample, `agentic_tool_call` opens a session on Miles' session server and +hands its OpenAI-compatible URL to `nemogym_agent_function.run`, which POSTs +the task to the NeMo-Gym server's `/run` with `policy_base_url` set to that +session URL and the sampling settings mapped onto `responses_create_params` +(`temperature`, `top_p`, `max_output_tokens`). NeMo-Gym's mini-swe-agent v2 +then talks to the policy exclusively through the session URL (litellm chat +completions), so Miles records every turn losslessly — token ids, logprobs, +and loss masks come from the session server, not from re-tokenizing message +text. The episode grade rides back in the `/run` response (`reward`, with the +SWE-bench eval report in `metadata` → `sample.metadata["eval_report"]`) and +enters training through `sample.metadata["reward"]`. + +Episodes that fail before the first model call produce no session records; the +sample is marked aborted and `check_no_aborted` drops its group from training. + +## Validating without a GPU + +Everything except the session server and the training loop can be validated +on CPU-only machines, in three independent layers (all three pass as of +2026-07-28, see [Validation status](#validation-status)): + +1. **Offline unit tests** — the `/run` request contract, response mapping, + failure semantics, and the data conversion. No network, no docker: + + ```bash + pytest examples/experimental/nemo-gym/tests/ -q + ``` + +2. **Golden scan** — the sandbox + per-task image + SWE-bench harness chain, + with no model involved at all: start the server with + `'++mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.run_golden=true'` + appended to the `gym env start` command, then + + ```bash + python eval_nemogym_via_api.py --input /root/swe_verified.jsonl --golden --limit 5 + ``` + + Every gold patch must score reward 1.0; the script exits non-zero + otherwise. + +3. **API-policy scan** — a real model drives full episodes through the same + `policy_base_url` override the trainer uses (so this also exercises the + NVIDIA-NeMo/Gym#2166 field end-to-end). Start the server *without* the + golden override and with `policy_model_name` set to the API model name + (e.g. `deepseek-chat`) in `env.yaml`, then: + + ```bash + export DEEPSEEK_API_KEY=... # or OPENAI_API_KEY + python eval_nemogym_via_api.py --input /root/swe_verified.jsonl --limit 2 \ + --policy-base-url https://api.deepseek.com/v1 + ``` + +## Troubleshooting + +1. `train.py: error: unrecognized arguments: --max-seq-len + --custom-agent-function-path`: `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1` is + missing from the training environment (it must reach the ray job — run.py + sets it via the ray runtime env). +2. `mini_swe_agent_2` dies at spin-up with an unresolvable-dependency error + (`openai==X` vs `nemo-gym depends on openai<=Y`): the main venv has extra + packages installed. Recreate it with `uv sync --extra dev` only — see the + setup note above. +3. Slow episodes are usually docker pulls (each task has its own image, + fetched on first use) or the in-container SWE-bench evaluation + (server-side `eval_timeout`, default 1800s). `NEMO_GYM_RUN_TIMEOUT` + (default 3600s) caps one episode end-to-end on the Miles side. +4. A failed episode surfaces as `sample.metadata["eval_report"]["error"]` with + a traceback from the NeMo-Gym server — check there before digging into + server logs. diff --git a/examples/experimental/nemo-gym/download_and_process_data.py b/examples/experimental/nemo-gym/download_and_process_data.py new file mode 100755 index 00000000000..9060c4a57ec --- /dev/null +++ b/examples/experimental/nemo-gym/download_and_process_data.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +"""Download and process data to Miles format.""" + +import argparse +import json +import tempfile +from pathlib import Path +from datasets import load_dataset + + +def convert_to_miles_format( + input_path: str, output_path: str, limit: int = None, split: str = "train", subset: str = "gym" +): + """Convert JSONL to Miles format. + + Args: + input_path: Path to input JSONL file + output_path: Path to output JSONL file in Miles format + limit: Optional limit on number of samples + split: Dataset split name (used in metadata) + subset: NeMo-Gym subset controlling per-task image selection and eval + ("gym" -> SWE-Gym images, "verified" -> official SWE-bench images) + """ + count = 0 + with open(input_path) as fin, open(output_path, "w") as fout: + for line in fin: + if limit and count >= limit: + break + + instance = json.loads(line) + + # Add subset and split to metadata for Gym API + metadata = dict(instance) + metadata["subset"] = subset + metadata["split"] = split + + miles_sample = { + "prompt": instance.get("problem_statement", ""), + "metadata": metadata, + } + + fout.write(json.dumps(miles_sample) + "\n") + count += 1 + + print(f"Converted {count} samples: {input_path} -> {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="Download HuggingFace dataset and convert to Miles format") + parser.add_argument("--input", type=str, required=True, help="HuggingFace dataset path or local JSONL file") + parser.add_argument("--output", type=str, required=True, help="Output JSONL file path") + parser.add_argument( + "--split", type=str, default="train", help="Dataset split (default: train, only for HF datasets)" + ) + parser.add_argument("--limit", type=int, help="Limit number of samples") + parser.add_argument( + "--subset", type=str, default="gym", help='NeMo-Gym subset: "gym" (SWE-Gym) or "verified" (SWE-bench Verified)' + ) + + args = parser.parse_args() + + input_path = Path(args.input) + + if input_path.exists() and input_path.suffix == ".jsonl": + print(f"Processing local file: {args.input}") + convert_to_miles_format(args.input, args.output, args.limit, args.split, args.subset) + else: + print(f"Loading HuggingFace dataset: {args.input} (split={args.split})") + ds = load_dataset(args.input, split=args.split) + + if args.limit: + ds = ds.select(range(min(args.limit, len(ds)))) + + tmp_path = None + try: + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as tmp: + tmp_path = tmp.name + + print(f"Downloading to temporary file: {tmp_path}") + ds.to_json(tmp_path) + + print(f"Converting to Miles format: {args.output}") + convert_to_miles_format(tmp_path, args.output, args.limit, split=args.split, subset=args.subset) + finally: + if tmp_path and Path(tmp_path).exists(): + Path(tmp_path).unlink() + + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/examples/experimental/nemo-gym/eval_nemogym_via_api.py b/examples/experimental/nemo-gym/eval_nemogym_via_api.py new file mode 100644 index 00000000000..5ccffab1011 --- /dev/null +++ b/examples/experimental/nemo-gym/eval_nemogym_via_api.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +"""Validate the NeMo-Gym leg without a GPU trainer (golden / API-policy scan). + +Drives a running mini_swe_agent_2 server through the same ``/run`` contract +the miles agent function uses, so a pass here validates everything except the +session server and training itself: + + Golden scan (no model at all — sandbox + image + SWE-bench harness): + # start the server with the golden override: + # gym env start ... '++mini_swe_agent_2.responses_api_agents.mini_swe_agent_2.run_golden=true' + python eval_nemogym_via_api.py --input swe_verified.jsonl --golden --limit 5 + + API-policy scan (a real model drives episodes via the policy_base_url + override — the exact NVIDIA-NeMo/Gym#2166 field the trainer relies on): + export DEEPSEEK_API_KEY=$(cat ~/.config/deepseek/api_key) + python eval_nemogym_via_api.py --input swe_verified.jsonl --limit 2 \ + --policy-base-url https://api.deepseek.com/v1 + +The server sends its own configured model name on every policy request, so +start it with ``policy_model_name`` (env.yaml) set to the name the policy +endpoint expects (e.g. ``deepseek-chat``). + +Input rows are miles prompt data (``{"prompt": ..., "metadata": {instance}}``, +the output of download_and_process_data.py) or raw SWE-bench instances. +""" + +import argparse +import asyncio +import json +import os +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from nemogym_agent_function import build_responses_create_params, post_json # noqa: E402 + + +def _load_instances(path: str, limit: int | None) -> list[dict]: + instances = [] + with open(path) as f: + for line in f: + if limit is not None and len(instances) >= limit: + break + row = json.loads(line) + instance = row.get("metadata", row) + instance.setdefault("subset", "gym") + instance.setdefault("split", "train") + instances.append(instance) + return instances + + +async def _run_one(args, instance: dict) -> dict: + request = { + **instance, + "responses_create_params": build_responses_create_params( + {"temperature": args.temperature, "top_p": args.top_p, "max_tokens": args.max_tokens} + ), + } + if not args.golden: + request["policy_base_url"] = args.policy_base_url + if args.policy_api_key: + request["policy_api_key"] = args.policy_api_key + + instance_id = instance.get("instance_id", "?") + t0 = time.monotonic() + try: + response = await asyncio.wait_for(post_json(f"{args.nemo_gym_url}/run", request), timeout=args.timeout) + reward = float(response.get("reward", 0.0)) + eval_report = response.get("metadata", {}) or {} + error = eval_report.get("error") + except Exception as e: # noqa: BLE001 - a scan reports failures, it doesn't crash + reward, eval_report, error = 0.0, {}, str(e) + elapsed = time.monotonic() - t0 + status = "OK " if (error is None and (not args.golden or reward == 1.0)) else "FAIL" + suffix = f" error={error}" if error else "" + print(f"[{status}] {instance_id} reward={reward} {elapsed:.0f}s{suffix}", flush=True) + return {"instance_id": instance_id, "reward": reward, "error": error, "eval_report": eval_report} + + +async def _main(args) -> int: + instances = _load_instances(args.input, args.limit) + print(f"{'golden' if args.golden else 'api-policy'} scan: {len(instances)} instance(s) via {args.nemo_gym_url}") + + sem = asyncio.Semaphore(args.concurrency) + + async def bounded(instance): + async with sem: + return await _run_one(args, instance) + + results = await asyncio.gather(*(bounded(i) for i in instances)) + + if args.output: + with open(args.output, "w") as f: + for r in results: + f.write(json.dumps(r) + "\n") + + rewards = [r["reward"] for r in results] + failures = [r for r in results if r["error"] or (args.golden and r["reward"] != 1.0)] + print(f"\nmean reward {sum(rewards) / max(len(rewards), 1):.2f} over {len(rewards)}; {len(failures)} failure(s)") + if args.golden and failures: + print("golden scan FAILED: every gold patch must score 1.0") + return 1 + return 0 + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--input", required=True, help="miles prompt data or raw SWE-bench jsonl") + parser.add_argument("--nemo-gym-url", default=os.getenv("NEMO_GYM_URL", "http://localhost:12000")) + parser.add_argument( + "--golden", action="store_true", help="server must run with run_golden=true; expect reward 1.0" + ) + parser.add_argument("--policy-base-url", default=None, help="OpenAI-compatible policy endpoint for the episode") + parser.add_argument("--policy-api-key", default=os.getenv("DEEPSEEK_API_KEY") or os.getenv("OPENAI_API_KEY")) + parser.add_argument("--limit", type=int, default=None) + parser.add_argument("--concurrency", type=int, default=2) + parser.add_argument("--temperature", type=float, default=0.6) + parser.add_argument("--top-p", type=float, default=0.95) + parser.add_argument("--max-tokens", type=int, default=8192) + parser.add_argument("--timeout", type=float, default=float(os.getenv("NEMO_GYM_RUN_TIMEOUT", "3600"))) + parser.add_argument("--output", default=None, help="write per-instance results jsonl here") + args = parser.parse_args() + + if not args.golden and not args.policy_base_url: + parser.error("either --golden or --policy-base-url is required") + + sys.exit(asyncio.run(_main(args))) + + +if __name__ == "__main__": + main() diff --git a/examples/experimental/nemo-gym/nemogym_agent_function.py b/examples/experimental/nemo-gym/nemogym_agent_function.py new file mode 100644 index 00000000000..58aee0f1426 --- /dev/null +++ b/examples/experimental/nemo-gym/nemogym_agent_function.py @@ -0,0 +1,151 @@ +"""NeMo-Gym <-> miles adapter (agent function). + +Targets upstream NVIDIA-NeMo/Gym's sandbox-backed ``mini_swe_agent_2`` agent, +which requires the per-request policy endpoint override from +https://github.com/NVIDIA-NeMo/Gym/pull/2166 (until it merges, run the +NeMo-Gym server from the PR branch ``nblintao/Gym@mini-swe-agent-per-request-policy-url``). + +miles calls ``run`` once per sample via +``--custom-agent-function-path nemogym_agent_function.run`` (with +``--custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate``). +Each call POSTs the task to the NeMo-Gym agent server's ``/run`` endpoint, +handing over the session's OpenAI-compatible URL as ``policy_base_url``. +NeMo-Gym runs mini-swe-agent v2 in a ``nemo_gym.sandbox`` container (docker / +daytona / apptainer / ecs_fargate / opensandbox providers) against that URL, +so every model call goes through miles' session server and is recorded +losslessly (token ids + logprobs + loss masks) — no re-tokenization. + +The NeMo-Gym environment grades the episode itself (SWE-bench harness); the +returned dict is merged into ``sample.metadata`` so the reward hook +(``--custom-rm-path nemogym_generate.reward_func``) can read +``metadata["reward"]``. + +Env vars: + NEMO_GYM_URL base URL of the NeMo-Gym agent server + (default: http://localhost:12000). + NEMO_GYM_RUN_TIMEOUT hard wall-clock cap in seconds for one /run call + (default: 3600). SWE episodes pull per-task docker images on + first use, which can dominate early rollouts. + MILES_ROUTER_EXTERNAL_HOST optional host rewrite for the session URL when + the NeMo-Gym server cannot resolve the trainer's hostname + (e.g. it runs outside the trainer's docker network). +""" + +import asyncio +import logging +import os +import random +from typing import Any +from urllib.parse import urlparse, urlunparse + +import httpx + +logger = logging.getLogger(__name__) + +# Deliberately no miles imports: importing miles pulls in torch, and this +# adapter (plus its offline tests and the eval_nemogym_via_api scan) must load +# on CPU-only machines too. + +_POST_ATTEMPTS = 3 +_POST_BACKOFF_S = (1.0, 5.0) + + +async def post_json(url: str, payload: dict[str, Any]) -> dict[str, Any]: + """POST JSON and return the JSON response, retrying transport errors. + + No HTTP timeout here: the caller bounds the whole episode with + asyncio.wait_for (a /run legitimately takes minutes). + """ + async with httpx.AsyncClient(timeout=None) as client: + for attempt in range(_POST_ATTEMPTS): + try: + response = await client.post(url, json=payload) + response.raise_for_status() + return response.json() + except httpx.TransportError: + if attempt == _POST_ATTEMPTS - 1: + raise + await asyncio.sleep(random.uniform(*_POST_BACKOFF_S)) + + +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 build_responses_create_params(request_kwargs: dict[str, Any]) -> dict[str, Any]: + """Map miles' chat-completions sampling kwargs onto NeMo-Gym's Responses-API params. + + mini_swe_agent_2 reads sampling settings exclusively from + ``responses_create_params`` (temperature / top_p / max_output_tokens, see + upstream ``_responses_create_params_to_model_kwargs``). + """ + params: dict[str, Any] = {"input": []} + for key in ("temperature", "top_p"): + if request_kwargs.get(key) is not None: + params[key] = request_kwargs[key] + if request_kwargs.get("max_tokens") is not None: + params["max_output_tokens"] = request_kwargs["max_tokens"] + return params + + +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 task instance via the NeMo-Gym mini_swe_agent_2 server. + + Returns the reward dict to merge into sample metadata, or None on a + transport failure (timeout, unreachable server). On None the recorded + session still becomes a sample; the reward hook then scores it 0.0 via its + default. Episodes that never reach the model produce no session records + and are ABORTED by the generate layer, so + ``--dynamic-sampling-filter-path .. check_no_aborted`` drops their group. + """ + metadata = metadata or {} + request_kwargs = request_kwargs or {} + + nemo_gym_url = os.getenv("NEMO_GYM_URL", "http://localhost:12000") + timeout_s = float(os.getenv("NEMO_GYM_RUN_TIMEOUT", "3600")) + + # The SWE-bench-format instance fields (instance_id, repo, base_commit, + # problem_statement, subset, split, ...) ride in metadata straight from the + # prompt data and must sit at the top level of the run body: the server + # uses body.model_dump() as the instance dict (image selection, eval). + request: dict[str, Any] = { + **metadata, + "responses_create_params": build_responses_create_params(request_kwargs), + "policy_base_url": _resolve_session_url(base_url), + } + + try: + response = await asyncio.wait_for( + post_json(f"{nemo_gym_url}/run", request), + timeout=timeout_s, + ) + except asyncio.TimeoutError: + logger.error(f"NeMo-Gym /run timed out after {timeout_s:.0f}s") + return None + except asyncio.CancelledError: + logger.warning("NeMo-Gym /run cancelled (sibling task failure?)") + return None + except Exception as e: + logger.error(f"NeMo-Gym /run failed: {e}") + return None + + return { + "reward": response.get("reward", 0.0), + # The SWE-bench eval report (tests_status, patch_successfully_applied, + # or {"error": ...} when the episode failed server-side) rides in the + # response's `metadata`. + "eval_report": response.get("metadata", {}) or {}, + } diff --git a/examples/experimental/nemo-gym/nemogym_generate.py b/examples/experimental/nemo-gym/nemogym_generate.py new file mode 100644 index 00000000000..62b8e060c51 --- /dev/null +++ b/examples/experimental/nemo-gym/nemogym_generate.py @@ -0,0 +1,18 @@ +"""NeMo-Gym example: reward hook. + +The generate function is provided by + miles.rollout.generate_hub.agentic_tool_call.generate +with --custom-agent-function-path pointing to nemogym_agent_function.run. + +Reward is pre-computed by the NeMo-Gym environment (SWE-bench harness) during +the episode and stored in sample.metadata["reward"]. +""" + +from miles.utils.types import Sample + + +async def reward_func(args, samples: Sample | list[Sample], **kwargs) -> float | list[float]: + """Reward is pre-computed by the NeMo-Gym environment during generate().""" + 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/nemo-gym/run.py b/examples/experimental/nemo-gym/run.py new file mode 100644 index 00000000000..25a8e43fc80 --- /dev/null +++ b/examples/experimental/nemo-gym/run.py @@ -0,0 +1,210 @@ +"""NeMo-Gym launcher (Qwen3-4B-Instruct-2507): Miles <-> mini_swe_agent_2 orchestration. + +Defaults are the exact configuration of the validated smoke run (4x H200, +2026-07-28): 3 GRPO steps at tiny scale against a NeMo-Gym server running the +docker sandbox provider. Scale up --num-rollout / batch sizes for real +training. + +Usage: + NEMO_GYM_URL=http://:12000 python run.py + python run.py --mode debug_rollout_only + python run.py --skip-prepare --prompt-data /my/data.jsonl +""" + +import os +import subprocess +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +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" + megatron_model_type: str = "qwen3-4B-Instruct-2507" + num_gpus_per_node: int = 4 + megatron_path: str = "/root/Megatron-LM" + + # Paths + skip_prepare: bool = False + base_dir: str = "/root" + model_name: str = "Qwen3-4B-Instruct-2507" + hf_checkpoint: str = "Qwen/Qwen3-4B-Instruct-2507" + ref_load: str = "/root/Qwen3-4B-Instruct-2507_torch_dist" + save_dir: str = "/root/Qwen3-4B-Instruct-2507_nemogym/" + prompt_data: str = "/root/swe_verified.jsonl" + + # Training settings (validated smoke scale) + max_seq_len: int = 16384 + rollout_max_response_len: int = 4096 + num_rollout: int = 3 + rollout_batch_size: int = 2 + n_samples_per_prompt: int = 4 + global_batch_size: int = 8 + save_interval: int = 1000 + + # NeMo-Gym settings + nemo_gym_url: str = os.environ.get("NEMO_GYM_URL", "http://localhost:12000") + # Trainer address reachable from the NeMo-Gym host; only needed when that + # host cannot resolve the trainer's hostname (e.g. it dials back over a + # tailnet). + router_external_host: str = os.environ.get("MILES_ROUTER_EXTERNAL_HOST", "") + + +def cleanup(): + """Kill old Ray jobs and stale processes to free GPU resources.""" + my_pid = os.getpid() + ppid = os.getppid() + exclude = f"grep -v '^{my_pid}$' | grep -v '^{ppid}$'" + for t in ["sglang", "train.py", "MegatronTrain"]: + subprocess.run( + f"pgrep -f '{t}' | {exclude} | xargs -r kill 2>/dev/null || true", + shell=True, + ) + time.sleep(5) + + +def prepare(args: ScriptArgs): + """Convert the 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} " + f"--save-interval {args.save_interval} " + ) + + rollout_args = ( + 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 " + f"--rollout-max-response-len {args.rollout_max_response_len} " + f"--max-seq-len {args.max_seq_len} " + f"--global-batch-size {args.global_batch_size} " + "--balance-data " + ) + + grpo_args = ( + "--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 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + perf_args = ( + "--tensor-model-parallel-size 2 " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + f"--max-tokens-per-gpu {args.max_seq_len} " + ) + + sglang_args = "--rollout-num-gpus-per-engine 1 --sglang-mem-fraction-static 0.7 " + + agent_args = ( + "--custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate " + "--custom-agent-function-path nemogym_agent_function.run " + "--custom-rm-path nemogym_generate.reward_func " + "--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_no_aborted " + "--use-session-server " + # 0.0.0.0 so the NeMo-Gym host can dial in on any interface (e.g. a + # tailnet address); internal calls resolve it to localhost. + "--session-server-ip 0.0.0.0 " + "--session-server-port 30000 " + "--tito-model qwen3 " + ) + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--colocate " + "--actor-num-nodes 1 " + f"--actor-num-gpus-per-node {args.num_gpus_per_node} " + ) + + debug_args = "--debug-rollout-only " if args.mode == "debug_rollout_only" else "" + + train_args = ( + f"{ckpt_args}" + f"{rollout_args}" + f"{grpo_args}" + f"{optimizer_args}" + f"{perf_args}" + f"{sglang_args}" + f"{agent_args}" + f"{misc_args}" + f"{debug_args}" + ) + + extra_env_vars = { + "PYTHONPATH": f"{args.megatron_path}:{SCRIPT_DIR}:{U.repo_base_dir}", + # Gates the dynamic registration of the agentic flags above; without + # it train.py fails with "unrecognized arguments". + "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1", + "NEMO_GYM_URL": args.nemo_gym_url, + } + if args.router_external_host: + extra_env_vars["MILES_ROUTER_EXTERNAL_HOST"] = args.router_external_host + + 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): + cleanup() + if not args.skip_prepare: + prepare(args) + execute(args) + + +if __name__ == "__main__": + typer.run(main) diff --git a/examples/experimental/nemo-gym/tests/test_nemogym_agent_function.py b/examples/experimental/nemo-gym/tests/test_nemogym_agent_function.py new file mode 100644 index 00000000000..f6cfccd3055 --- /dev/null +++ b/examples/experimental/nemo-gym/tests/test_nemogym_agent_function.py @@ -0,0 +1,152 @@ +"""Offline unit tests for the NeMo-Gym adapter (no network, no GPU). + +Not collected by the repo-level pytest run (testpaths = ./tests); run manually +when touching the adapter: + + pytest examples/experimental/nemo-gym/tests/ -q + +Covers the /run request contract the mini_swe_agent_2 server expects (instance +fields at the top level, policy_base_url override, sampling mapped onto +responses_create_params), the response mapping, and the failure semantics +(transport failure -> None so the sample keeps its recorded session). +""" + +import asyncio +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +import download_and_process_data # noqa: E402 +import nemogym_agent_function as naf # noqa: E402 + + +def run_async(coro): + return asyncio.run(coro) + + +_INSTANCE_METADATA = { + "instance_id": "django__django-10973", + "repo": "django/django", + "base_commit": "ddb2936", + "problem_statement": "Use subprocess.run ...", + "subset": "gym", + "split": "train", +} + + +def _capture_post(captured, response=None): + async def fake_post(url, payload, **kwargs): + captured["url"] = url + captured["payload"] = payload + return response if response is not None else {"reward": 0.0, "metadata": {}} + + return fake_post + + +# --- request contract ----------------------------------------------------- + + +def test_run_body_carries_instance_fields_and_policy_override(monkeypatch): + captured = {} + monkeypatch.setattr(naf, "post_json", _capture_post(captured)) + monkeypatch.setenv("NEMO_GYM_URL", "http://gym:12000") + monkeypatch.delenv("MILES_ROUTER_EXTERNAL_HOST", raising=False) + + run_async( + naf.run( + base_url="http://trainer:30000/sessions/abc", + prompt="ignored", + request_kwargs={"temperature": 0.8, "top_p": 0.9, "max_tokens": 4096}, + metadata=dict(_INSTANCE_METADATA), + ) + ) + + assert captured["url"] == "http://gym:12000/run" + body = captured["payload"] + # Instance fields at the top level: the server uses body.model_dump() as + # the instance dict (image selection, eval). + for key, value in _INSTANCE_METADATA.items(): + assert body[key] == value + assert body["policy_base_url"] == "http://trainer:30000/sessions/abc/v1" + assert body["responses_create_params"] == { + "input": [], + "temperature": 0.8, + "top_p": 0.9, + "max_output_tokens": 4096, + } + # The fork-era channels must be gone. + assert "sglang_url" not in body + assert "sampling_params" not in body + + +def test_sampling_params_omitted_when_unset(monkeypatch): + captured = {} + monkeypatch.setattr(naf, "post_json", _capture_post(captured)) + + run_async(naf.run(base_url="http://t:1/sessions/s", prompt="", request_kwargs={}, metadata={})) + + assert captured["payload"]["responses_create_params"] == {"input": []} + + +def test_external_host_rewrites_session_url(monkeypatch): + captured = {} + monkeypatch.setattr(naf, "post_json", _capture_post(captured)) + monkeypatch.setenv("MILES_ROUTER_EXTERNAL_HOST", "100.64.0.7") + + run_async(naf.run(base_url="http://pod-hostname:30000/sessions/s1", prompt="", metadata={})) + + assert captured["payload"]["policy_base_url"] == "http://100.64.0.7:30000/sessions/s1/v1" + + +# --- response mapping ----------------------------------------------------- + + +def test_reward_and_eval_report_mapping(monkeypatch): + eval_report = {"django__django-10973": {"patch_successfully_applied": True}} + captured = {} + monkeypatch.setattr(naf, "post_json", _capture_post(captured, {"reward": 1.0, "metadata": eval_report})) + + result = run_async(naf.run(base_url="http://t:1/sessions/s", prompt="", metadata={})) + + assert result == {"reward": 1.0, "eval_report": eval_report} + + +def test_transport_failure_returns_none(monkeypatch): + async def failing_post(url, payload, **kwargs): + raise RuntimeError("connection refused") + + monkeypatch.setattr(naf, "post_json", failing_post) + + assert run_async(naf.run(base_url="http://t:1/sessions/s", prompt="", metadata={})) is None + + +def test_timeout_returns_none(monkeypatch): + async def slow_post(url, payload, **kwargs): + await asyncio.sleep(10) + + monkeypatch.setattr(naf, "post_json", slow_post) + monkeypatch.setenv("NEMO_GYM_RUN_TIMEOUT", "0.01") + + assert run_async(naf.run(base_url="http://t:1/sessions/s", prompt="", metadata={})) is None + + +# --- data conversion ------------------------------------------------------ + + +def test_convert_to_miles_format(tmp_path): + src = tmp_path / "raw.jsonl" + instance = {"instance_id": "x__y-1", "repo": "x/y", "problem_statement": "fix it", "patch": "diff"} + src.write_text(json.dumps(instance) + "\n") + dst = tmp_path / "miles.jsonl" + + download_and_process_data.convert_to_miles_format(str(src), str(dst), split="train") + + row = json.loads(dst.read_text()) + assert row["prompt"] == "fix it" + # Full instance preserved in metadata, plus the subset/split the server's + # image selection and eval need. + assert row["metadata"]["instance_id"] == "x__y-1" + assert row["metadata"]["patch"] == "diff" + assert row["metadata"]["subset"] == "gym" + assert row["metadata"]["split"] == "train"