From 235e475a5c26371eeda18d3075e2405da79e5edd Mon Sep 17 00:00:00 2001 From: Lin Jia Date: Tue, 14 Jul 2026 13:08:47 -0700 Subject: [PATCH 1/7] docs: add switchyard openhands implementation handoff Signed-off-by: Lin Jia --- gym_switchyard_integration.md | 459 ++++++++++++++++++++++++++++++++++ 1 file changed, 459 insertions(+) create mode 100644 gym_switchyard_integration.md diff --git a/gym_switchyard_integration.md b/gym_switchyard_integration.md new file mode 100644 index 0000000000..7aa47acf61 --- /dev/null +++ b/gym_switchyard_integration.md @@ -0,0 +1,459 @@ +# Gym and Switchyard OpenHands MVP + +**Status:** Implementation handoff + +**Last reviewed:** 2026-07-14 + +**Gym base:** `main` at `1b7d98a4` + +**Gym implementation branch:** `feature/switchyard-openhands-integration` + +**Switchyard implementation:** `drifold/feature/token-capture` at `1365ca0` +([NVIDIA-NeMo/Switchyard#63](https://github.com/NVIDIA-NeMo/Switchyard/pull/63)) + +**OpenHands base:** Gym's pinned +[`nv-OpenHands` commit](https://github.com/sdevare-nv/nv-OpenHands/tree/5f0180054732945df08ad2293903e6873f0492b6) + +## Objective + +Run Gym's existing OpenHands SWE harness through Switchyard, capture exact vLLM +prompt tokens, generated tokens, and sampled-token log probabilities for every policy +call, and return them in Gym's existing trainer-facing Responses API representation. + +The MVP keeps the existing responsibilities: + +- OpenHands owns the agent loop, tools, patch production, and completion logs. +- Gym owns SWE evaluation, reward, trace validation, and format adaptation. +- Switchyard owns proxying, exact token capture, durable per-call records, and session + retrieval. + +The expected path is: + +~~~text +Gym SWEBenchWrapper + -> OpenHands agent in Apptainer + -> Switchyard + -> one vLLM target +~~~ + +## Important baseline decision + +This implementation starts from Gym `main`. It does not depend on +`feat/model-server-observability`. + +Do not use or port: + +- `CaptureStore` +- `ModelCallRecord` +- model-call capture middleware +- `/ng-rollout//...` URL routing +- `ng_model_call_capture` + +Those types describe Gym model-server observability. Switchyard is the capture source +for this integration, and its records contain the token data needed for training. + +The main branch already has a unique per-run setup path in +`SWEBenchWrapper._setup_params`. The MVP adds a Switchyard UUID there instead of +changing generic rollout collection or importing retry-ID plumbing from the +observability branch. + +## Existing code to reuse + +### Gym + +The current implementation is concentrated in +`responses_api_agents/swe_agents/app.py`: + +- `SWEBenchWrapperConfig` owns harness configuration. +- `SWEBenchWrapperInstanceConfig` is serialized into response metadata and returned + as `instance_config`. +- `OpenHandsHarnessProcessor` builds the agent-container script. +- `SWEBenchWrapper._setup_params` creates unique per-run state. +- `SWEBenchWrapper._inner_responses` builds the current OpenHands trajectory. +- `SWEBenchWrapper.run` creates the final + `responses_create_params`, `response`, `reward`, and `instance_config`. +- `VLLMConverter(return_token_id_information=True)` converts Chat Completions + messages into Gym Responses API items. +- `TokenIDLogProbMixin` is the existing training contract for + `prompt_token_ids`, `generation_token_ids`, and + `generation_log_probs`. + +The current OpenHands configuration pins +`sdevare-nv/nv-OpenHands@5f0180054732945df08ad2293903e6873f0492b6`. +The fork is cloned at setup time; its client is not vendored in Gym. + +### Switchyard + +PR #63 already provides: + +- `target.token_capture_engine: vllm` +- capture activation through `--enable-rl-logging` and `--rl-log-dir` +- correlation through the `proxy_x_session_id` request header +- one JSON completion record per model call +- exact prompt IDs, generation IDs, and generation log probabilities +- buffered upstream responses with a normal synthesized client stream +- `GET /v1/sessions/{session_id}/completions` + +The completion endpoint returns records ordered by capture time and UUID. The Gym +adapter still validates message-history continuity and never repairs a broken chain +using timestamps. + +## MVP scope + +The first implementation supports: + +- `responses_api_agents/swe_agents` +- `agent_framework: openhands` +- the pinned `nv-OpenHands` fork +- non-replay seed rollouts +- the sequential main-agent session only +- one Switchyard route and one vLLM target +- post-agent HTTP retrieval +- one normal Gym rollout with one token-annotated generated item per model call + +The first implementation excludes: + +- OpenCode and other Gym harnesses +- replay-prefix or continuation rollouts +- subagent reconstruction +- parallel or branching policy calls +- context compaction or earlier-message rewriting +- SGLang and multiple routed targets +- Polar-style flattened arrays or token-prefix merging +- changes to Gym's generic rollout collection + +Unsupported trace shapes fail closed and mask the sample. They do not produce a +partially annotated training trajectory. + +## Runtime configuration + +### Switchyard + +Use the existing Gym model-server name as the Switchyard route ID. Gym's current +OpenHands configs use `policy_model`. + +~~~yaml +routes: + policy_model: + type: model + target: + model: + base_url: http://:/v1 + api_key: dummy + format: openai + token_capture_engine: vllm +~~~ + +Start the legacy route-bundle path so the Python processor chain is installed: + +~~~bash +switchyard \ + --routing-profiles single-vllm.yaml \ + --enable-rl-logging \ + --rl-log-dir /data/switchyard-traces/ \ + -- serve +~~~ + +Do not replace this with `serve --config`; that code path does not install the +capture processors in the current Switchyard branch. + +### Gym + +Add one optional setting to `SWEBenchWrapperConfig`: + +~~~yaml +swe_agents: + responses_api_agents: + swe_agents: + agent_framework: openhands + switchyard_base_url: http://:4000 + model_server: + name: policy_model + type: responses_api_models +~~~ + +When `switchyard_base_url` is absent, behavior remains unchanged. When present for +an OpenHands run, policy calls go directly to Switchyard. + +The URL must be reachable from both: + +- the Gym wrapper process, for retrieval; and +- the OpenHands Apptainer container, for model calls. + +## Session identity + +Do not derive the Switchyard session from task indices, the SWE instance ID, or the +observability branch's rollout ID helper. + +In `SWEBenchWrapper._setup_params`, generate: + +~~~python +switchyard_session_id = uuid.uuid4().hex +~~~ + +Add `switchyard_session_id: Optional[str] = None` to +`SWEBenchWrapperInstanceConfig`. Populate it only when all of these are true: + +- `switchyard_base_url` is configured; +- `agent_framework == "openhands"`; and +- this is a normal agent run rather than golden-patch verification. + +A lowercase UUID hex string is URL-safe, maps safely through Switchyard's current +session-directory sanitizer, and is unique across retries. It is serialized in the +existing `instance_config`, so `SWEBenchWrapper.run` can retrieve the same session +after `responses` returns without introducing shared mutable state. + +## OpenHands transport change + +Change the existing `NemoGymClient._post_completion` in the external +`nv-OpenHands` fork. Do not add a second client or a local compatibility proxy. + +The direct path is: + +1. Build `params` exactly as the client does today. +2. If `SWITCHYARD_BASE_URL` is absent, keep the existing + `ServerClient.post` behavior. +3. If it is present: + - require `SWITCHYARD_SESSION_ID`; + - copy `params`; + - replace `params["model"]` with `NEMO_GYM_MODEL_SERVER_NAME`; + - POST to `/v1/chat/completions`; + - send `proxy_x_session_id: `. +4. Reuse the client's existing response-status handling, JSON parsing, cookies, + metrics, and completion-log writing. + +The model rewrite is required. Switchyard dispatches on the incoming OpenAI +`model` field, while OpenHands normally sends the underlying LLM model string. +`NEMO_GYM_MODEL_SERVER_NAME` is already exported by the Gym harness and exactly +matches the `policy_model` route in the example. + +Gym must export two additional values only into the agent container: + +~~~text +SWITCHYARD_BASE_URL= +SWITCHYARD_SESSION_ID= +~~~ + +The evaluator container does not make policy calls and does not receive these values. + +After the fork change is pushed, update every OpenHands pin in Gym: + +- `responses_api_agents/swe_agents/configs/swebench_openhands.yaml` +- `responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml` +- `responses_api_agents/swe_agents/configs/swebench_multi_tools.yaml` + +## Gym retrieval and adaptation + +Add a small module at +`responses_api_agents/swe_agents/switchyard_trace.py`. Keep network orchestration +in `SWEBenchWrapper.run`; keep schema validation and conversion in the module. + +After `SWEBenchWrapper.responses` returns, `run` already reads the serialized +`instance_config`. If it contains a Switchyard URL and session UUID, request: + +~~~text +GET /v1/sessions//completions +~~~ + +Use Gym's existing async HTTP and response helpers with a bounded timeout. No polling +or session-finalization protocol is needed: OpenHands has exited, and Switchyard +writes each record before returning that call's model response. + +### Required session validation + +Reject the session unless: + +- the envelope has `schema_version == 1`; +- the envelope session ID equals the requested UUID; +- at least one record exists; +- every record has `schema_version == 1`; +- every record has the requested session ID and a unique UUID; +- every record has `is_valid == true`; +- every record has non-empty `request_id` and `model`; +- prompt and generation IDs are non-empty integer arrays; +- generation IDs and finite generation log probabilities have equal lengths; +- messages are non-empty and end with an assistant message; and +- model, tools, and tool choice are consistent across the session. + +Process records in the endpoint's returned order, but require each later prompt +history to strictly extend the already reconstructed history. Capture time and UUID +are not permitted to repair a divergent or reordered history. + +### Reconstruction algorithm + +For each record: + +1. Split `messages` into `prompt_messages = messages[:-1]` and the final + assistant message. +2. Copy the record's token triple onto that assistant Chat Completions message. +3. For the first record, initialize the sequence with its prompt plus assistant. +4. For each later record, require its prompt to begin with the entire assembled + sequence and to contain a non-empty environment-message suffix. +5. Append only that suffix and the new assistant message. +6. Convert the assembled messages with the existing + `VLLMConverter(return_token_id_information=True)`. +7. Use the existing `split_responses_input_output_items` helper. + +Gym's converter places the token triple on the final generated Responses API item +for one assistant message. A message containing tool calls can produce multiple +Responses items, but exactly one generated item retains that call's token triple. + +Map Switchyard tools as follows: + +| Switchyard | Gym function tool | +|---|---| +| `id` | `name` | +| `description` | `description` | +| `inputSchema.jsonSchema` | `parameters` | +| implicit tool kind | `type: "function"` | + +On success: + +- replace `responses_create_params.input` with reconstructed input items; +- replace `responses_create_params.tools` with mapped Switchyard tools; +- replace `response.output` with reconstructed output items; +- keep Gym's existing response identity and reward; +- attach compact response metadata containing the Switchyard session ID, source, + record UUIDs, and captured model. Keep metadata values as strings; JSON-encode the + UUID list. + +Do not attach the raw completion envelope to the rollout. Switchyard remains the +durable diagnostic source. + +### Failure behavior + +Retrieval, schema, or reconstruction failure must not erase the existing OpenHands +patch, evaluation metrics, response text, or reward. + +Instead: + +- keep the current OpenHands-derived response; +- set `instance_config["mask_sample"] = true`; +- add `ng_switchyard_trace_error: Optional[str]` to + `SWEBenchVerifyResponse`; +- store a concise error containing the session ID and reason; and +- emit no Switchyard token annotations. + +This is fail-closed for training and fail-open for rollout diagnostics. + +## Trace format rationale + +Switchyard emits one lossless capture record per model call. Gym trainers consume +Responses API items carrying `TokenIDLogProbMixin`. + +The MVP therefore returns one ordinary top-level Gym rollout: + +| Source | Gym destination | +|---|---| +| first record's initial prompt | `responses_create_params.input` | +| later user/tool-result suffixes | ordered `response.output` items | +| each captured assistant generation | generated `response.output` item(s) | +| each record's token triple | final generated item for that assistant message | +| OpenHands evaluation | existing `reward` | +| session and record identity | compact response metadata | + +Do not add a nested trainer-facing trace list. Do not retokenize messages. Do not +concatenate token arrays across calls. User and tool-result items remain context and +carry no token triple. + +Polar-style prefix merging can be added later at the trainer boundary. It is not part +of this MVP. + +## Files to change + +### Gym + +1. `responses_api_agents/swe_agents/app.py` + - add `switchyard_base_url`; + - add and generate `switchyard_session_id`; + - export the two Switchyard environment variables for OpenHands only; + - retrieve and apply the reconstructed trace in `run`; + - add `ng_switchyard_trace_error`. +2. `responses_api_agents/swe_agents/switchyard_trace.py` + - DTOs, validation, tool mapping, and reconstruction. +3. `responses_api_agents/swe_agents/tests/test_app.py` + - configuration, UUID creation, environment export, retrieval integration, and + masking behavior. +4. `responses_api_agents/swe_agents/tests/test_switchyard_trace.py` + - focused adapter and validation tests. +5. The three OpenHands YAML files listed above + - update the fork commit after the external change lands. + +Do not change `nemo_gym/rollout_collection.py` or generic model-server code. + +### nv-OpenHands + +1. Update `openhands/agenthub/nemo_gym_client.py`. +2. Add focused client tests for: + - fallback to Gym `ServerClient`; + - direct URL construction; + - route-model rewrite; + - session header; + - missing-session failure; + - response parsing, cookies, metrics, and completion logging. +3. Push the fork commit before updating Gym's pin. + +### Switchyard + +No Gym-specific Switchyard feature is required beyond PR #63. + +Two correctness limitations remain in that branch: + +- arbitrary session IDs can collide after directory-name sanitization; the Gym MVP + avoids this by sending UUID hex only; +- missing `request_id` or `model` does not currently make Switchyard's + `is_valid` false; Gym explicitly validates both fields. + +The parser and target schema already agree on nested +`target.token_capture_engine`; that earlier mismatch is resolved. + +## Tests + +Minimum coverage: + +1. No Switchyard setting preserves current OpenHands behavior exactly. +2. A configured run creates a unique UUID and exports it only to the agent + container. +3. The OpenHands client selects direct HTTP, rewrites the route model, and sends the + session header. +4. A valid two-call tool trajectory reconstructs into one Gym rollout with two + exact token triples and unchanged reward. +5. Invalid schema, wrong session, duplicate UUID, empty tokens, non-finite + log probabilities, mismatched token lengths, and divergent histories mask the + sample without partial annotations. +6. One built-in OpenHands example rollout runs end to end through Switchyard and + preserves its patch/evaluation result. + +Focused existing regression suites: + +~~~bash +uv run pytest -q tests/unit_tests/test_responses_converter.py +uv run pytest -q responses_api_agents/swe_agents/tests/test_app.py +~~~ + +Switchyard capture regression suite: + +~~~bash +uv run pytest -q \ + tests/test_token_capture.py \ + tests/test_token_capture_target.py \ + tests/test_token_capture_translation.py +~~~ + +## Implementation order + +1. Change and test the pinned `nv-OpenHands` client. +2. Push the fork commit and update Gym's three pins. +3. Add Gym config/session plumbing and agent-only environment export. +4. Add the Gym trace DTO, validation, and reconstruction module. +5. Integrate retrieval and fail-closed masking in `SWEBenchWrapper.run`. +6. Run focused unit tests. +7. Run one example end to end with a fresh Switchyard log directory. + +## Done when + +The MVP is complete when a built-in Gym OpenHands rollout sends every policy call +through Switchyard under one unique UUID, retrieves all valid records after the agent +finishes, reconstructs one normal Gym training rollout with unchanged token IDs and +log probabilities, preserves the existing patch/evaluation/reward result, and masks +the sample rather than returning partial training data whenever capture is invalid. From 712f16b643970200cec184ecac78f8c09abec157 Mon Sep 17 00:00:00 2001 From: Lin Jia Date: Tue, 14 Jul 2026 13:46:06 -0700 Subject: [PATCH 2/7] docs: polish handoff after verifying claims against both repos - name the retrieval envelope's record list field (completions) and exact ordering keys (captured_at, uuid) - note serve --config rejects --enable-rl-logging and that retrieval works from any capture-enabled process sharing the rl-log-dir - require exactly one HTTP attempt per direct policy call: transport retries can duplicate capture records and mask the sample - document that token-field stripping is a no-op on the direct path - rename ng_switchyard_trace_error -> switchyard_trace_error to match existing agent_error_kind-style naming (no ng_ prefix convention) - cite the existing mask_sample flag and verify_golden_patch gate Co-Authored-By: Claude Fable 5 Signed-off-by: Lin Jia --- gym_switchyard_integration.md | 63 ++++++++++++++++++++++++++--------- 1 file changed, 47 insertions(+), 16 deletions(-) diff --git a/gym_switchyard_integration.md b/gym_switchyard_integration.md index 7aa47acf61..ec4f86e258 100644 --- a/gym_switchyard_integration.md +++ b/gym_switchyard_integration.md @@ -94,9 +94,10 @@ PR #63 already provides: - buffered upstream responses with a normal synthesized client stream - `GET /v1/sessions/{session_id}/completions` -The completion endpoint returns records ordered by capture time and UUID. The Gym -adapter still validates message-history continuity and never repairs a broken chain -using timestamps. +The retrieval envelope is `{schema_version, session_id, completions}` — the record +list field is named `completions`. Records are returned ordered by +(`captured_at`, `uuid`). The Gym adapter still validates message-history continuity +and never repairs a broken chain using timestamps. ## MVP scope @@ -154,8 +155,14 @@ switchyard \ -- serve ~~~ -Do not replace this with `serve --config`; that code path does not install the -capture processors in the current Switchyard branch. +Do not replace this with `serve --config`; that path runs the Rust profile server, +which has no Python processor chain and rejects `--enable-rl-logging` outright. + +The session-retrieval endpoint is registered only on this capture-enabled serve +path. Records are plain JSON files under `--rl-log-dir`, so retrieval does not +require the exact process that captured them: any capture-enabled Switchyard +process pointed at the same log directory can serve +`GET /v1/sessions/...`, including after a proxy restart. ### Gym @@ -196,7 +203,8 @@ Add `switchyard_session_id: Optional[str] = None` to - `switchyard_base_url` is configured; - `agent_framework == "openhands"`; and -- this is a normal agent run rather than golden-patch verification. +- this is a normal agent run rather than golden-patch verification + (`verify_golden_patch` is false). A lowercase UUID hex string is URL-safe, maps safely through Switchyard's current session-directory sanitizer, and is unique across retries. It is serialized in the @@ -218,14 +226,30 @@ The direct path is: - copy `params`; - replace `params["model"]` with `NEMO_GYM_MODEL_SERVER_NAME`; - POST to `/v1/chat/completions`; - - send `proxy_x_session_id: `. + - send `proxy_x_session_id: `; + - make exactly one HTTP attempt — no transport retries. 4. Reuse the client's existing response-status handling, JSON parsing, cookies, metrics, and completion-log writing. +The single-attempt rule exists because an attempt that reaches vLLM writes a +capture record even when the client never receives the response. A transport +retry can therefore leave two records with the same prompt history; strict-history +reconstruction would refuse the session and mask the sample. Gym's +`ServerClient.post` and `nemo_gym.server_utils.request` both retry with backoff, +so the direct path must use the shared aiohttp client without the retry wrapper +and let a failed call fail the run through existing OpenHands error handling. + The model rewrite is required. Switchyard dispatches on the incoming OpenAI `model` field, while OpenHands normally sends the underlying LLM model string. `NEMO_GYM_MODEL_SERVER_NAME` is already exported by the Gym harness and exactly -matches the `policy_model` route in the example. +matches the `policy_model` route in the example. `params["model"]` always exists: +`_nemo_gym_llm_kwargs` is built with `model=self.config.model`. + +No message-format change is needed. On the direct path the response message never +carries `prompt_token_ids`, `generation_token_ids`, or `generation_log_probs` — +those are attached by Gym's model server, not vLLM — so the client's existing +token-field stripping loop is a no-op and every captured prompt history stays +strictly extending. Gym must export two additional values only into the agent container: @@ -255,9 +279,11 @@ After `SWEBenchWrapper.responses` returns, `run` already reads the serialized GET /v1/sessions//completions ~~~ -Use Gym's existing async HTTP and response helpers with a bounded timeout. No polling +Use Gym's existing async HTTP helper (`nemo_gym.server_utils.request`) with a +bounded timeout; its retries are safe here because the GET is idempotent. No polling or session-finalization protocol is needed: OpenHands has exited, and Switchyard -writes each record before returning that call's model response. +durably writes each record (atomic rename) before returning that call's model +response. ### Required session validation @@ -265,7 +291,7 @@ Reject the session unless: - the envelope has `schema_version == 1`; - the envelope session ID equals the requested UUID; -- at least one record exists; +- the `completions` list is non-empty; - every record has `schema_version == 1`; - every record has the requested session ID and a unique UUID; - every record has `is_valid == true`; @@ -328,9 +354,10 @@ patch, evaluation metrics, response text, or reward. Instead: - keep the current OpenHands-derived response; -- set `instance_config["mask_sample"] = true`; -- add `ng_switchyard_trace_error: Optional[str]` to - `SWEBenchVerifyResponse`; +- set the existing `SWEBenchWrapperInstanceConfig.mask_sample` flag to true — this + is Gym's established training-mask convention, already used for agent failures; +- add `switchyard_trace_error: Optional[str]` to `SWEBenchVerifyResponse`, + following the existing `agent_error_kind`-style metrics naming; - store a concise error containing the session ID and reason; and - emit no Switchyard token annotations. @@ -366,9 +393,12 @@ of this MVP. 1. `responses_api_agents/swe_agents/app.py` - add `switchyard_base_url`; - add and generate `switchyard_session_id`; - - export the two Switchyard environment variables for OpenHands only; + - export the two Switchyard environment variables next to the existing + `NEMO_GYM_MODEL_SERVER_NAME` export in `OpenHandsHarnessProcessor`'s + agent command (agent container only; the opencode harness and the eval + container command are untouched); - retrieve and apply the reconstructed trace in `run`; - - add `ng_switchyard_trace_error`. + - add `switchyard_trace_error`. 2. `responses_api_agents/swe_agents/switchyard_trace.py` - DTOs, validation, tool mapping, and reconstruction. 3. `responses_api_agents/swe_agents/tests/test_app.py` @@ -390,6 +420,7 @@ Do not change `nemo_gym/rollout_collection.py` or generic model-server code. - route-model rewrite; - session header; - missing-session failure; + - exactly one HTTP attempt per call (no transport retries); - response parsing, cookies, metrics, and completion logging. 3. Push the fork commit before updating Gym's pin. From 99e38456394447dd6f535d61725e07f538466ae1 Mon Sep 17 00:00:00 2001 From: Lin Jia Date: Tue, 14 Jul 2026 14:13:52 -0700 Subject: [PATCH 3/7] feat: Switchyard token-capture integration for OpenHands SWE runs Implements the MVP from gym_switchyard_integration.md: - SWEBenchWrapperConfig.switchyard_base_url (optional; absent = unchanged behavior) and SWEBenchWrapperInstanceConfig.switchyard_session_id (uuid4 hex, generated per OpenHands agent run, never for golden-patch verification) - OpenHandsHarnessProcessor exports SWITCHYARD_BASE_URL/SESSION_ID into the agent container script only; the eval container is untouched - switchyard_trace.py validates a retrieved capture session fail-closed (schema, per-record token triples, session consistency, strict-history extension) and reconstructs one token-annotated Gym rollout via the existing VLLMConverter/split helpers - SWEBenchWrapper.run retrieves GET /v1/sessions/{id}/completions after the agent exits and replaces input/tools/output on success; any failure masks the sample (mask_sample=true, switchyard_trace_error) while preserving the OpenHands patch/eval/reward Requires the nv-OpenHands fork commit adding direct Switchyard transport to NemoGymClient (branch switchyard-direct-transport); the three config pins will be bumped once that commit is pushed. Co-Authored-By: Claude Fable 5 Signed-off-by: Lin Jia --- responses_api_agents/swe_agents/app.py | 92 +++++- .../swe_agents/switchyard_trace.py | 269 ++++++++++++++++++ .../swe_agents/tests/test_app.py | 256 +++++++++++++++++ .../swe_agents/tests/test_switchyard_trace.py | 260 +++++++++++++++++ 4 files changed, 872 insertions(+), 5 deletions(-) create mode 100644 responses_api_agents/swe_agents/switchyard_trace.py create mode 100644 responses_api_agents/swe_agents/tests/test_switchyard_trace.py diff --git a/responses_api_agents/swe_agents/app.py b/responses_api_agents/swe_agents/app.py index 316fc13ffa..3f35626da2 100644 --- a/responses_api_agents/swe_agents/app.py +++ b/responses_api_agents/swe_agents/app.py @@ -40,6 +40,7 @@ import ray import tomlkit +from aiohttp import ClientTimeout from gprof2dot import main as gprof2dot_main from openai.types.responses.function_tool import FunctionTool from pydantic import BaseModel, ConfigDict, Field @@ -62,7 +63,8 @@ NeMoGymResponseCreateParamsNonStreaming, ) from nemo_gym.profiling import Profiler -from nemo_gym.server_utils import get_first_server_config_dict +from nemo_gym.server_utils import get_first_server_config_dict, get_response_json, raise_for_status, request +from responses_api_agents.swe_agents.switchyard_trace import SwitchyardTrace, reconstruct_switchyard_rollout from responses_api_models.vllm_model.app import VLLMConverter, split_responses_input_output_items @@ -177,6 +179,16 @@ class SWEBenchWrapperConfig(BaseResponsesAPIAgentConfig): "If False (default), selection is deterministic per instance_id.", ) + switchyard_base_url: Optional[str] = Field( + default=None, + description=( + "Base URL of a token-capture-enabled Switchyard proxy (e.g. http://host:4000). When set for an " + "OpenHands run, agent policy calls go directly to Switchyard and the captured per-call token " + "records are retrieved after the agent finishes to build the training rollout. Must be reachable " + "from both this process and the agent container. When absent, behavior is unchanged." + ), + ) + openhands_should_log: bool = False debug: bool = False @@ -255,6 +267,10 @@ class SWEBenchWrapperInstanceConfig(SWEBenchWrapperServerConfig, SWEBenchWrapper # GRPO related fields mask_sample: bool = False + # Switchyard capture session for this run (uuid4 hex). Set only for OpenHands + # agent runs with switchyard_base_url configured; None leaves behavior unchanged. + switchyard_session_id: Optional[str] = None + @property def instance_id(self) -> str: return self.problem_info["instance_id"] @@ -299,6 +315,8 @@ class SWEBenchMetrics(BaseModel): class SWEBenchVerifyResponse(SWEBenchMetrics, BaseVerifyResponse): instance_config: SWEBenchWrapperInstanceConfig subagent_trajectories: Optional[List[Dict[str, Any]]] = None + # Why the Switchyard trace could not be applied (the sample is then masked). + switchyard_trace_error: Optional[str] = None ######################################## @@ -1684,6 +1702,17 @@ def get_run_command(self) -> ExecuteContainerCommandArgs: else: camel_case_tool_names_cmd = "" + # Agent container only: the eval container makes no policy calls and must + # not receive these. The in-container NemoGymClient switches to direct + # Switchyard HTTP when SWITCHYARD_BASE_URL is set. + if self.config.switchyard_base_url and self.config.switchyard_session_id: + switchyard_cmd = ( + f"export SWITCHYARD_BASE_URL={shlex.quote(self.config.switchyard_base_url)} && " + f"export SWITCHYARD_SESSION_ID={self.config.switchyard_session_id} && " + ) + else: + switchyard_cmd = "" + workspace_check_cmd = "" # Run the same baseline dependency/env repair the eval uses (if any), in a @@ -1715,7 +1744,8 @@ def get_run_command(self) -> ExecuteContainerCommandArgs: f"{profiling_cmd}" f"export NEMO_GYM_METRICS_FPATH={self.config.base_mounted_dir}/nemo_gym_metrics.json && " f"export NEMO_GYM_CONFIG_DICT={self.config.ng_global_config_dict_str} && " - f"export NEMO_GYM_MODEL_SERVER_NAME={self.config.model_server_name} &&" + f"export NEMO_GYM_MODEL_SERVER_NAME={self.config.model_server_name} && " + f"{switchyard_cmd}" "export VIRTUAL_ENV=/openhands_setup/OpenHands/.venv && " "export PATH=$PATH:/openhands_setup/OpenHands/.venv/bin && " # CRITICAL: Configure poetry to only use the OpenHands venv (ignore external venvs) @@ -3398,6 +3428,17 @@ def _setup_params( agent_run_id = f"{instance_id}_{int(time.time())}_{str(uuid.uuid4())[:8]}" + # One Switchyard capture session per agent run. uuid4 hex is URL-safe, maps + # 1:1 through Switchyard's session-directory sanitizer, and is unique across + # retries. Golden-patch verification makes no policy calls, so no session. + switchyard_session_id = None + if ( + self.config.switchyard_base_url + and self.config.agent_framework == "openhands" + and not self.config.verify_golden_patch + ): + switchyard_session_id = uuid.uuid4().hex + # Ground-truth artifacts live here; hidden from the container by an empty overmount. eval_private_dir = persistent_dir / "eval_private" (eval_private_dir / ".empty").mkdir(parents=True, exist_ok=True) @@ -3458,6 +3499,7 @@ def _setup_params( ray_queue_timestamp=time.time(), inference_params=inference_params, agent_run_id=agent_run_id, + switchyard_session_id=switchyard_session_id, instance_dataset_path=instance_dataset_path, agent_instance_dataset_path=agent_instance_dataset_path, trajectories_root=trajectories_root, @@ -3685,17 +3727,57 @@ async def run(self, body: BaseRunRequest) -> SWEBenchVerifyResponse: if "subagent_trajectories" in metadata: subagent_trajectories = json.loads(metadata["subagent_trajectories"]) + instance_config = SWEBenchWrapperInstanceConfig.model_validate_json(metadata["instance_config"]) + switchyard_trace_error = None + if instance_config.switchyard_base_url and instance_config.switchyard_session_id: + try: + trace = await self._retrieve_switchyard_trace(instance_config) + switchyard_input = [item.model_dump() for item in trace.input_items] + switchyard_tools = [tool.model_dump() for tool in trace.tools] + responses_create_params["input"] = switchyard_input + responses_create_params["tools"] = switchyard_tools + response.output = trace.output_items + response.metadata = { + "switchyard_source": "switchyard", + "switchyard_session_id": instance_config.switchyard_session_id, + "switchyard_record_uuids": json.dumps(trace.record_uuids), + "switchyard_model": trace.model, + } + except Exception as e: + # Fail closed for training, open for diagnostics: keep the + # OpenHands-derived rollout/patch/reward, but mask the sample + # rather than emit a partially token-annotated trajectory. + instance_config.mask_sample = True + switchyard_trace_error = ( + f"session {instance_config.switchyard_session_id}: {type(e).__name__}: {e}" + ) + return SWEBenchVerifyResponse( responses_create_params=responses_create_params, response=response, reward=1.0 if metrics.resolved else 0.0, **metrics.model_dump(), - instance_config=SWEBenchWrapperInstanceConfig.model_validate_json( - metadata["instance_config"] - ).model_dump(), + instance_config=instance_config.model_dump(), subagent_trajectories=subagent_trajectories, + switchyard_trace_error=switchyard_trace_error, ) + async def _retrieve_switchyard_trace(self, instance_config: SWEBenchWrapperInstanceConfig) -> SwitchyardTrace: + """Fetch this run's captured completions from Switchyard and rebuild the rollout. + + No polling is needed: the agent has exited and Switchyard durably writes + each record before returning that call's model response. The GET is + idempotent, so the helper's transport retries are safe. + """ + url = ( + f"{instance_config.switchyard_base_url.rstrip('/')}" + f"/v1/sessions/{instance_config.switchyard_session_id}/completions" + ) + response = await request("GET", url, timeout=ClientTimeout(total=60)) + await raise_for_status(response) + envelope = await get_response_json(response) + return reconstruct_switchyard_rollout(envelope, instance_config.switchyard_session_id, self._vllm_converter) + if __name__ == "__main__": SWEBenchWrapper.run_webserver() diff --git a/responses_api_agents/swe_agents/switchyard_trace.py b/responses_api_agents/swe_agents/switchyard_trace.py new file mode 100644 index 0000000000..73f6f2b274 --- /dev/null +++ b/responses_api_agents/swe_agents/switchyard_trace.py @@ -0,0 +1,269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Validation and reconstruction of Switchyard token-capture sessions. + +Switchyard's `GET /v1/sessions/{session_id}/completions` returns one lossless +record per policy call: the OpenAI chat messages it proxied (prompt plus the +generated assistant turn) and the exact vLLM token triple +(`prompt_token_ids`, `generation_token_ids`, `generation_log_probs`). + +This module validates a retrieved session fail-closed and reconstructs it into +one ordinary Gym rollout: Responses API input/output items where each captured +assistant generation carries its token triple on the final generated item. +Any schema or history divergence raises `SwitchyardTraceError` — the caller +masks the sample rather than emitting partially annotated training data. +""" + +import math +from typing import Any, Dict, List, NamedTuple + +from openai.types.responses.function_tool import FunctionTool + +from nemo_gym.responses_converter import ResponsesConverter, split_responses_input_output_items + + +SWITCHYARD_SCHEMA_VERSION = 1 + +# Message roles OpenHands can legitimately send through Switchyard. +_KNOWN_ROLES = {"system", "developer", "user", "assistant", "tool"} +# Roles allowed in the between-call suffix a later prompt appends: environment +# feedback only. An assistant message here means a policy call was not captured. +_ENV_ROLES = {"user", "tool"} + + +class SwitchyardTraceError(Exception): + """A Switchyard session failed validation or reconstruction.""" + + +class SwitchyardTrace(NamedTuple): + input_items: List[Any] + output_items: List[Any] + tools: List[FunctionTool] + record_uuids: List[str] + model: str + + +def map_switchyard_tools(raw_tools: Any) -> List[FunctionTool]: + """Map Switchyard trace tools `{id, description, inputSchema.jsonSchema}` to Gym function tools.""" + if not isinstance(raw_tools, list): + raise SwitchyardTraceError(f"tools is not a list: {type(raw_tools).__name__}") + + tools = [] + for i, entry in enumerate(raw_tools): + if not isinstance(entry, dict) or not isinstance(entry.get("id"), str) or not entry["id"]: + raise SwitchyardTraceError(f"tool {i} has no usable id") + input_schema = entry.get("inputSchema") + parameters = input_schema.get("jsonSchema") if isinstance(input_schema, dict) else None + tools.append( + FunctionTool( + type="function", + name=entry["id"], + description=entry.get("description") or None, + parameters=parameters, + strict=None, + ) + ) + return tools + + +def reconstruct_switchyard_rollout(envelope: Any, session_id: str, converter: ResponsesConverter) -> SwitchyardTrace: + """Validate a retrieval envelope and rebuild one token-annotated Gym rollout. + + Records are processed in the endpoint's returned order; each later prompt + must strictly extend the already assembled history with a non-empty + environment-message suffix. Raises `SwitchyardTraceError` on any deviation. + """ + records = _validate_envelope(envelope, session_id) + + first = records[0] + model, tools, tool_choice = first["model"], first["tools"], first.get("tool_choice") + + # `assembled` is the canonical history used for strict-extension comparison — + # later prompts never carry token triples, so it must stay triple-free. + # `annotated` is the conversion list, where each captured assistant message + # additionally carries its token triple. + assembled: List[Dict[str, Any]] = [] + annotated: List[Dict[str, Any]] = [] + record_uuids: List[str] = [] + for i, record in enumerate(records): + _validate_record(record, i, session_id, record_uuids) + if record["model"] != model: + raise SwitchyardTraceError(f"record {i} model {record['model']!r} != session model {model!r}") + if record["tools"] != tools or record.get("tool_choice") != tool_choice: + raise SwitchyardTraceError(f"record {i} tools/tool_choice differ from the session's first record") + + messages = [_normalize_message(m, i) for m in record["messages"]] + prompt, assistant = messages[:-1], messages[-1] + if assistant["role"] != "assistant": + raise SwitchyardTraceError(f"record {i} messages do not end with an assistant message") + + if i == 0: + new_context = prompt + else: + if prompt[: len(assembled)] != assembled: + raise SwitchyardTraceError(f"record {i} prompt does not extend the reconstructed history") + new_context = prompt[len(assembled) :] + if not new_context: + raise SwitchyardTraceError(f"record {i} prompt adds no environment messages") + bad_roles = {m["role"] for m in new_context} - _ENV_ROLES + if bad_roles: + raise SwitchyardTraceError(f"record {i} suffix contains non-environment roles: {sorted(bad_roles)}") + assembled.extend(new_context) + assembled.append(assistant) + + annotated.extend(new_context) + # The converter moves the triple onto the final generated Responses item. + annotated.append( + assistant + | { + "prompt_token_ids": record["prompt_token_ids"], + "generation_token_ids": record["generation_token_ids"], + "generation_log_probs": record["generation_log_probs"], + } + ) + record_uuids.append(record["uuid"]) + + items = converter.chat_completions_messages_to_responses_items(annotated) + input_items, output_items = split_responses_input_output_items(items) + + return SwitchyardTrace( + input_items=input_items, + output_items=output_items, + tools=map_switchyard_tools(tools), + record_uuids=record_uuids, + model=model, + ) + + +def _validate_envelope(envelope: Any, session_id: str) -> List[Dict[str, Any]]: + if not isinstance(envelope, dict): + raise SwitchyardTraceError(f"envelope is not an object: {type(envelope).__name__}") + if envelope.get("schema_version") != SWITCHYARD_SCHEMA_VERSION: + raise SwitchyardTraceError(f"unsupported envelope schema_version: {envelope.get('schema_version')!r}") + if envelope.get("session_id") != session_id: + raise SwitchyardTraceError(f"envelope session_id {envelope.get('session_id')!r} != requested {session_id!r}") + completions = envelope.get("completions") + if not isinstance(completions, list) or not completions: + raise SwitchyardTraceError("envelope has no completions") + return completions + + +def _validate_record(record: Any, i: int, session_id: str, seen_uuids: List[str]) -> None: + if not isinstance(record, dict): + raise SwitchyardTraceError(f"record {i} is not an object") + if record.get("schema_version") != SWITCHYARD_SCHEMA_VERSION: + raise SwitchyardTraceError(f"record {i} has unsupported schema_version: {record.get('schema_version')!r}") + if record.get("session_id") != session_id: + raise SwitchyardTraceError(f"record {i} session_id {record.get('session_id')!r} != requested {session_id!r}") + record_uuid = record.get("uuid") + if not isinstance(record_uuid, str) or not record_uuid: + raise SwitchyardTraceError(f"record {i} has no uuid") + if record_uuid in seen_uuids: + raise SwitchyardTraceError(f"record {i} duplicates uuid {record_uuid}") + if record.get("is_valid") is not True: + raise SwitchyardTraceError(f"record {record_uuid} is_valid is not true") + for field in ("request_id", "model"): + if not isinstance(record.get(field), str) or not record[field]: + raise SwitchyardTraceError(f"record {record_uuid} has empty {field}") + for field in ("prompt_token_ids", "generation_token_ids"): + if not _is_token_id_list(record.get(field)): + raise SwitchyardTraceError(f"record {record_uuid} {field} is not a non-empty list of ints") + log_probs = record.get("generation_log_probs") + if not isinstance(log_probs, list) or not all(_is_finite_number(v) for v in log_probs): + raise SwitchyardTraceError(f"record {record_uuid} generation_log_probs is not a list of finite floats") + if len(record["generation_token_ids"]) != len(log_probs): + raise SwitchyardTraceError(f"record {record_uuid} generation token ids and log probs have different lengths") + if not isinstance(record.get("messages"), list) or not record["messages"]: + raise SwitchyardTraceError(f"record {record_uuid} has no messages") + + +def _is_token_id_list(value: Any) -> bool: + return isinstance(value, list) and bool(value) and all(_is_int(v) for v in value) + + +def _is_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _is_finite_number(value: Any) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) + + +def _normalize_message(message: Any, i: int) -> Dict[str, Any]: + """Canonicalize one chat message for history comparison and conversion. + + OpenHands serializes content as either a string or a list of text parts, + while Switchyard's captured assistant turn always carries string content — + the same logical message can appear in both shapes across records. Reduce + every message to `{role, content: str}` plus tool identity so strict + history comparison and the Responses converter see one canonical shape. + """ + if not isinstance(message, dict): + raise SwitchyardTraceError(f"record {i} contains a non-object message") + role = message.get("role") + if role not in _KNOWN_ROLES: + raise SwitchyardTraceError(f"record {i} contains a message with unsupported role {role!r}") + + normalized: Dict[str, Any] = {"role": role, "content": _content_to_text(message.get("content"), i)} + if role == "assistant": + tool_calls = message.get("tool_calls") + if tool_calls: + normalized["tool_calls"] = _normalize_tool_calls(tool_calls, i) + elif role == "tool": + tool_call_id = message.get("tool_call_id") + if not isinstance(tool_call_id, str) or not tool_call_id: + raise SwitchyardTraceError(f"record {i} contains a tool message without tool_call_id") + normalized["tool_call_id"] = tool_call_id + return normalized + + +def _content_to_text(content: Any, i: int) -> str: + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + texts = [] + for part in content: + if not isinstance(part, dict) or not isinstance(part.get("text"), str): + raise SwitchyardTraceError(f"record {i} contains an unsupported content part") + texts.append(part["text"]) + return "".join(texts) + raise SwitchyardTraceError(f"record {i} contains unsupported content of type {type(content).__name__}") + + +def _normalize_tool_calls(tool_calls: Any, i: int) -> List[Dict[str, Any]]: + if not isinstance(tool_calls, list): + raise SwitchyardTraceError(f"record {i} contains non-list tool_calls") + normalized = [] + for call in tool_calls: + function = call.get("function") if isinstance(call, dict) else None + if ( + not isinstance(function, dict) + or not isinstance(call.get("id"), str) + or not call["id"] + or not isinstance(function.get("name"), str) + or not function["name"] + or not isinstance(function.get("arguments"), str) + ): + raise SwitchyardTraceError(f"record {i} contains a malformed tool call") + normalized.append( + { + "id": call["id"], + "type": "function", + "function": {"name": function["name"], "arguments": function["arguments"]}, + } + ) + return normalized diff --git a/responses_api_agents/swe_agents/tests/test_app.py b/responses_api_agents/swe_agents/tests/test_app.py index 6f1402f24f..c90b3ddc6c 100644 --- a/responses_api_agents/swe_agents/tests/test_app.py +++ b/responses_api_agents/swe_agents/tests/test_app.py @@ -22,12 +22,16 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from openai.types.responses.function_tool import FunctionTool import responses_api_agents.swe_agents.app as swe_app from nemo_gym.config_types import ModelServerRef, OmegaConf from nemo_gym.openai_utils import ( + NeMoGymEasyInputMessage, NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming, + NeMoGymResponseOutputMessageForTraining, + NeMoGymResponseOutputText, ) from nemo_gym.server_utils import ServerClient from responses_api_agents.swe_agents.app import ( @@ -56,6 +60,7 @@ runner_ray_remote, update_and_read_metrics, ) +from responses_api_agents.swe_agents.switchyard_trace import SwitchyardTrace, SwitchyardTraceError SWE_AGENTS_DIR = Path(__file__).resolve().parent.parent @@ -933,6 +938,28 @@ def test_get_run_command_camel_case_tool_names(self) -> None: processor.get_run_command() assert "CAMEL_CASE_TOOL_NAMES=true" in self._read_agent_script(config) + def test_get_run_command_switchyard_env(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = _make_instance_config( + tmpdir, + switchyard_base_url="http://switchyard:4000", + switchyard_session_id="0123456789abcdef0123456789abcdef", + ) + config.persistent_dir.mkdir(parents=True, exist_ok=True) + processor = OpenHandsHarnessProcessor(config=config) + processor.get_run_command() + script = self._read_agent_script(config) + assert "export SWITCHYARD_BASE_URL=http://switchyard:4000 && " in script + assert "export SWITCHYARD_SESSION_ID=0123456789abcdef0123456789abcdef && " in script + + def test_get_run_command_no_switchyard_env_by_default(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = _make_instance_config(tmpdir) + config.persistent_dir.mkdir(parents=True, exist_ok=True) + processor = OpenHandsHarnessProcessor(config=config) + processor.get_run_command() + assert "SWITCHYARD" not in self._read_agent_script(config) + ######################################## # Workspace path + user-message resolver tests @@ -2286,6 +2313,67 @@ def test_setup_params_with_prompt_overrides(self, monkeypatch) -> None: # deterministic selection based on instance_id assert params.resolved_agent_cls in ["CodexAgent", "OpenCodeAgent"] + def _switchyard_body(self) -> NeMoGymResponseCreateParamsNonStreaming: + return NeMoGymResponseCreateParamsNonStreaming( + model="test-model", + input=[], + temperature=1.0, + top_p=1.0, + metadata={ + "problem_statement": "Fix bug", + "instance_id": "django__django-12345", + "base_commit": "abc123", + "dataset_name": "SWE-bench", + "split": "test", + "instance_dict": json.dumps({"repo": "django/django"}), + }, + ) + + def test_setup_params_switchyard_session(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "django__django-12345.sif").touch() + wrapper.config.container_formatter = [str(Path(tmpdir) / "{instance_id}.sif")] + self._setup_oh_dirs(wrapper) + wrapper.config.switchyard_base_url = "http://switchyard:4000" + + params, _ = wrapper._setup_params(self._switchyard_body()) + session_id = params.switchyard_session_id + assert session_id is not None + assert len(session_id) == 32 + assert set(session_id) <= set("0123456789abcdef") + # Exported into the agent container script. + assert f"export SWITCHYARD_SESSION_ID={session_id} && " in params.agent_script + assert "export SWITCHYARD_BASE_URL=http://switchyard:4000 && " in params.agent_script + + # Unique per run. + params2, _ = wrapper._setup_params(self._switchyard_body()) + assert params2.switchyard_session_id != session_id + + def test_setup_params_no_switchyard_by_default(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "django__django-12345.sif").touch() + wrapper.config.container_formatter = [str(Path(tmpdir) / "{instance_id}.sif")] + self._setup_oh_dirs(wrapper) + + params, _ = wrapper._setup_params(self._switchyard_body()) + assert params.switchyard_session_id is None + assert "SWITCHYARD" not in params.agent_script + + def test_setup_params_no_switchyard_for_golden_patch(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "django__django-12345.sif").touch() + wrapper.config.container_formatter = [str(Path(tmpdir) / "{instance_id}.sif")] + self._setup_oh_dirs(wrapper) + wrapper.config.switchyard_base_url = "http://switchyard:4000" + wrapper.config.verify_golden_patch = True + + params, _ = wrapper._setup_params(self._switchyard_body()) + assert params.switchyard_session_id is None + assert "SWITCHYARD" not in params.agent_script + class TestSWEBenchWrapperResponses: def _setup_oh_dirs(self, wrapper): @@ -2460,6 +2548,174 @@ async def test_run_not_resolved(self, monkeypatch) -> None: assert isinstance(result, SWEBenchVerifyResponse) assert result.reward == 0.0 + @staticmethod + def _run_body(): + from nemo_gym.base_resources_server import BaseRunRequest + + return BaseRunRequest( + responses_create_params=NeMoGymResponseCreateParamsNonStreaming( + model="test-model", + input=[], + metadata={ + "problem_statement": "Fix", + "instance_id": "test-1", + "base_commit": "abc", + "dataset_name": "SWE-bench", + "split": "test", + "instance_dict": "{}", + }, + ) + ) + + def _switchyard_response(self, **instance_config_overrides) -> NeMoGymResponse: + return NeMoGymResponse( + id="swebench-test", + created_at=123, + model="test-model", + object="response", + output=[], + parallel_tool_calls=True, + tool_choice="auto", + tools=[], + metadata={ + "input": "[]", + "metrics": json.dumps({"resolved": True, "patch_exists": True}), + "instance_config": _make_instance_config( + tempfile.mkdtemp(), + switchyard_base_url="http://switchyard:4000", + switchyard_session_id="0123456789abcdef0123456789abcdef", + **instance_config_overrides, + ).model_dump_json(), + }, + ) + + @pytest.mark.asyncio + async def test_run_switchyard_success_replaces_rollout(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + + trace = SwitchyardTrace( + input_items=[NeMoGymEasyInputMessage(role="user", content="Fix the bug.")], + output_items=[ + NeMoGymResponseOutputMessageForTraining( + id="msg_1", + content=[NeMoGymResponseOutputText(type="output_text", text="All fixed.", annotations=[])], + prompt_token_ids=[1, 2], + generation_token_ids=[3], + generation_log_probs=[-0.5], + ) + ], + tools=[FunctionTool(type="function", name="editor", description=None, parameters=None, strict=None)], + record_uuids=["u1", "u2"], + model="Qwen/Qwen3-0.6B", + ) + + with ( + patch.object( + SWEBenchWrapper, "responses", new_callable=AsyncMock, return_value=self._switchyard_response() + ), + patch.object(SWEBenchWrapper, "_retrieve_switchyard_trace", new_callable=AsyncMock, return_value=trace), + ): + result = await wrapper.run(self._run_body()) + + assert result.switchyard_trace_error is None + assert result.instance_config.mask_sample is False + assert result.reward == 1.0 + assert [item.model_dump() for item in result.responses_create_params.input] == [ + trace.input_items[0].model_dump() + ] + assert result.responses_create_params.tools == [trace.tools[0].model_dump()] + assert result.response.output == trace.output_items + assert result.response.metadata == { + "switchyard_source": "switchyard", + "switchyard_session_id": "0123456789abcdef0123456789abcdef", + "switchyard_record_uuids": json.dumps(["u1", "u2"]), + "switchyard_model": "Qwen/Qwen3-0.6B", + } + + @pytest.mark.asyncio + async def test_run_switchyard_failure_masks_sample(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + + with ( + patch.object( + SWEBenchWrapper, "responses", new_callable=AsyncMock, return_value=self._switchyard_response() + ), + patch.object( + SWEBenchWrapper, + "_retrieve_switchyard_trace", + new_callable=AsyncMock, + side_effect=SwitchyardTraceError("record 1 prompt does not extend the reconstructed history"), + ), + ): + result = await wrapper.run(self._run_body()) + + # Fail closed for training, open for diagnostics. + assert result.instance_config.mask_sample is True + assert result.switchyard_trace_error == ( + "session 0123456789abcdef0123456789abcdef: SwitchyardTraceError: " + "record 1 prompt does not extend the reconstructed history" + ) + assert result.reward == 1.0 + assert result.response.output == [] + assert result.responses_create_params.input == [] + assert result.response.metadata is None + + @pytest.mark.asyncio + async def test_run_without_switchyard_does_not_retrieve(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + + mock_response = NeMoGymResponse( + id="swebench-test", + created_at=123, + model="test-model", + object="response", + output=[], + parallel_tool_calls=True, + tool_choice="auto", + tools=[], + metadata={ + "input": "[]", + "metrics": json.dumps({"resolved": True, "patch_exists": True}), + "instance_config": _make_instance_config(tempfile.mkdtemp()).model_dump_json(), + }, + ) + + with ( + patch.object(SWEBenchWrapper, "responses", new_callable=AsyncMock, return_value=mock_response), + patch.object(SWEBenchWrapper, "_retrieve_switchyard_trace", new_callable=AsyncMock) as retrieve, + ): + result = await wrapper.run(self._run_body()) + + retrieve.assert_not_called() + assert result.switchyard_trace_error is None + assert result.instance_config.mask_sample is False + + @pytest.mark.asyncio + async def test_retrieve_switchyard_trace_url(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + with tempfile.TemporaryDirectory() as tmpdir: + instance_config = _make_instance_config( + tmpdir, + switchyard_base_url="http://switchyard:4000/", + switchyard_session_id="0123456789abcdef0123456789abcdef", + ) + + request_mock = AsyncMock(return_value=MagicMock()) + envelope = {"schema_version": 1} + monkeypatch.setattr(swe_app, "request", request_mock) + monkeypatch.setattr(swe_app, "raise_for_status", AsyncMock()) + monkeypatch.setattr(swe_app, "get_response_json", AsyncMock(return_value=envelope)) + reconstruct_mock = MagicMock(return_value="trace") + monkeypatch.setattr(swe_app, "reconstruct_switchyard_rollout", reconstruct_mock) + + result = await wrapper._retrieve_switchyard_trace(instance_config) + + assert result == "trace" + args, kwargs = request_mock.await_args + assert args == ("GET", "http://switchyard:4000/v1/sessions/0123456789abcdef0123456789abcdef/completions") + assert kwargs["timeout"].total == 60 + reconstruct_mock.assert_called_once_with(envelope, "0123456789abcdef0123456789abcdef", wrapper._vllm_converter) + ######################################## # _load_rebench_log_parsers tests diff --git a/responses_api_agents/swe_agents/tests/test_switchyard_trace.py b/responses_api_agents/swe_agents/tests/test_switchyard_trace.py new file mode 100644 index 0000000000..63a65a8b62 --- /dev/null +++ b/responses_api_agents/swe_agents/tests/test_switchyard_trace.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import pytest + +from nemo_gym.openai_utils import ( + NeMoGymEasyInputMessage, + NeMoGymFunctionCallOutput, + NeMoGymResponseFunctionToolCallForTraining, + NeMoGymResponseOutputMessageForTraining, +) +from nemo_gym.responses_converter import ResponsesConverter +from responses_api_agents.swe_agents.switchyard_trace import ( + SwitchyardTraceError, + map_switchyard_tools, + reconstruct_switchyard_rollout, +) + + +SESSION_ID = "abc123" + +TOOL_SCHEMA = {"type": "object", "properties": {"path": {"type": "string"}}} +TOOLS = [{"id": "str_replace_editor", "description": "Edit files", "inputSchema": {"jsonSchema": TOOL_SCHEMA}}] +TOOL_CALL = {"id": "call_1", "type": "function", "function": {"name": "str_replace_editor", "arguments": "{}"}} + +TRIPLE_0 = {"prompt_token_ids": [1, 2, 3], "generation_token_ids": [4, 5], "generation_log_probs": [-0.1, -0.2]} +TRIPLE_1 = {"prompt_token_ids": [1, 2, 3, 4, 5, 6], "generation_token_ids": [7], "generation_log_probs": [-0.3]} + + +def _record(i: int, messages: list, triple: dict, **overrides) -> dict: + record = { + "schema_version": 1, + "uuid": f"uuid-{i}", + "session_id": SESSION_ID, + "captured_at": f"2026-07-14T00:00:0{i}+00:00", + "request_id": f"req-{i}", + "model": "Qwen/Qwen3-0.6B", + "messages": messages, + "tools": TOOLS, + "tool_choice": "auto", + "finish_reason": "stop", + "is_valid": True, + **triple, + } + record.update(overrides) + return record + + +def _two_call_records() -> list: + """A two-call tool trajectory, with the same logical messages appearing in + both OpenHands' list-of-text-parts shape and Switchyard's captured string shape.""" + first = _record( + 0, + [ + {"role": "system", "content": "You are a SWE agent."}, + {"role": "user", "content": "Fix the bug."}, + {"role": "assistant", "tool_calls": [TOOL_CALL]}, + ], + TRIPLE_0, + ) + second = _record( + 1, + [ + {"role": "system", "content": [{"type": "text", "text": "You are a SWE agent."}]}, + {"role": "user", "content": [{"type": "text", "text": "Fix the bug."}]}, + {"role": "assistant", "content": None, "tool_calls": [TOOL_CALL | {"index": 0}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "edited ok"}, + {"role": "assistant", "content": "All fixed."}, + ], + TRIPLE_1, + ) + return [first, second] + + +def _envelope(records: list, **overrides) -> dict: + envelope = {"schema_version": 1, "session_id": SESSION_ID, "completions": records} + envelope.update(overrides) + return envelope + + +def _reconstruct(envelope: dict): + converter = ResponsesConverter(return_token_id_information=True) + return reconstruct_switchyard_rollout(envelope, SESSION_ID, converter) + + +class TestReconstruction: + def test_two_call_tool_trajectory(self) -> None: + trace = _reconstruct(_envelope(_two_call_records())) + + assert trace.record_uuids == ["uuid-0", "uuid-1"] + assert trace.model == "Qwen/Qwen3-0.6B" + + assert [type(item) for item in trace.input_items] == [NeMoGymEasyInputMessage, NeMoGymEasyInputMessage] + assert trace.input_items[0].role == "system" + assert trace.input_items[1].content == "Fix the bug." + + assert [type(item) for item in trace.output_items] == [ + NeMoGymResponseFunctionToolCallForTraining, + NeMoGymFunctionCallOutput, + NeMoGymResponseOutputMessageForTraining, + ] + tool_call, tool_output, final_message = trace.output_items + assert tool_call.call_id == "call_1" + assert tool_call.prompt_token_ids == TRIPLE_0["prompt_token_ids"] + assert tool_call.generation_token_ids == TRIPLE_0["generation_token_ids"] + assert tool_call.generation_log_probs == TRIPLE_0["generation_log_probs"] + assert tool_output.call_id == "call_1" + assert tool_output.output == "edited ok" + assert final_message.content[0].text == "All fixed." + assert final_message.prompt_token_ids == TRIPLE_1["prompt_token_ids"] + assert final_message.generation_token_ids == TRIPLE_1["generation_token_ids"] + assert final_message.generation_log_probs == TRIPLE_1["generation_log_probs"] + + assert len(trace.tools) == 1 + assert trace.tools[0].name == "str_replace_editor" + assert trace.tools[0].description == "Edit files" + assert trace.tools[0].parameters == TOOL_SCHEMA + assert trace.tools[0].type == "function" + + def test_single_record_session(self) -> None: + records = [_two_call_records()[0]] + trace = _reconstruct(_envelope(records)) + assert trace.record_uuids == ["uuid-0"] + assert len(trace.input_items) == 2 + assert len(trace.output_items) == 1 + + def test_divergent_history_rejected(self) -> None: + records = _two_call_records() + records[1]["messages"][1]["content"] = [{"type": "text", "text": "Fix a DIFFERENT bug."}] + with pytest.raises(SwitchyardTraceError, match="does not extend"): + _reconstruct(_envelope(records)) + + def test_missing_environment_suffix_rejected(self) -> None: + records = _two_call_records() + # Second prompt replays the first prompt + assistant but adds no tool result. + records[1]["messages"] = records[1]["messages"][:3] + [{"role": "assistant", "content": "All fixed."}] + with pytest.raises(SwitchyardTraceError, match="adds no environment messages"): + _reconstruct(_envelope(records)) + + def test_assistant_in_suffix_rejected(self) -> None: + records = _two_call_records() + records[1]["messages"].insert(4, {"role": "assistant", "content": "uncaptured turn"}) + records[1]["messages"].insert(5, {"role": "user", "content": "go on"}) + with pytest.raises(SwitchyardTraceError, match="non-environment roles"): + _reconstruct(_envelope(records)) + + def test_reordered_records_rejected(self) -> None: + records = list(reversed(_two_call_records())) + with pytest.raises(SwitchyardTraceError, match="does not extend"): + _reconstruct(_envelope(records)) + + +class TestValidation: + @pytest.mark.parametrize( + "mutate, match", + [ + (lambda e: e.update(schema_version=2), "schema_version"), + (lambda e: e.update(session_id="other"), "session_id"), + (lambda e: e.update(completions=[]), "no completions"), + (lambda e: e.update(completions="nope"), "no completions"), + ], + ) + def test_bad_envelope(self, mutate, match) -> None: + envelope = _envelope(_two_call_records()) + mutate(envelope) + with pytest.raises(SwitchyardTraceError, match=match): + _reconstruct(envelope) + + def test_non_object_envelope(self) -> None: + with pytest.raises(SwitchyardTraceError, match="not an object"): + _reconstruct([]) + + @pytest.mark.parametrize( + "field, value, match", + [ + ("schema_version", 2, "schema_version"), + ("session_id", "other", "session_id"), + ("uuid", "", "no uuid"), + ("uuid", "uuid-0", "duplicates uuid"), + ("is_valid", False, "is_valid"), + ("request_id", "", "empty request_id"), + ("model", "", "empty model"), + ("prompt_token_ids", [], "prompt_token_ids"), + ("prompt_token_ids", [1, "2"], "prompt_token_ids"), + ("generation_token_ids", None, "generation_token_ids"), + ("generation_log_probs", [-0.1, float("nan")], "finite"), + ("generation_log_probs", [-0.1, -0.2], "different lengths"), + ("messages", [], "no messages"), + ], + ) + def test_bad_record_fields(self, field, value, match) -> None: + records = _two_call_records() + records[1][field] = value + with pytest.raises(SwitchyardTraceError, match=match): + _reconstruct(_envelope(records)) + + def test_messages_must_end_with_assistant(self) -> None: + records = _two_call_records() + records[0]["messages"].append({"role": "user", "content": "trailing"}) + with pytest.raises(SwitchyardTraceError, match="end with an assistant"): + _reconstruct(_envelope(records)) + + @pytest.mark.parametrize( + "field, value", + [("model", "other-model"), ("tools", []), ("tool_choice", "required")], + ) + def test_session_consistency(self, field, value) -> None: + records = _two_call_records() + records[1][field] = value + with pytest.raises(SwitchyardTraceError): + _reconstruct(_envelope(records)) + + @pytest.mark.parametrize( + "message, match", + [ + ("nope", "non-object message"), + ({"role": "narrator", "content": "hm"}, "unsupported role"), + ({"role": "user", "content": {"nested": True}}, "unsupported content"), + ({"role": "user", "content": [{"type": "image_url", "image_url": {}}]}, "content part"), + ({"role": "tool", "content": "result"}, "without tool_call_id"), + ], + ) + def test_bad_messages(self, message, match) -> None: + records = _two_call_records() + records[1]["messages"].insert(3, message) + with pytest.raises(SwitchyardTraceError, match=match): + _reconstruct(_envelope(records)) + + def test_malformed_tool_call(self) -> None: + records = _two_call_records() + records[0]["messages"][-1]["tool_calls"] = [{"function": {"name": "x", "arguments": "{}"}}] + with pytest.raises(SwitchyardTraceError, match="malformed tool call"): + _reconstruct(_envelope(records)) + + +class TestToolMapping: + def test_tool_without_schema(self) -> None: + tools = map_switchyard_tools([{"id": "noop", "description": ""}]) + assert tools[0].name == "noop" + assert tools[0].parameters is None + assert tools[0].description is None + + def test_tool_without_id_rejected(self) -> None: + with pytest.raises(SwitchyardTraceError, match="no usable id"): + map_switchyard_tools([{"description": "no id"}]) + + def test_non_list_rejected(self) -> None: + with pytest.raises(SwitchyardTraceError, match="not a list"): + map_switchyard_tools({"id": "x"}) From 00258fe6c5efdaab12146bb465a68d0936a78768 Mon Sep 17 00:00:00 2001 From: Lin Jia Date: Tue, 14 Jul 2026 14:57:56 -0700 Subject: [PATCH 4/7] feat: zero-fork Switchyard transport for OpenHands runs Replace the planned nv-OpenHands client change with configuration Gym already controls, so the pinned fork commit needs no modification and no pins move: - write the Switchyard route id as the agent TOML model (Switchyard dispatches on it and substitutes the route target's model upstream) - ride the capture session on completion_kwargs.proxy_x_session_id, which lands verbatim in the pinned client's request body and which Switchyard strips before forwarding (body-field session support is added on the Switchyard PR #63 branch) - rewrite the model-server entry of the agent container's NEMO_GYM_CONFIG_DICT to Switchyard's host/port; the eval container keeps the original - validate switchyard_base_url as http://: at server startup (ServerClient can only target plain host:port) The pinned client keeps writing completion logs and metrics, so the fallback rollout and non-Switchyard behavior are unchanged. Doc updated to the zero-fork architecture. Co-Authored-By: Claude Fable 5 Signed-off-by: Lin Jia --- gym_switchyard_integration.md | 158 ++++++++---------- responses_api_agents/swe_agents/app.py | 80 +++++++-- .../swe_agents/tests/test_app.py | 76 +++++++-- 3 files changed, 199 insertions(+), 115 deletions(-) diff --git a/gym_switchyard_integration.md b/gym_switchyard_integration.md index ec4f86e258..c6635a559a 100644 --- a/gym_switchyard_integration.md +++ b/gym_switchyard_integration.md @@ -8,11 +8,12 @@ **Gym implementation branch:** `feature/switchyard-openhands-integration` -**Switchyard implementation:** `drifold/feature/token-capture` at `1365ca0` +**Switchyard implementation:** `drifold/feature/token-capture` at `eb351a7` ([NVIDIA-NeMo/Switchyard#63](https://github.com/NVIDIA-NeMo/Switchyard/pull/63)) **OpenHands base:** Gym's pinned [`nv-OpenHands` commit](https://github.com/sdevare-nv/nv-OpenHands/tree/5f0180054732945df08ad2293903e6873f0492b6) +— unchanged; this integration requires no OpenHands-side modification ## Objective @@ -99,6 +100,13 @@ list field is named `completions`. Records are returned ordered by (`captured_at`, `uuid`). The Gym adapter still validates message-history continuity and never repairs a broken chain using timestamps. +Since `eb351a7`, the branch also accepts the session id as a top-level +request-body field `proxy_x_session_id` (the header takes precedence; the field +is always stripped before the request is forwarded upstream). This serves clients whose only reachable knob is extra JSON body +fields — the pinned OpenHands client among them — and generalizes to any harness +that can set OpenAI-SDK/litellm `extra_body`, with no custom-header surface +required. + ## MVP scope The first implementation supports: @@ -182,7 +190,8 @@ swe_agents: When `switchyard_base_url` is absent, behavior remains unchanged. When present for an OpenHands run, policy calls go directly to Switchyard. -The URL must be reachable from both: +The URL must be exactly `http://:` (validated at server startup) and +reachable from both: - the Gym wrapper process, for retrieval; and - the OpenHands Apptainer container, for model calls. @@ -211,60 +220,42 @@ session-directory sanitizer, and is unique across retries. It is serialized in t existing `instance_config`, so `SWEBenchWrapper.run` can retrieve the same session after `responses` returns without introducing shared mutable state. -## OpenHands transport change - -Change the existing `NemoGymClient._post_completion` in the external -`nv-OpenHands` fork. Do not add a second client or a local compatibility proxy. - -The direct path is: - -1. Build `params` exactly as the client does today. -2. If `SWITCHYARD_BASE_URL` is absent, keep the existing - `ServerClient.post` behavior. -3. If it is present: - - require `SWITCHYARD_SESSION_ID`; - - copy `params`; - - replace `params["model"]` with `NEMO_GYM_MODEL_SERVER_NAME`; - - POST to `/v1/chat/completions`; - - send `proxy_x_session_id: `; - - make exactly one HTTP attempt — no transport retries. -4. Reuse the client's existing response-status handling, JSON parsing, cookies, - metrics, and completion-log writing. - -The single-attempt rule exists because an attempt that reaches vLLM writes a -capture record even when the client never receives the response. A transport -retry can therefore leave two records with the same prompt history; strict-history -reconstruction would refuse the session and mask the sample. Gym's -`ServerClient.post` and `nemo_gym.server_utils.request` both retry with backoff, -so the direct path must use the shared aiohttp client without the retry wrapper -and let a failed call fail the run through existing OpenHands error handling. - -The model rewrite is required. Switchyard dispatches on the incoming OpenAI -`model` field, while OpenHands normally sends the underlying LLM model string. -`NEMO_GYM_MODEL_SERVER_NAME` is already exported by the Gym harness and exactly -matches the `policy_model` route in the example. `params["model"]` always exists: -`_nemo_gym_llm_kwargs` is built with `model=self.config.model`. - -No message-format change is needed. On the direct path the response message never -carries `prompt_token_ids`, `generation_token_ids`, or `generation_log_probs` — -those are attached by Gym's model server, not vLLM — so the client's existing -token-field stripping loop is a no-op and every captured prompt history stays -strictly extending. - -Gym must export two additional values only into the agent container: - -~~~text -SWITCHYARD_BASE_URL= -SWITCHYARD_SESSION_ID= -~~~ - -The evaluator container does not make policy calls and does not receive these values. - -After the fork change is pushed, update every OpenHands pin in Gym: - -- `responses_api_agents/swe_agents/configs/swebench_openhands.yaml` -- `responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml` -- `responses_api_agents/swe_agents/configs/swebench_multi_tools.yaml` +## OpenHands transport (no fork change) + +No nv-OpenHands modification is required, and the three OpenHands YAML pins do +not move. The pinned client already provides everything the transport needs, +reachable entirely through configuration Gym writes on every run: + +1. `NemoGymClient._post_completion` posts to the server named + `NEMO_GYM_MODEL_SERVER_NAME`, resolving its host/port from the + `NEMO_GYM_CONFIG_DICT` the harness exports. Gym rewrites that model-server + entry to Switchyard's host/port — in the agent container's copy only. The + eval container keeps the original config dict and makes no policy calls. +2. The posted `model` field echoes the TOML `[llm.model] model`. Gym writes the + Switchyard route id there; Switchyard dispatches on it and substitutes the + route target's real model upstream. +3. The TOML `completion_kwargs` (a stock OpenHands `LLMConfig` field) merge + verbatim into the posted JSON body. Gym rides the capture session on the + `proxy_x_session_id` body field, which Switchyard strips before forwarding. + +Because `ServerClient` can only target plain `http://:`, +`switchyard_base_url` must have exactly that form; Gym validates it at server +startup and fails fast rather than masking every sample. + +The pinned client keeps writing its per-call completion logs and +`NEMO_GYM_METRICS_FPATH` metrics, so the existing fallback rollout, metrics, and +non-Switchyard behavior are unchanged. + +Token-field handling needs no attention: Switchyard's translation layer drops +vLLM's engine token fields from the client-facing response, so the client's +token bookkeeping is inert and every captured prompt history stays strictly +extending. + +Transport retries: the client's `ServerClient` path retries only on +connection-level errors. A retry after a call that actually reached vLLM writes +a second capture record whose prompt history does not extend the first; +reconstruction then refuses the session and masks the sample — fail-closed, +never partially annotated. ## Gym retrieval and adaptation @@ -391,42 +382,34 @@ of this MVP. ### Gym 1. `responses_api_agents/swe_agents/app.py` - - add `switchyard_base_url`; + - add `switchyard_base_url` (validated as `http://:` at startup); - add and generate `switchyard_session_id`; - - export the two Switchyard environment variables next to the existing - `NEMO_GYM_MODEL_SERVER_NAME` export in `OpenHandsHarnessProcessor`'s - agent command (agent container only; the opencode harness and the eval - container command are untouched); + - in `OpenHandsHarnessProcessor.get_run_command`, when configured: write the + route id as the TOML model, add `completion_kwargs.proxy_x_session_id`, and + export a config dict whose model-server entry points at Switchyard + (`_parse_switchyard_base_url`, `_switchyard_ng_config_dict_str`; agent + container only — the opencode harness and the eval container are untouched); - retrieve and apply the reconstructed trace in `run`; - add `switchyard_trace_error`. 2. `responses_api_agents/swe_agents/switchyard_trace.py` - DTOs, validation, tool mapping, and reconstruction. 3. `responses_api_agents/swe_agents/tests/test_app.py` - - configuration, UUID creation, environment export, retrieval integration, and - masking behavior. + - configuration and URL validation, UUID creation, TOML/config-dict routing, + retrieval integration, and masking behavior. 4. `responses_api_agents/swe_agents/tests/test_switchyard_trace.py` - focused adapter and validation tests. -5. The three OpenHands YAML files listed above - - update the fork commit after the external change lands. Do not change `nemo_gym/rollout_collection.py` or generic model-server code. ### nv-OpenHands -1. Update `openhands/agenthub/nemo_gym_client.py`. -2. Add focused client tests for: - - fallback to Gym `ServerClient`; - - direct URL construction; - - route-model rewrite; - - session header; - - missing-session failure; - - exactly one HTTP attempt per call (no transport retries); - - response parsing, cookies, metrics, and completion logging. -3. Push the fork commit before updating Gym's pin. +Nothing. The pinned fork commit is unchanged and no pins move. ### Switchyard -No Gym-specific Switchyard feature is required beyond PR #63. +One addition on the PR #63 branch: `TokenCaptureRequestProcessor` also resolves +the session from the `proxy_x_session_id` request-body field (header takes +precedence) and always strips it before forwarding upstream, with focused tests. Two correctness limitations remain in that branch: @@ -443,10 +426,12 @@ The parser and target schema already agree on nested Minimum coverage: 1. No Switchyard setting preserves current OpenHands behavior exactly. -2. A configured run creates a unique UUID and exports it only to the agent - container. -3. The OpenHands client selects direct HTTP, rewrites the route model, and sends the - session header. +2. A configured run creates a unique UUID, writes it into the agent TOML as + `completion_kwargs.proxy_x_session_id`, and points the agent container's + config-dict model-server entry at Switchyard; the eval container keeps the + original config dict. +3. Switchyard resolves the session from the body field, strips it before + forwarding upstream, and the header keeps precedence. 4. A valid two-call tool trajectory reconstructs into one Gym rollout with two exact token triples and unchanged reward. 5. Invalid schema, wrong session, duplicate UUID, empty tokens, non-finite @@ -473,13 +458,12 @@ uv run pytest -q \ ## Implementation order -1. Change and test the pinned `nv-OpenHands` client. -2. Push the fork commit and update Gym's three pins. -3. Add Gym config/session plumbing and agent-only environment export. -4. Add the Gym trace DTO, validation, and reconstruction module. -5. Integrate retrieval and fail-closed masking in `SWEBenchWrapper.run`. -6. Run focused unit tests. -7. Run one example end to end with a fresh Switchyard log directory. +1. Add the Switchyard body-field session change and run the capture suite. +2. Add Gym config/session plumbing and the agent-only TOML/config-dict routing. +3. Add the Gym trace DTO, validation, and reconstruction module. +4. Integrate retrieval and fail-closed masking in `SWEBenchWrapper.run`. +5. Run focused unit tests. +6. Run one example end to end with a fresh Switchyard log directory. ## Done when diff --git a/responses_api_agents/swe_agents/app.py b/responses_api_agents/swe_agents/app.py index 3f35626da2..0692ea272f 100644 --- a/responses_api_agents/swe_agents/app.py +++ b/responses_api_agents/swe_agents/app.py @@ -37,6 +37,7 @@ from subprocess import run as subprocess_run from traceback import format_exc from typing import Any, Dict, List, Literal, NamedTuple, Optional, Tuple, Union +from urllib.parse import urlparse import ray import tomlkit @@ -182,10 +183,12 @@ class SWEBenchWrapperConfig(BaseResponsesAPIAgentConfig): switchyard_base_url: Optional[str] = Field( default=None, description=( - "Base URL of a token-capture-enabled Switchyard proxy (e.g. http://host:4000). When set for an " - "OpenHands run, agent policy calls go directly to Switchyard and the captured per-call token " - "records are retrieved after the agent finishes to build the training rollout. Must be reachable " - "from both this process and the agent container. When absent, behavior is unchanged." + "Base URL of a token-capture-enabled Switchyard proxy, exactly http://:. When set " + "for an OpenHands run, agent policy calls are routed to Switchyard (via the agent container's " + "NEMO_GYM_CONFIG_DICT and oh_config.toml — no OpenHands change needed) and the captured " + "per-call token records are retrieved after the agent finishes to build the training rollout. " + "Must be reachable from both this process and the agent container. When absent, behavior is " + "unchanged." ), ) @@ -1515,6 +1518,36 @@ def postprocess_after_run(self, report_file: Path) -> None: report_path.write_text(json.dumps(report, indent=2)) +def _parse_switchyard_base_url(switchyard_base_url: str) -> Tuple[str, int]: + """Host/port of the Switchyard proxy. + + The agent container reaches Switchyard through the pinned OpenHands client's + ServerClient path, which can only target ``http://:`` — so the + configured base URL must be exactly that shape. + """ + parsed = urlparse(switchyard_base_url) + if parsed.scheme != "http" or not parsed.hostname or not parsed.port or parsed.path.strip("/") or parsed.query: + raise ValueError(f"switchyard_base_url must have the form http://:, got {switchyard_base_url!r}") + return parsed.hostname, parsed.port + + +def _switchyard_ng_config_dict_str( + ng_global_config_dict_str: str, model_server_name: str, switchyard_base_url: str +) -> str: + """NEMO_GYM_CONFIG_DICT (shell-quoted YAML) for a Switchyard-routed agent container. + + Rewrites the model server's host/port to Switchyard's, so the in-container + client's existing ``ServerClient.post(server_name=...)`` lands on the proxy + with no OpenHands change. + """ + host, port = _parse_switchyard_base_url(switchyard_base_url) + global_config_dict = OmegaConf.create(shlex.split(ng_global_config_dict_str)[0]) + server_config_dict = get_first_server_config_dict(global_config_dict, model_server_name) + server_config_dict.host = host + server_config_dict.port = port + return shlex.quote(OmegaConf.to_yaml(global_config_dict)) + + class OpenHandsHarnessProcessor(BaseDatasetHarnessProcessor): def _sync_openhands_to_config_commit(self, openhands_dir: Path) -> None: """Ensure OpenHands checkout matches config.agent_framework_commit. @@ -1610,6 +1643,27 @@ def get_run_command(self) -> ExecuteContainerCommandArgs: "top_p": self.config.inference_params["top_p"], } + # Zero-fork Switchyard transport. The pinned OpenHands client posts to the + # server named NEMO_GYM_MODEL_SERVER_NAME with the TOML model string, and + # litellm `completion_kwargs` land verbatim in its request body. So: make + # the model string the Switchyard route id (Switchyard dispatches on it), + # ride the capture session on the body field Switchyard strips before + # forwarding, and point the model-server entry of the agent container's + # NEMO_GYM_CONFIG_DICT at Switchyard. Agent container only: the eval + # container makes no policy calls and keeps the original config dict. + if self.config.switchyard_base_url and self.config.switchyard_session_id: + config["llm"]["model"]["model"] = self.config.model_server_name + config["llm"]["model"]["completion_kwargs"] = { + "proxy_x_session_id": self.config.switchyard_session_id, + } + ng_config_dict_str = _switchyard_ng_config_dict_str( + self.config.ng_global_config_dict_str, + self.config.model_server_name, + self.config.switchyard_base_url, + ) + else: + ng_config_dict_str = self.config.ng_global_config_dict_str + config_str = tomlkit.dumps(config) eval_dir_in_openhands = self.config.eval_dir_in_openhands @@ -1702,17 +1756,6 @@ def get_run_command(self) -> ExecuteContainerCommandArgs: else: camel_case_tool_names_cmd = "" - # Agent container only: the eval container makes no policy calls and must - # not receive these. The in-container NemoGymClient switches to direct - # Switchyard HTTP when SWITCHYARD_BASE_URL is set. - if self.config.switchyard_base_url and self.config.switchyard_session_id: - switchyard_cmd = ( - f"export SWITCHYARD_BASE_URL={shlex.quote(self.config.switchyard_base_url)} && " - f"export SWITCHYARD_SESSION_ID={self.config.switchyard_session_id} && " - ) - else: - switchyard_cmd = "" - workspace_check_cmd = "" # Run the same baseline dependency/env repair the eval uses (if any), in a @@ -1743,9 +1786,8 @@ def get_run_command(self) -> ExecuteContainerCommandArgs: f"{log_cmd}" f"{profiling_cmd}" f"export NEMO_GYM_METRICS_FPATH={self.config.base_mounted_dir}/nemo_gym_metrics.json && " - f"export NEMO_GYM_CONFIG_DICT={self.config.ng_global_config_dict_str} && " + f"export NEMO_GYM_CONFIG_DICT={ng_config_dict_str} && " f"export NEMO_GYM_MODEL_SERVER_NAME={self.config.model_server_name} && " - f"{switchyard_cmd}" "export VIRTUAL_ENV=/openhands_setup/OpenHands/.venv && " "export PATH=$PATH:/openhands_setup/OpenHands/.venv/bin && " # CRITICAL: Configure poetry to only use the OpenHands venv (ignore external venvs) @@ -2790,6 +2832,10 @@ class SWEBenchWrapper(SimpleResponsesAPIAgent): ######################################## def model_post_init(self, context: Any) -> None: + # Fail fast on a malformed Switchyard URL rather than masking every sample. + if self.config.switchyard_base_url: + _parse_switchyard_base_url(self.config.switchyard_base_url) + run_session_id = f"{int(time.time() * 1000)}_{str(uuid.uuid4())[:8]}" workspace_root = Path(__file__).parent # Only set up the agent harness that's actually selected. Both share the diff --git a/responses_api_agents/swe_agents/tests/test_app.py b/responses_api_agents/swe_agents/tests/test_app.py index c90b3ddc6c..484450f0cb 100644 --- a/responses_api_agents/swe_agents/tests/test_app.py +++ b/responses_api_agents/swe_agents/tests/test_app.py @@ -15,6 +15,7 @@ import asyncio import json import os +import shlex import shutil import tempfile import time @@ -54,8 +55,10 @@ SWEBenchWrapperServerConfig, SWERebenchDatasetProcessor, _extract_instance_dict, + _parse_switchyard_base_url, _render_opencode_user_message, _resolve_opencode_workspace_path, + _switchyard_ng_config_dict_str, file_lock, runner_ray_remote, update_and_read_metrics, @@ -176,6 +179,15 @@ def _make_instance_config(tmpdir: str, **overrides) -> SWEBenchWrapperInstanceCo return SWEBenchWrapperInstanceConfig(**defaults) +def _ng_config_dict_str(host: str = "model-host", port: int = 9999) -> str: + """Shell-quoted NEMO_GYM_CONFIG_DICT YAML with a resolvable test_model entry.""" + return shlex.quote( + OmegaConf.to_yaml( + OmegaConf.create({"test_model": {"responses_api_models": {"vllm_model": {"host": host, "port": port}}}}) + ) + ) + + ######################################## # Config model tests ######################################## @@ -938,27 +950,64 @@ def test_get_run_command_camel_case_tool_names(self) -> None: processor.get_run_command() assert "CAMEL_CASE_TOOL_NAMES=true" in self._read_agent_script(config) - def test_get_run_command_switchyard_env(self) -> None: + def test_get_run_command_switchyard_transport(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: config = _make_instance_config( tmpdir, switchyard_base_url="http://switchyard:4000", switchyard_session_id="0123456789abcdef0123456789abcdef", + ng_global_config_dict_str=_ng_config_dict_str(), ) config.persistent_dir.mkdir(parents=True, exist_ok=True) processor = OpenHandsHarnessProcessor(config=config) processor.get_run_command() script = self._read_agent_script(config) - assert "export SWITCHYARD_BASE_URL=http://switchyard:4000 && " in script - assert "export SWITCHYARD_SESSION_ID=0123456789abcdef0123456789abcdef && " in script + # The TOML model becomes the Switchyard route id and carries the + # capture session as a body field via litellm completion_kwargs. + assert 'model = "test_model"' in script + assert 'proxy_x_session_id = "0123456789abcdef0123456789abcdef"' in script + # The agent's NEMO_GYM_CONFIG_DICT points the model server at Switchyard. + assert "host: switchyard" in script + assert "port: 4000" in script - def test_get_run_command_no_switchyard_env_by_default(self) -> None: + def test_get_run_command_no_switchyard_by_default(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: - config = _make_instance_config(tmpdir) + config = _make_instance_config(tmpdir, ng_global_config_dict_str=_ng_config_dict_str()) config.persistent_dir.mkdir(parents=True, exist_ok=True) processor = OpenHandsHarnessProcessor(config=config) processor.get_run_command() - assert "SWITCHYARD" not in self._read_agent_script(config) + script = self._read_agent_script(config) + assert 'model = "test-model"' in script + assert "proxy_x_session_id" not in script + assert "host: switchyard" not in script + + +class TestSwitchyardTransportConfig: + @pytest.mark.parametrize( + "url", + ["https://switchyard:4000", "http://switchyard:4000/v1", "http://switchyard", "switchyard:4000"], + ) + def test_rejects_non_host_port_urls(self, url) -> None: + with pytest.raises(ValueError, match="http://:"): + _parse_switchyard_base_url(url) + + def test_accepts_host_port(self) -> None: + assert _parse_switchyard_base_url("http://switchyard:4000") == ("switchyard", 4000) + assert _parse_switchyard_base_url("http://10.0.0.5:4000/") == ("10.0.0.5", 4000) + + def test_config_dict_rewrite_points_model_server_at_switchyard(self) -> None: + rewritten = _switchyard_ng_config_dict_str(_ng_config_dict_str(), "test_model", "http://switchyard:4000") + loaded = OmegaConf.create(shlex.split(rewritten)[0]) + assert loaded.test_model.responses_api_models.vllm_model.host == "switchyard" + assert loaded.test_model.responses_api_models.vllm_model.port == 4000 + + def test_wrapper_fails_fast_on_bad_url(self, monkeypatch) -> None: + monkeypatch.setattr(swe_app, "get_global_config_dict", MagicMock(return_value=OmegaConf.create({}))) + monkeypatch.setattr(BaseDatasetHarnessProcessor, "_run_setup_command", MagicMock(return_value=None)) + config = _minimal_server_config() + config.switchyard_base_url = "http://switchyard:4000/v1" + with pytest.raises(ValueError, match="http://:"): + SWEBenchWrapper(config=config, server_client=MagicMock(spec=ServerClient)) ######################################## @@ -2336,15 +2385,17 @@ def test_setup_params_switchyard_session(self, monkeypatch) -> None: wrapper.config.container_formatter = [str(Path(tmpdir) / "{instance_id}.sif")] self._setup_oh_dirs(wrapper) wrapper.config.switchyard_base_url = "http://switchyard:4000" + wrapper._swe_bench_wrapper_server_config.ng_global_config_dict_str = _ng_config_dict_str() params, _ = wrapper._setup_params(self._switchyard_body()) session_id = params.switchyard_session_id assert session_id is not None assert len(session_id) == 32 assert set(session_id) <= set("0123456789abcdef") - # Exported into the agent container script. - assert f"export SWITCHYARD_SESSION_ID={session_id} && " in params.agent_script - assert "export SWITCHYARD_BASE_URL=http://switchyard:4000 && " in params.agent_script + # Routed into the agent container via the TOML and the rewritten + # NEMO_GYM_CONFIG_DICT — no OpenHands-side changes involved. + assert f'proxy_x_session_id = "{session_id}"' in params.agent_script + assert "host: switchyard" in params.agent_script # Unique per run. params2, _ = wrapper._setup_params(self._switchyard_body()) @@ -2356,10 +2407,12 @@ def test_setup_params_no_switchyard_by_default(self, monkeypatch) -> None: (Path(tmpdir) / "django__django-12345.sif").touch() wrapper.config.container_formatter = [str(Path(tmpdir) / "{instance_id}.sif")] self._setup_oh_dirs(wrapper) + wrapper._swe_bench_wrapper_server_config.ng_global_config_dict_str = _ng_config_dict_str() params, _ = wrapper._setup_params(self._switchyard_body()) assert params.switchyard_session_id is None - assert "SWITCHYARD" not in params.agent_script + assert "proxy_x_session_id" not in params.agent_script + assert "host: switchyard" not in params.agent_script def test_setup_params_no_switchyard_for_golden_patch(self, monkeypatch) -> None: wrapper = _create_wrapper(monkeypatch) @@ -2369,10 +2422,11 @@ def test_setup_params_no_switchyard_for_golden_patch(self, monkeypatch) -> None: self._setup_oh_dirs(wrapper) wrapper.config.switchyard_base_url = "http://switchyard:4000" wrapper.config.verify_golden_patch = True + wrapper._swe_bench_wrapper_server_config.ng_global_config_dict_str = _ng_config_dict_str() params, _ = wrapper._setup_params(self._switchyard_body()) assert params.switchyard_session_id is None - assert "SWITCHYARD" not in params.agent_script + assert "proxy_x_session_id" not in params.agent_script class TestSWEBenchWrapperResponses: From 3d88104abad98bb59cffc1b60525e703f4706aa7 Mon Sep 17 00:00:00 2001 From: Lin Jia Date: Tue, 14 Jul 2026 15:19:06 -0700 Subject: [PATCH 5/7] docs: advertise switchyard_base_url opt-in in OpenHands configs Co-Authored-By: Claude Fable 5 Signed-off-by: Lin Jia --- .../swe_agents/configs/swebench_multi_tools.yaml | 3 +++ .../swe_agents/configs/swebench_openhands.yaml | 3 +++ .../swe_agents/configs/swebench_openhands_training.yaml | 6 ++++++ 3 files changed, 12 insertions(+) diff --git a/responses_api_agents/swe_agents/configs/swebench_multi_tools.yaml b/responses_api_agents/swe_agents/configs/swebench_multi_tools.yaml index 637cfcf3cd..4f833bb4db 100644 --- a/responses_api_agents/swe_agents/configs/swebench_multi_tools.yaml +++ b/responses_api_agents/swe_agents/configs/swebench_multi_tools.yaml @@ -13,6 +13,9 @@ swe_agents: agent_max_turns: 100 agent_framework_repo: https://github.com/sdevare-nv/nv-OpenHands.git agent_framework_commit: 5f0180054732945df08ad2293903e6873f0492b6 # pragma: allowlist secret + # Optional: route policy calls through a token-capture Switchyard proxy and + # rebuild the rollout with exact token ids (see gym_switchyard_integration.md). + # switchyard_base_url: http://: # Container configuration container_formatter: ??? diff --git a/responses_api_agents/swe_agents/configs/swebench_openhands.yaml b/responses_api_agents/swe_agents/configs/swebench_openhands.yaml index bdfa6cc27d..d0dae762e4 100644 --- a/responses_api_agents/swe_agents/configs/swebench_openhands.yaml +++ b/responses_api_agents/swe_agents/configs/swebench_openhands.yaml @@ -10,6 +10,9 @@ swe_agents: agent_max_turns: 100 agent_framework_repo: https://github.com/sdevare-nv/nv-OpenHands.git agent_framework_commit: 5f0180054732945df08ad2293903e6873f0492b6 # pragma: allowlist secret + # Optional: route policy calls through a token-capture Switchyard proxy and + # rebuild the rollout with exact token ids (see gym_switchyard_integration.md). + # switchyard_base_url: http://: # Container configuration container_formatter: ??? diff --git a/responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml b/responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml index a8531d3439..329f242ccb 100644 --- a/responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml +++ b/responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml @@ -9,6 +9,9 @@ swe_agents_train: agent_max_turns: 100 agent_framework_repo: https://github.com/sdevare-nv/nv-OpenHands.git agent_framework_commit: 5f0180054732945df08ad2293903e6873f0492b6 # pragma: allowlist secret + # Optional: route policy calls through a token-capture Switchyard proxy and + # rebuild the rollout with exact token ids (see gym_switchyard_integration.md). + # switchyard_base_url: http://: # Container configuration container_formatter: ??? container_folder_path: null @@ -61,6 +64,9 @@ swe_agents_val: agent_max_turns: 100 agent_framework_repo: https://github.com/sdevare-nv/nv-OpenHands.git agent_framework_commit: 5f0180054732945df08ad2293903e6873f0492b6 # pragma: allowlist secret + # Optional: route policy calls through a token-capture Switchyard proxy and + # rebuild the rollout with exact token ids (see gym_switchyard_integration.md). + # switchyard_base_url: http://: # Container configuration container_formatter: ??? container_folder_path: null From 6bb23e024dc1b3307e22836036ea71bab874ea55 Mon Sep 17 00:00:00 2001 From: Lin Jia Date: Tue, 14 Jul 2026 15:24:24 -0700 Subject: [PATCH 6/7] refactor: promote switchyard_trace to nemo_gym as shared infra The module is harness-neutral by design (depends only on Switchyard's record schema and the shared Responses converter), so it lives in the core library where any agent server can reuse it for Switchyard trace collection; its tests move to tests/unit_tests accordingly. Co-Authored-By: Claude Fable 5 Signed-off-by: Lin Jia --- gym_switchyard_integration.md | 13 ++++++++----- .../swe_agents => nemo_gym}/switchyard_trace.py | 0 responses_api_agents/swe_agents/app.py | 2 +- responses_api_agents/swe_agents/tests/test_app.py | 2 +- .../unit_tests}/test_switchyard_trace.py | 2 +- 5 files changed, 11 insertions(+), 8 deletions(-) rename {responses_api_agents/swe_agents => nemo_gym}/switchyard_trace.py (100%) rename {responses_api_agents/swe_agents/tests => tests/unit_tests}/test_switchyard_trace.py (99%) diff --git a/gym_switchyard_integration.md b/gym_switchyard_integration.md index c6635a559a..cc9c3d4f6c 100644 --- a/gym_switchyard_integration.md +++ b/gym_switchyard_integration.md @@ -259,8 +259,10 @@ never partially annotated. ## Gym retrieval and adaptation -Add a small module at -`responses_api_agents/swe_agents/switchyard_trace.py`. Keep network orchestration +Add a small module at `nemo_gym/switchyard_trace.py`. It is shared +infrastructure: it depends only on Switchyard's record schema and Gym's +converter — nothing harness-specific — so any agent server can reuse it for +Switchyard trace collection. Keep network orchestration in `SWEBenchWrapper.run`; keep schema validation and conversion in the module. After `SWEBenchWrapper.responses` returns, `run` already reads the serialized @@ -391,12 +393,13 @@ of this MVP. container only — the opencode harness and the eval container are untouched); - retrieve and apply the reconstructed trace in `run`; - add `switchyard_trace_error`. -2. `responses_api_agents/swe_agents/switchyard_trace.py` - - DTOs, validation, tool mapping, and reconstruction. +2. `nemo_gym/switchyard_trace.py` + - DTOs, validation, tool mapping, and reconstruction — harness-neutral + shared infrastructure. 3. `responses_api_agents/swe_agents/tests/test_app.py` - configuration and URL validation, UUID creation, TOML/config-dict routing, retrieval integration, and masking behavior. -4. `responses_api_agents/swe_agents/tests/test_switchyard_trace.py` +4. `tests/unit_tests/test_switchyard_trace.py` - focused adapter and validation tests. Do not change `nemo_gym/rollout_collection.py` or generic model-server code. diff --git a/responses_api_agents/swe_agents/switchyard_trace.py b/nemo_gym/switchyard_trace.py similarity index 100% rename from responses_api_agents/swe_agents/switchyard_trace.py rename to nemo_gym/switchyard_trace.py diff --git a/responses_api_agents/swe_agents/app.py b/responses_api_agents/swe_agents/app.py index 0692ea272f..4acdeb3b3e 100644 --- a/responses_api_agents/swe_agents/app.py +++ b/responses_api_agents/swe_agents/app.py @@ -65,7 +65,7 @@ ) from nemo_gym.profiling import Profiler from nemo_gym.server_utils import get_first_server_config_dict, get_response_json, raise_for_status, request -from responses_api_agents.swe_agents.switchyard_trace import SwitchyardTrace, reconstruct_switchyard_rollout +from nemo_gym.switchyard_trace import SwitchyardTrace, reconstruct_switchyard_rollout from responses_api_models.vllm_model.app import VLLMConverter, split_responses_input_output_items diff --git a/responses_api_agents/swe_agents/tests/test_app.py b/responses_api_agents/swe_agents/tests/test_app.py index 484450f0cb..567939fbc7 100644 --- a/responses_api_agents/swe_agents/tests/test_app.py +++ b/responses_api_agents/swe_agents/tests/test_app.py @@ -35,6 +35,7 @@ NeMoGymResponseOutputText, ) from nemo_gym.server_utils import ServerClient +from nemo_gym.switchyard_trace import SwitchyardTrace, SwitchyardTraceError from responses_api_agents.swe_agents.app import ( ActiveContainerCommand, AgentPromptOverride, @@ -63,7 +64,6 @@ runner_ray_remote, update_and_read_metrics, ) -from responses_api_agents.swe_agents.switchyard_trace import SwitchyardTrace, SwitchyardTraceError SWE_AGENTS_DIR = Path(__file__).resolve().parent.parent diff --git a/responses_api_agents/swe_agents/tests/test_switchyard_trace.py b/tests/unit_tests/test_switchyard_trace.py similarity index 99% rename from responses_api_agents/swe_agents/tests/test_switchyard_trace.py rename to tests/unit_tests/test_switchyard_trace.py index 63a65a8b62..38b16ff71a 100644 --- a/responses_api_agents/swe_agents/tests/test_switchyard_trace.py +++ b/tests/unit_tests/test_switchyard_trace.py @@ -21,7 +21,7 @@ NeMoGymResponseOutputMessageForTraining, ) from nemo_gym.responses_converter import ResponsesConverter -from responses_api_agents.swe_agents.switchyard_trace import ( +from nemo_gym.switchyard_trace import ( SwitchyardTraceError, map_switchyard_tools, reconstruct_switchyard_rollout, From 188061ec1b5f50afe29a2083b4620c09a1d82d8c Mon Sep 17 00:00:00 2001 From: Lin Jia Date: Tue, 14 Jul 2026 16:09:52 -0700 Subject: [PATCH 7/7] feat: surface mask_sample on BaseVerifyResponse and guard token contiguity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two trainer-boundary hardenings: - mask_sample becomes a first-class top-level field on BaseVerifyResponse (default false). swe_agents populates it from its instance config; anyterminal_agent already surfaced it via TerminalBenchMetrics. This gives trainers one documented place to read the unreliable-reward signal instead of agent-specific config layouts (NeMo RL maps it to a zero loss multiplier). - Switchyard trace reconstruction now validates that each record's prompt token ids extend the prior record's prompt+generation tokens — the exact contiguity the trainer hard-asserts when concatenating the triples. A violating session masks the sample instead of crashing the training step. Co-Authored-By: Claude Fable 5 Signed-off-by: Lin Jia --- gym_switchyard_integration.md | 6 +++++- nemo_gym/base_resources_server.py | 5 +++++ nemo_gym/switchyard_trace.py | 11 +++++++++++ responses_api_agents/swe_agents/app.py | 1 + responses_api_agents/swe_agents/tests/test_app.py | 3 +++ tests/unit_tests/test_switchyard_trace.py | 11 ++++++++++- 6 files changed, 35 insertions(+), 2 deletions(-) diff --git a/gym_switchyard_integration.md b/gym_switchyard_integration.md index cc9c3d4f6c..df5a5f10df 100644 --- a/gym_switchyard_integration.md +++ b/gym_switchyard_integration.md @@ -291,6 +291,8 @@ Reject the session unless: - every record has non-empty `request_id` and `model`; - prompt and generation IDs are non-empty integer arrays; - generation IDs and finite generation log probabilities have equal lengths; +- each record's prompt token ids extend the prior record's prompt plus + generation token ids (the trainer's contiguity requirement); - messages are non-empty and end with an assistant message; and - model, tools, and tool choice are consistent across the session. @@ -348,7 +350,9 @@ Instead: - keep the current OpenHands-derived response; - set the existing `SWEBenchWrapperInstanceConfig.mask_sample` flag to true — this - is Gym's established training-mask convention, already used for agent failures; + is Gym's established training-mask convention, already used for agent failures — + and surface it as the top-level `mask_sample` field on the verify response + (`BaseVerifyResponse`), which NeMo RL maps to a zero loss multiplier; - add `switchyard_trace_error: Optional[str]` to `SWEBenchVerifyResponse`, following the existing `agent_error_kind`-style metrics naming; - store a concise error containing the session ID and reason; and diff --git a/nemo_gym/base_resources_server.py b/nemo_gym/base_resources_server.py index 4d519aca7b..73c70a1433 100644 --- a/nemo_gym/base_resources_server.py +++ b/nemo_gym/base_resources_server.py @@ -90,6 +90,11 @@ class BaseVerifyRequest(BaseRunRequest): class BaseVerifyResponse(BaseVerifyRequest): reward: float + # Set by agent servers when the rollout's reward is unreliable (e.g. agent or + # eval timeout, OOM-killed container, failed trace capture). Trainers must + # exclude the sample from the gradient (NeMo RL zeroes its loss multiplier) + # rather than train on a garbage reward. + mask_sample: bool = False class BaseSeedSessionRequest(BaseModel): diff --git a/nemo_gym/switchyard_trace.py b/nemo_gym/switchyard_trace.py index 73f6f2b274..6bbe41c917 100644 --- a/nemo_gym/switchyard_trace.py +++ b/nemo_gym/switchyard_trace.py @@ -97,8 +97,19 @@ def reconstruct_switchyard_rollout(envelope: Any, session_id: str, converter: Re assembled: List[Dict[str, Any]] = [] annotated: List[Dict[str, Any]] = [] record_uuids: List[str] = [] + # Token-level counterpart of the message-history check: trainers concatenate + # the per-call triples into one sequence and hard-assert that each call's + # prompt tokens extend the previous prompt+generation. Validate it here so a + # violation (e.g. a template that re-renders history differently) masks the + # sample instead of crashing the training step. + token_history: List[int] = [] for i, record in enumerate(records): _validate_record(record, i, session_id, record_uuids) + if record["prompt_token_ids"][: len(token_history)] != token_history: + raise SwitchyardTraceError( + f"record {i} prompt_token_ids do not extend the prior prompt+generation tokens" + ) + token_history = record["prompt_token_ids"] + record["generation_token_ids"] if record["model"] != model: raise SwitchyardTraceError(f"record {i} model {record['model']!r} != session model {model!r}") if record["tools"] != tools or record.get("tool_choice") != tool_choice: diff --git a/responses_api_agents/swe_agents/app.py b/responses_api_agents/swe_agents/app.py index 4acdeb3b3e..e52dc9e8eb 100644 --- a/responses_api_agents/swe_agents/app.py +++ b/responses_api_agents/swe_agents/app.py @@ -3802,6 +3802,7 @@ async def run(self, body: BaseRunRequest) -> SWEBenchVerifyResponse: responses_create_params=responses_create_params, response=response, reward=1.0 if metrics.resolved else 0.0, + mask_sample=instance_config.mask_sample, **metrics.model_dump(), instance_config=instance_config.model_dump(), subagent_trajectories=subagent_trajectories, diff --git a/responses_api_agents/swe_agents/tests/test_app.py b/responses_api_agents/swe_agents/tests/test_app.py index 567939fbc7..fd39f56b5d 100644 --- a/responses_api_agents/swe_agents/tests/test_app.py +++ b/responses_api_agents/swe_agents/tests/test_app.py @@ -2672,6 +2672,7 @@ async def test_run_switchyard_success_replaces_rollout(self, monkeypatch) -> Non result = await wrapper.run(self._run_body()) assert result.switchyard_trace_error is None + assert result.mask_sample is False assert result.instance_config.mask_sample is False assert result.reward == 1.0 assert [item.model_dump() for item in result.responses_create_params.input] == [ @@ -2704,6 +2705,7 @@ async def test_run_switchyard_failure_masks_sample(self, monkeypatch) -> None: result = await wrapper.run(self._run_body()) # Fail closed for training, open for diagnostics. + assert result.mask_sample is True assert result.instance_config.mask_sample is True assert result.switchyard_trace_error == ( "session 0123456789abcdef0123456789abcdef: SwitchyardTraceError: " @@ -2742,6 +2744,7 @@ async def test_run_without_switchyard_does_not_retrieve(self, monkeypatch) -> No retrieve.assert_not_called() assert result.switchyard_trace_error is None + assert result.mask_sample is False assert result.instance_config.mask_sample is False @pytest.mark.asyncio diff --git a/tests/unit_tests/test_switchyard_trace.py b/tests/unit_tests/test_switchyard_trace.py index 38b16ff71a..0057a879a2 100644 --- a/tests/unit_tests/test_switchyard_trace.py +++ b/tests/unit_tests/test_switchyard_trace.py @@ -157,7 +157,16 @@ def test_assistant_in_suffix_rejected(self) -> None: def test_reordered_records_rejected(self) -> None: records = list(reversed(_two_call_records())) - with pytest.raises(SwitchyardTraceError, match="does not extend"): + with pytest.raises(SwitchyardTraceError, match="not extend"): + _reconstruct(_envelope(records)) + + def test_non_extending_token_history_rejected(self) -> None: + # Messages extend but the token ids don't (e.g. a chat template that + # re-renders earlier history differently) — the trainer would hard-assert + # on this, so reconstruction must mask it instead. + records = _two_call_records() + records[1]["prompt_token_ids"] = [9, 9, 9, 9, 9, 9] + with pytest.raises(SwitchyardTraceError, match="do not extend the prior prompt"): _reconstruct(_envelope(records))