diff --git a/gym_switchyard_integration.md b/gym_switchyard_integration.md new file mode 100644 index 0000000000..df5a5f10df --- /dev/null +++ b/gym_switchyard_integration.md @@ -0,0 +1,481 @@ +# 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 `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 + +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 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. + +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: + +- `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 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 + +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 exactly `http://:` (validated at server startup) and +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 + (`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 +existing `instance_config`, so `SWEBenchWrapper.run` can retrieve the same session +after `responses` returns without introducing shared mutable state. + +## 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 + +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 +`instance_config`. If it contains a Switchyard URL and session UUID, request: + +~~~text +GET /v1/sessions//completions +~~~ + +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 +durably writes each record (atomic rename) 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; +- 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`; +- 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. + +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 the existing `SWEBenchWrapperInstanceConfig.mask_sample` flag to true — this + 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 +- 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` (validated as `http://:` at startup); + - add and generate `switchyard_session_id`; + - 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. `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. `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. + +### nv-OpenHands + +Nothing. The pinned fork commit is unchanged and no pins move. + +### Switchyard + +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: + +- 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, 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 + 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. 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 + +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. 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 new file mode 100644 index 0000000000..6bbe41c917 --- /dev/null +++ b/nemo_gym/switchyard_trace.py @@ -0,0 +1,280 @@ +# 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] = [] + # 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: + 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/app.py b/responses_api_agents/swe_agents/app.py index 316fc13ffa..e52dc9e8eb 100644 --- a/responses_api_agents/swe_agents/app.py +++ b/responses_api_agents/swe_agents/app.py @@ -37,9 +37,11 @@ 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 +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 +64,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 nemo_gym.switchyard_trace import SwitchyardTrace, reconstruct_switchyard_rollout from responses_api_models.vllm_model.app import VLLMConverter, split_responses_input_output_items @@ -177,6 +180,18 @@ 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, 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." + ), + ) + openhands_should_log: bool = False debug: bool = False @@ -255,6 +270,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 +318,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 ######################################## @@ -1497,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. @@ -1592,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 @@ -1714,8 +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_MODEL_SERVER_NAME={self.config.model_server_name} &&" + f"export NEMO_GYM_CONFIG_DICT={ng_config_dict_str} && " + f"export NEMO_GYM_MODEL_SERVER_NAME={self.config.model_server_name} && " "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) @@ -2760,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 @@ -3398,6 +3474,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 +3545,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 +3773,58 @@ 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, + mask_sample=instance_config.mask_sample, **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/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 diff --git a/responses_api_agents/swe_agents/tests/test_app.py b/responses_api_agents/swe_agents/tests/test_app.py index 6f1402f24f..fd39f56b5d 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 @@ -22,14 +23,19 @@ 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 nemo_gym.switchyard_trace import SwitchyardTrace, SwitchyardTraceError from responses_api_agents.swe_agents.app import ( ActiveContainerCommand, AgentPromptOverride, @@ -50,8 +56,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, @@ -171,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 ######################################## @@ -933,6 +950,65 @@ 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_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) + # 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_by_default(self) -> None: + with tempfile.TemporaryDirectory() as 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() + 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)) + ######################################## # Workspace path + user-message resolver tests @@ -2286,6 +2362,72 @@ 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" + 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") + # 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()) + 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) + 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 "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) + 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 + 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 "proxy_x_session_id" not in params.agent_script + class TestSWEBenchWrapperResponses: def _setup_oh_dirs(self, wrapper): @@ -2460,6 +2602,177 @@ 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.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] == [ + 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.mask_sample is True + 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.mask_sample is False + 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/tests/unit_tests/test_switchyard_trace.py b/tests/unit_tests/test_switchyard_trace.py new file mode 100644 index 0000000000..0057a879a2 --- /dev/null +++ b/tests/unit_tests/test_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. +import pytest + +from nemo_gym.openai_utils import ( + NeMoGymEasyInputMessage, + NeMoGymFunctionCallOutput, + NeMoGymResponseFunctionToolCallForTraining, + NeMoGymResponseOutputMessageForTraining, +) +from nemo_gym.responses_converter import ResponsesConverter +from nemo_gym.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="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)) + + +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"})