diff --git a/README.md b/README.md index 69185488bb..0dccba4ed3 100644 --- a/README.md +++ b/README.md @@ -317,8 +317,10 @@ The Dataset column links to publicly available datasets (e.g., on HuggingFace). | Swe Agents | | - | - | ✓ | ✓ | Apache 2.0 | swebench_multi_tools.yaml | - | | Swe Agents | | - | - | ✓ | ✓ | Apache 2.0 | swebench_openhands.yaml | - | | Swe Agents | | - | - | ✓ | ✓ | Apache 2.0 | swebench_openhands_training.yaml | - | -| Swe Agents | coding | SWE-bench driven by the opencode agent framework for RL training. | Train software engineering capabilities on SWE tasks using opencode rollouts. | ✓ | - | Apache 2.0 | swebench_opencode_training.yaml | - | -| Swe Agents | coding | SWE-bench driven by the opencode agent framework. | Eval software engineering capabilities on SWE-bench using opencode. | ✓ | ✓ | Apache 2.0 | swebench_opencode.yaml | - | +| Swe Agents | coding | SWE-bench driven by the opencode agent framework for RL training. | Train software engineering capabilities on SWE tasks using opencode rollouts. | ✓ | - | Apache 2.0 | swebench_nv_opencode_training.yaml | - | +| Swe Agents | coding | SWE-bench driven by the opencode agent framework. | Eval software engineering capabilities on SWE-bench using opencode. | ✓ | ✓ | Apache 2.0 | swebench_nv_opencode.yaml | - | +| Swe Agents | coding | SWE-bench driven by upstream opencode for RL training. | Train software engineering capabilities on SWE tasks using upstream opencode rollouts. | ✓ | - | Apache 2.0 | swebench_opencode_training.yaml | - | +| Swe Agents | coding | SWE-bench driven by upstream opencode. | Eval software engineering capabilities on SWE-bench using upstream opencode. | ✓ | ✓ | Apache 2.0 | swebench_opencode.yaml | - | | Swe Pivot | agent | SWE pivot verifier for PivotRL on coding agent trajectories | Improve coding agent fix-design decisions | ✓ | ✓ | Apache 2.0 | swe_pivot.yaml | - | | Swerl Gen | coding | Running sandboxed evaluation for SWE-style tasks (either patch generation or reproduction test generation) | Improve SWE capabilities useful for benchmarks like SWE-bench | ✓ | ✓ | Apache 2.0 | swerl_gen.yaml | - | | Swerl Llm Judge | coding | SWE-style multiple-choice LLM-judge tasks scored via ... choice. | Improve SWE capabilities useful for benchmarks like SWE-bench | ✓ | ✓ | MIT | swerl_llm_judge.yaml | - | diff --git a/nemo_gym/switchyard_trace.py b/nemo_gym/switchyard_trace.py index 6bbe41c917..ea4d6271c6 100644 --- a/nemo_gym/switchyard_trace.py +++ b/nemo_gym/switchyard_trace.py @@ -26,8 +26,9 @@ masks the sample rather than emitting partially annotated training data. """ +import json import math -from typing import Any, Dict, List, NamedTuple +from typing import Any, Dict, List, NamedTuple, Optional from openai.types.responses.function_tool import FunctionTool @@ -53,6 +54,9 @@ class SwitchyardTrace(NamedTuple): tools: List[FunctionTool] record_uuids: List[str] model: str + # Set when reconstruction stopped at a validation break and emitted only the + # longest valid prefix. Carriers of this MUST be masked from the loss. + partial_reason: Optional[str] = None def map_switchyard_tools(raw_tools: Any) -> List[FunctionTool]: @@ -78,12 +82,22 @@ def map_switchyard_tools(raw_tools: Any) -> List[FunctionTool]: return tools -def reconstruct_switchyard_rollout(envelope: Any, session_id: str, converter: ResponsesConverter) -> SwitchyardTrace: +def reconstruct_switchyard_rollout( + envelope: Any, session_id: str, converter: ResponsesConverter, allow_partial: bool = True +) -> 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. + environment-message suffix. By default (``allow_partial=True`` — the + training pipeline is the primary consumer, and masking a partial sample is + its correct semantics) a deviation at record i > 0 stops the walk and emits + the longest valid prefix with ``partial_reason`` set; pass + ``allow_partial=False`` for fail-fast certification (analysis, tests). The prefix is the only token-continuous object that exists once the + chain breaks (each record's tokens are deltas over the accumulated history), + and emitting it lets the sample mask itself from the loss instead of + aborting its whole prompt group downstream. A record-0 failure still raises: + there is no valid prefix to emit. """ records = _validate_envelope(envelope, session_id) @@ -103,12 +117,12 @@ def reconstruct_switchyard_rollout(envelope: Any, session_id: str, converter: Re # violation (e.g. a template that re-renders history differently) masks the # sample instead of crashing the training step. token_history: List[int] = [] + partial_reason: Optional[str] = None for i, record in enumerate(records): + try: _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" - ) + 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}") @@ -123,7 +137,11 @@ def reconstruct_switchyard_rollout(envelope: Any, session_id: str, converter: Re if i == 0: new_context = prompt else: - if prompt[: len(assembled)] != assembled: + # Compare with tool-call arguments canonicalized so a client that + # re-serializes JSON differently across turns (compact vs. pretty) + # does not fail the extension check. The reconstructed items keep the + # original arguments — this normalization is comparison-only. + if _canonicalize_tool_args(prompt[: len(assembled)]) != _canonicalize_tool_args(assembled): raise SwitchyardTraceError(f"record {i} prompt does not extend the reconstructed history") new_context = prompt[len(assembled) :] if not new_context: @@ -145,6 +163,11 @@ def reconstruct_switchyard_rollout(envelope: Any, session_id: str, converter: Re } ) record_uuids.append(record["uuid"]) + except SwitchyardTraceError as e: + if not allow_partial or i == 0: + raise + partial_reason = f"record {i} of {len(records)}: {e}" + break items = converter.chat_completions_messages_to_responses_items(annotated) input_items, output_items = split_responses_input_output_items(items) @@ -155,6 +178,7 @@ def reconstruct_switchyard_rollout(envelope: Any, session_id: str, converter: Re tools=map_switchyard_tools(tools), record_uuids=record_uuids, model=model, + partial_reason=partial_reason, ) @@ -255,6 +279,34 @@ def _content_to_text(content: Any, i: int) -> str: raise SwitchyardTraceError(f"record {i} contains unsupported content of type {type(content).__name__}") +def _canonicalize_tool_args(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Copy of *messages* with assistant tool-call ``arguments`` canonicalized to + sorted-key JSON, used only for the history-extension equality check. + + A client may re-serialize the same tool-call JSON differently across turns + (compact vs. pretty-printed); that whitespace difference must not fail the + extension check. The reconstructed training items keep the original arguments + — this normalization never reaches them. + """ + out: List[Dict[str, Any]] = [] + for message in messages: + tool_calls = message.get("tool_calls") + if not tool_calls: + out.append(message) + continue + canonical_calls = [] + for call in tool_calls: + function = call["function"] + arguments = function["arguments"] + try: + arguments = json.dumps(json.loads(arguments), sort_keys=True) + except (json.JSONDecodeError, TypeError): + pass # not valid JSON — compare as-is + canonical_calls.append({**call, "function": {**function, "arguments": arguments}}) + out.append({**message, "tool_calls": canonical_calls}) + return out + + 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") diff --git a/responses_api_agents/swe_agents/README.md b/responses_api_agents/swe_agents/README.md index 75d3ec7929..179a38d143 100644 --- a/responses_api_agents/swe_agents/README.md +++ b/responses_api_agents/swe_agents/README.md @@ -409,8 +409,10 @@ Bundled YAML configs: - `configs/swebench_openhands_training.yaml` — same shape as above but tuned for training. - `configs/swe_agent_config.yaml` — alternative SWE-agent path (uses `agent_tools_file`). - `configs/swebench_multi_tools.yaml` — full 15-way prompt × agent-class × tool-name bundle (OpenHands). -- `configs/swebench_opencode.yaml` — opencode harness (`agent_framework: opencode`), pinned fork + commit, no prompt overrides. -- `configs/swebench_opencode_training.yaml` — opencode harness tuned for training. +- `configs/swebench_opencode.yaml` — upstream opencode harness (`opencode_source: opencode`), routed through Switchyard, no prompt overrides. +- `configs/swebench_opencode_training.yaml` — upstream opencode harness tuned for training. +- `configs/swebench_nv_opencode.yaml` — NeMo-Gym opencode fork (`opencode_source: nv-opencode`), pinned repo + commit. Deprecated; prefer the upstream configs above. +- `configs/swebench_nv_opencode_training.yaml` — fork harness tuned for training. Deprecated; prefer the upstream config above. - `configs/swebench_opencode_multi_tools.yaml` — opencode harness with the 11-way `prompts/opencode_harness/*.txt` prompt bundle. See [Prompt overrides](#opencode-integration). - `configs/swebench_opencode_no_instruction.yaml` / `configs/swebench_opencode_empty.yaml` — opencode ablation baselines (minimal / no methodology instructions). - `configs/swebench_deepswe.yaml` / `configs/swebench_denovoswe.yaml` — the `deepswe` / `denovoswe` datasets (⚠️ WIP, see [Supported datasets and harnesses](#supported-datasets-and-harnesses)). diff --git a/responses_api_agents/swe_agents/app.py b/responses_api_agents/swe_agents/app.py index e52dc9e8eb..2ef59437ab 100644 --- a/responses_api_agents/swe_agents/app.py +++ b/responses_api_agents/swe_agents/app.py @@ -24,6 +24,7 @@ import shlex import shutil import signal +import socket import sys import time import uuid @@ -65,7 +66,12 @@ ) 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 nemo_gym.switchyard_trace import SwitchyardTrace, reconstruct_switchyard_rollout +from nemo_gym.switchyard_trace import ( + SWITCHYARD_SCHEMA_VERSION, + SwitchyardTrace, + SwitchyardTraceError, + reconstruct_switchyard_rollout, +) from responses_api_models.vllm_model.app import VLLMConverter, split_responses_input_output_items @@ -107,6 +113,13 @@ class SWEBenchWrapperConfig(BaseResponsesAPIAgentConfig): "fork at swe_openhands_setup/. 'opencode' uses the opencode fork at swe_opencode_setup/ via " "its bench/ entry point.", ) + opencode_source: Literal["opencode", "nv-opencode"] = Field( + default="opencode", + description="For agent_framework='opencode', which opencode to run. 'opencode' (default) runs " + "upstream opencode headless (`opencode run`) with an opencode.json provider pointed at Switchyard " + "— no fork, so opencode's native X-Session-Id is captured. 'nv-opencode' uses the pinned fork's " + "bench entry point.", + ) agent_config: Optional[str] = Field(default=None, description="Path to agent configuration file") agent_tools_file: Optional[str] = Field( default=None, description="Path to JSON file containing tool definitions in OpenAI format (for SWE-agent)" @@ -180,15 +193,50 @@ class SWEBenchWrapperConfig(BaseResponsesAPIAgentConfig): "If False (default), selection is deterministic per instance_id.", ) - switchyard_base_url: Optional[str] = Field( + switchyard_spawn_routing_profile: Optional[str] = Field( + default=None, + description=( + "Path to a Switchyard routing-profiles YAML; the shipped swe_agents/switchyard_profile.yaml " + "works for any deployment (endpoint, model, and parsers resolve from env vars this server " + "exports at spawn). When set, one dedicated token-capture Switchyard " + "instance is spawned per agent run on a free port (stateful token injection requires every " + "call of a session to reach the same process) and torn down after trace retrieval in run(). " + "Records land under the run's persistent_dir, and retrieval reads them from there, so a " + "proxy that exits mid-run does not strand its captured tokens. The `switchyard` CLI must be " + "on PATH. Instances are reaped in run(); calling responses() directly leaks the instance." + ), + ) + switchyard_spawn_host: Optional[str] = Field( + default=None, + description=( + "Host advertised to agent containers for spawned Switchyard instances (the instance binds " + "0.0.0.0). Defaults to this node's hostname, which peer nodes resolve on Slurm clusters. " + "Set explicitly when containers must use a specific routable IP." + ), + ) + switchyard_tool_parser: Optional[str] = Field( + default=None, + description=( + "vLLM tool-call parser for the served model family (e.g. 'qwen3_coder'), exported as " + "${SWITCHYARD_TOOL_PARSER} to the spawned instance. Required by profiles that reference " + "that variable — the shipped switchyard_profile.yaml does." + ), + ) + switchyard_reasoning_parser: Optional[str] = Field( + default=None, + description=( + "vLLM reasoning parser for the served model family (e.g. 'nano_v3'), exported as " + "${SWITCHYARD_REASONING_PARSER} to the spawned instance. Required by profiles that " + "reference that variable — the shipped switchyard_profile.yaml does." + ), + ) + switchyard_parser_pythonpath: 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." + "Colon-separated paths prepended to the spawned instance's PYTHONPATH so Switchyard's " + "parsers can import the serving vLLM. Deployment-specific (e.g. the training container's " + "/opt/nemo-rl/3rdparty/vllm:/opt/nemo_rl_venv/lib/python3.12/site-packages); unset leaves " + "PYTHONPATH untouched." ), ) @@ -266,14 +314,27 @@ class SWEBenchWrapperInstanceConfig(SWEBenchWrapperServerConfig, SWEBenchWrapper agent_command: Optional[ExecuteContainerCommandArgs] = None agent_apptainer_command_str: Optional[str] = None agent_script: Optional[str] = None + # Served model context length, fetched from Switchyard's /v1/models before the + # agent command is built; None leaves the opencode provider entry limit-less. + opencode_context_len: Optional[int] = None # 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. + # agent runs in spawn mode; None leaves behavior unchanged. switchyard_session_id: Optional[str] = None + # URL of the Switchyard spawned for this run, exactly http://:. + # Internal plumbing, not user-settable: _spawn_switchyard_local fills it in on + # the agent's node, and the harness command builder reads it from here. + switchyard_spawned_base_url: Optional[str] = None + + # Generation endpoint this run's Switchyard proxies to. Chosen round-robin by + # the server (which holds the counter) and carried here because the instance + # is spawned on the agent's own node, where that counter is not available. + switchyard_backend_url: Optional[str] = None + @property def instance_id(self) -> str: return self.problem_info["instance_id"] @@ -1518,6 +1579,22 @@ def postprocess_after_run(self, report_file: Path) -> None: report_path.write_text(json.dumps(report, indent=2)) +def _validate_switchyard_config(config: SWEBenchWrapperConfig) -> None: + """Fail fast on a malformed Switchyard routing profile.""" + if config.switchyard_spawn_routing_profile: + if not Path(config.switchyard_spawn_routing_profile).is_file(): + raise ValueError( + f"switchyard_spawn_routing_profile does not exist: {config.switchyard_spawn_routing_profile!r}" + ) + + +def _find_free_port() -> int: + """An OS-assigned free TCP port (small bind-to-use race, acceptable per run).""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("", 0)) + return int(sock.getsockname()[1]) + + def _parse_switchyard_base_url(switchyard_base_url: str) -> Tuple[str, int]: """Host/port of the Switchyard proxy. @@ -1651,7 +1728,7 @@ def get_run_command(self) -> ExecuteContainerCommandArgs: # 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: + if self.config.switchyard_spawned_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, @@ -1659,7 +1736,7 @@ def get_run_command(self) -> ExecuteContainerCommandArgs: 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, + self.config.switchyard_spawned_base_url, ) else: ng_config_dict_str = self.config.ng_global_config_dict_str @@ -1975,10 +2052,158 @@ def _render_opencode_user_message( ) +# Container path for the system-prompt instruction file we hand opencode via the +# config's ``instructions`` list. Kept outside the workspace so it never shows up +# in the ``git diff HEAD`` that becomes the task patch. +_OPENCODE_INSTRUCTIONS_PATH = "/tmp/switchyard_opencode_instructions.md" + +# Steer the model away from opencode's `task` tool crash: upstream opencode +# hard-fails when the model supplies a `task_id` that does not begin with `ses` +# (a synchronous SessionID.make throw defeats its graceful fallback — see +# sst/opencode task tool). Smaller models routinely invent one; omitting it lets +# opencode spawn a fresh subagent session, which Switchyard captures cleanly. +_OPENCODE_TASK_ID_GUIDANCE = ( + "When you use the `task` tool to start a new subagent, do not set the " + "`task_id` parameter — omit it so a fresh subagent session is created. " + "Only pass `task_id` to resume a subagent session you started earlier in " + "this conversation; a valid session id begins with `ses`. Never invent or " + "guess a `task_id` value." +) + + +def _opencode_switchyard_config( + switchyard_base_url: str, model: str, context_len: Optional[int] = None +) -> Dict[str, Any]: + """opencode.json for upstream opencode routed through Switchyard (no fork). + + A custom ``@ai-sdk/openai-compatible`` provider — whose id does NOT start with + ``opencode`` — points at Switchyard. That id prefix is load-bearing: upstream + opencode only emits its native ``X-Session-Id`` / ``x-parent-session-id`` + correlation headers for non-``opencode`` providers, and those headers are how + Switchyard captures the session tree. Set as the default model so opencode does + not fall back to its hosted free models. ``instructions`` appends the task-tool + guidance to opencode's system prompt (see ``_OPENCODE_TASK_ID_GUIDANCE``). + + ``context_len`` (the served model's max_model_len) populates the model's + ``limit`` metadata: opencode budgets its ``max_tokens`` from ``limit.output`` + (falling back to a flat 32000 without it) and enables compaction only when + ``limit.context`` is set. Output is reserved as ``min(32000, context // 4)`` + so compaction (triggered at ``context - output``) keeps most of the window. + """ + model_entry: Dict[str, Any] = {"name": model} + if context_len is not None: + model_entry["limit"] = { + "context": context_len, + "output": min(32000, context_len // 4), + } + return { + "$schema": "https://opencode.ai/config.json", + "provider": { + "switchyard": { + "npm": "@ai-sdk/openai-compatible", + "name": "Switchyard", + "options": {"baseURL": f"{switchyard_base_url.rstrip('/')}/v1", "apiKey": "dummy"}, + "models": {model: model_entry}, + } + }, + "model": f"switchyard/{model}", + "instructions": [_OPENCODE_INSTRUCTIONS_PATH], + "autoupdate": False, + } + + +def _filter_title_gen_records(envelope: dict) -> dict: + """Strip the title-generation record opencode fires before the agent loop. + + opencode calls the model once with system='You are a title generator' to + name each session. Removing it keeps reconstruction to agent turns only. + """ + filtered = [ + r + for r in envelope.get("completions", []) + if not any( + msg.get("role") == "system" and "title generator" in str(msg.get("content", "")).lower() + for msg in r.get("messages", []) + ) + ] + return {**envelope, "completions": filtered} + + + +def _retrieve_switchyard_trace_from_records( + records_by_session: Dict[str, List[Dict[str, Any]]], + session_id: str, + converter: "VLLMConverter", + filter_title_gen: bool = False, + allow_partial: bool = True, +) -> SwitchyardTrace: + """Rebuild one session's rollout from its captured records. + + Module-level and converter-explicit so the Ray task — which owns the capture + directory but has no server instance — can run retrieval on the rollout's own + node instead of centralizing it on the gym server. + """ + envelope = { + "schema_version": SWITCHYARD_SCHEMA_VERSION, + "session_id": session_id, + "completions": records_by_session.get(session_id, []), + } + if filter_title_gen: + envelope = _filter_title_gen_records(envelope) + return reconstruct_switchyard_rollout(envelope, session_id, converter, allow_partial=allow_partial) + + +def _reconstruct_switchyard_sessions_from_records( + records_by_session: Dict[str, List[Dict[str, Any]]], + sessions: List[Dict[str, Any]], + converter: "VLLMConverter", + filter_title_gen: bool = False, + allow_partial: bool = True, +) -> Tuple[SwitchyardTrace, str, List[Dict[str, Any]], Optional[str]]: + """Reconstruct every session in an OpenCode session tree. See the wrapper + method of the same name for full semantics (root selection, degradation on + ambiguous trees); this is the converter-explicit form the Ray task uses. + """ + roots = [se for se in sessions if not se.get("parent_session_id")] + if not roots: + raise SwitchyardTraceError("no root session captured") + + degraded: Optional[str] = None + if len(roots) > 1: + roots.sort(key=lambda se: len(records_by_session.get(str(se["session_id"]), [])), reverse=True) + degraded = f"ambiguous session tree: {len(roots)} root sessions (likely context compaction)" + + root_id = str(roots[0]["session_id"]) + root_trace = _retrieve_switchyard_trace_from_records( + records_by_session, root_id, converter, filter_title_gen=filter_title_gen, allow_partial=allow_partial + ) + + subagents: List[Dict[str, Any]] = [] + for entry in sessions: + session_id = str(entry["session_id"]) + if not entry.get("parent_session_id") or session_id == root_id: + continue + trace = _retrieve_switchyard_trace_from_records( + records_by_session, session_id, converter, filter_title_gen=filter_title_gen + ) + subagents.append( + { + "session_id": session_id, + "parent_session_id": entry["parent_session_id"], + "input": [item.model_dump() for item in trace.input_items], + "output": [item.model_dump() for item in trace.output_items], + "model": trace.model, + "record_uuids": trace.record_uuids, + } + ) + return root_trace, root_id, subagents, degraded + class OpenCodeHarnessProcessor(BaseDatasetHarnessProcessor): """Drives the opencode fork; mirrors OpenHandsHarnessProcessor.""" def setup(self) -> Path: + if self.config.opencode_source == "opencode": + return self._setup_upstream() setup_dir = self.parent_dir / "swe_opencode_setup" with self._setup_directory_lock(setup_dir, "opencode"): @@ -2009,7 +2234,28 @@ def setup(self) -> Path: self._run_setup_command(command) return setup_dir + def _setup_upstream(self) -> Path: + """Install upstream opencode-ai via bun; no fork clone.""" + setup_dir = self.parent_dir / "swe_upstream_opencode_setup" + with file_lock(setup_dir, "upstream opencode setup"): + bun_dir = setup_dir / "bun" + install_dir = setup_dir / "opencode" + opencode_bin = install_dir / "node_modules" / ".bin" / "opencode" + if opencode_bin.exists() and (bun_dir / "bin" / "bun").exists(): + print(f"upstream opencode already installed at {setup_dir}", flush=True) + return setup_dir + print(f"Installing upstream opencode-ai at {setup_dir}...", flush=True) + setup_dir.mkdir(parents=True, exist_ok=True) + install_dir.mkdir(parents=True, exist_ok=True) + # The env assignment must ride the pipeline stage that runs the + # installer: `VAR=x cmd1 | cmd2` binds VAR to cmd1 only. + self._run_setup_command(f"curl -fsSL https://bun.sh/install | BUN_INSTALL={bun_dir} bash") + self._run_setup_command(f"cd {install_dir} && PATH={bun_dir}/bin:$PATH bun add opencode-ai") + return setup_dir + def get_run_command(self) -> ExecuteContainerCommandArgs: + if self.config.opencode_source == "opencode": + return self._get_upstream_run_command() data_point = self.config.problem_info agent_run_id = self.config.agent_run_id @@ -2195,6 +2441,157 @@ def get_run_command(self) -> ExecuteContainerCommandArgs: timeout=self.config.swebench_agent_timeout + 60, ) + def _get_upstream_run_command(self) -> ExecuteContainerCommandArgs: + """Run command for opencode_source='opencode': upstream stock opencode via opencode.json.""" + data_point = self.config.problem_info + agent_run_id = self.config.agent_run_id + instance_id = data_point["instance_id"] + eval_dir_in_opencode = self.config.eval_dir_in_openhands + + assert self.config.opencode_setup_dir is not None, ( + "opencode setup directory is not set; opencode_source='opencode' requires _setup_upstream() to have run." + ) + assert self.config.switchyard_spawned_base_url is not None, ( + "opencode_source='opencode' requires a spawned Switchyard — it captures the session." + ) + + try: + model_server_cfg = get_first_server_config_dict(get_global_config_dict(), self.config.model_server_name) + default_model_name = ( + getattr(model_server_cfg, "openai_model", None) or getattr(model_server_cfg, "model", None) or "" + ) + except Exception as e: + raise RuntimeError( + f"Could not resolve model server '{self.config.model_server_name}' for upstream opencode: {e}" + ) + + effective_model = self.config.body.model or default_model_name + workspace_path = _resolve_opencode_workspace_path(data_point) + user_message = _render_opencode_user_message( + data_point, + workspace_path, + template_override_path=self.config.resolved_user_prompt_template, + ) + user_message_host_path = self.config.persistent_dir / f"user_message_{agent_run_id}.txt" + user_message_host_path.write_text(user_message) + + opencode_cfg_json = json.dumps( + _opencode_switchyard_config( + self.config.switchyard_spawned_base_url, effective_model, self.config.opencode_context_len + ) + ) + + # Output goes to the same eval-dir layout as the fork so _openhands_dir_copy_from_host works unchanged. + output_dir_in_container = f"/opencode_setup/opencode/{eval_dir_in_opencode}/{instance_id}/bench_run" + output_jsonl_in_container = f"{output_dir_in_container}/output.jsonl" + + dataset_name = str(data_point.get("dataset_name", "")) + if "SWE-Gym" in dataset_name: + conda_activate_cmd = ( + "{ deactivate >/dev/null 2>&1 || true; unset VIRTUAL_ENV; " + "if [ -d /opt/miniconda3 ]; then " + ". /opt/miniconda3/etc/profile.d/conda.sh && conda activate testbed || true; " + "fi; } && " + ) + elif "R2E-Gym" in dataset_name: + conda_activate_cmd = ( + "{ deactivate >/dev/null 2>&1 || true; unset VIRTUAL_ENV; " + "if [ -f /testbed/.venv/bin/activate ]; then " + ". /testbed/.venv/bin/activate || true; " + "fi; } && " + ) + elif dataset_name in ("nv-internal-1", "swe-bench-ext") or "SWE-rebench-V2" in dataset_name: + conda_activate_cmd = "" + else: + conda_activate_cmd = ( + "if [ -d /opt/miniconda3 ]; then " + ". /opt/miniconda3/etc/profile.d/conda.sh && conda activate testbed || true; " + "fi && " + ) + + baseline_fix = _extract_instance_dict(data_point).get("baseline_fix", "") + baseline_fix_cmd = f"{{ {baseline_fix} >/tmp/baseline_fix.log 2>&1 || true; }} && " if baseline_fix else "" + + # args dict encoded as JSON so the extraction Python never sees shell metacharacters + py_args_json = json.dumps( + { + "workspace": workspace_path, + "instance_id": instance_id, + "model": effective_model, + "output_dir": output_dir_in_container, + "output_jsonl": output_jsonl_in_container, + } + ) + + agent_main_cmd = ( + "mkdir -p /tmp/ && " + "export PATH=/opencode_setup/bun/bin:$PATH && " + f'date +"%s.%N" > {self.config.generation_apptainer_spinup_timestamp_mounted_fpath} && ' + 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} && " + "export OPENCODE_DISABLE_MODELS_FETCH=1 && " + f"export SWITCHYARD_BASE_URL={shlex.quote(self.config.switchyard_spawned_base_url)} && " + "mkdir -p /root/.cache/opencode && " + "echo '{}' >/root/.cache/opencode/models.json && " + f"{conda_activate_cmd}" + f"{baseline_fix_cmd}" + f"cd {shlex.quote(workspace_path)} && " + f"echo {shlex.quote(opencode_cfg_json)} > opencode.json && " + f"echo {shlex.quote(_OPENCODE_TASK_ID_GUIDANCE)} > {_OPENCODE_INSTRUCTIONS_PATH} && " + "_OC_EXIT=0 && " + f"timeout {self.config.swebench_agent_timeout} " + "/opencode_setup/opencode/node_modules/.bin/opencode run " + f"--model {shlex.quote(f'switchyard/{effective_model}')} " + '"$(cat /opencode_setup/opencode/user_message.txt)" ' + "|| _OC_EXIT=$?" + ) + + # Post-run extraction: git diff → output.jsonl in bench format. + # Uses _OC_ARGS JSON so no shell metacharacter issues in paths/ids. + extract_py = ( + "import json,os,subprocess;" + "a=json.loads(os.environ['_OC_ARGS']);" + "r=subprocess.run(['git','-C',a['workspace'],'diff','HEAD'],capture_output=True,text=True,errors='replace');" + "patch=r.stdout.strip() or None;" + "ec=int(os.environ.get('_OC_EXIT','0'));" + "os.makedirs(a['output_dir'],exist_ok=True);" + "f=open(a['output_jsonl'],'w');" + "json.dump({'instance_id':a['instance_id'],'test_result':{'git_patch':patch}," + "'metadata':{'llm_config':{'model':a['model']}}," + "'error':None if ec==0 else f'opencode exit {ec}','metrics':None},f);" + "f.close()" + ) + + agent_script_name = f"agent_script_{agent_run_id}.sh" + agent_script_path = self.config.persistent_dir / agent_script_name + with open(agent_script_path, "w") as f: + f.write("#!/bin/bash\nset -e\n") + f.write(agent_main_cmd) + f.write(f"\nexport _OC_ARGS={shlex.quote(py_args_json)}\n") + f.write(f"python3 -c {shlex.quote(extract_py)}\n") + f.flush() + os.fsync(f.fileno()) + + agent_timeout_seconds = self.config.swebench_agent_timeout + opencode_cmd = ( + f"timeout --signal=TERM --kill-after=30 {agent_timeout_seconds} " + f"bash /trajectories_mount/{agent_script_name}" + ) + + search_path = os.path.join( + self.config.opencode_setup_dir / "opencode" / eval_dir_in_opencode, + "**", + "output.jsonl", + ) + + return ExecuteContainerCommandArgs( + command=opencode_cmd, + expected_file_pattern=search_path, + mode="agent", + timeout=self.config.swebench_agent_timeout + 60, + ) + ######################################## # START Ray worker logic @@ -2214,6 +2611,217 @@ def _classify_agent_error(err: Optional[str]) -> Optional[str]: return "other" +def _switchyard_log_path(params: SWEBenchWrapperInstanceConfig) -> Path: + """Where a spawned Switchyard's stdout/stderr lands for this run.""" + return params.persistent_dir / "switchyard.log" + + +async def _spawn_switchyard_local(params: SWEBenchWrapperInstanceConfig) -> Tuple[str, Process]: + """Start this run's Switchyard on the local node; returns ``(base_url, process)``. + + Runs wherever its agent runs, so proxies spread across the cluster the same + way ``runner_ray_remote`` does rather than piling onto the server's node. + Records go under the run's ``persistent_dir`` (shared storage), so the server + can rebuild the trajectory from disk after this process is gone. + """ + port = _find_free_port() + host = params.switchyard_spawn_host or socket.gethostname() + rl_log_dir = params.persistent_dir / "switchyard_traces" + rl_log_dir.mkdir(parents=True, exist_ok=True) + log_path = _switchyard_log_path(params) + + env = dict(os.environ) + if params.switchyard_backend_url: + env["SWITCHYARD_VLLM_BASE_URL"] = params.switchyard_backend_url + + # Everything the routing profile references via ${...} must be in the child's + # env: endpoint and model derived from the run's config (same source the gym + # vllm_model server reads), parsers and vLLM import paths from config fields. + cfg = OmegaConf.create(shlex.split(params.ng_global_config_dict_str)[0]) + if not env.get("SWITCHYARD_VLLM_BASE_URL"): + urls = [u for u in (cfg.get("policy_base_url") or []) if u] + if urls: + env["SWITCHYARD_VLLM_BASE_URL"] = str(urls[0]) + if not env.get("SWITCHYARD_POLICY_MODEL"): + model = cfg.get("policy_model_name") + if model: + env["SWITCHYARD_POLICY_MODEL"] = str(model) + if params.switchyard_tool_parser: + env["SWITCHYARD_TOOL_PARSER"] = params.switchyard_tool_parser + if params.switchyard_reasoning_parser: + env["SWITCHYARD_REASONING_PARSER"] = params.switchyard_reasoning_parser + if params.switchyard_parser_pythonpath: + existing = env.get("PYTHONPATH") + env["PYTHONPATH"] = params.switchyard_parser_pythonpath + (f":{existing}" if existing else "") + + # The gym venv's bin dir is on PATH for the server process but not for this Ray + # task, so resolve the CLI next to the interpreter running the task instead. + switchyard_bin = shutil.which("switchyard") or str(Path(sys.executable).parent / "switchyard") + + with log_path.open("w") as log_file: + process = await asyncio.create_subprocess_exec( + switchyard_bin, + "--routing-profiles", + str(params.switchyard_spawn_routing_profile), + "--enable-rl-logging", + "--rl-log-dir", + str(rl_log_dir), + "--", + "serve", + "--port", + str(port), + stdout=log_file, + stderr=log_file, + env=env, + ) + try: + await _wait_switchyard_ready_local(port, process) + except Exception: + await _teardown_switchyard_process(process) + raise + return f"http://{host}:{port}", process + + +async def _wait_switchyard_ready_local(port: int, process: Process, timeout_s: float = 240.0) -> None: + """Wait until *port* accepts TCP connections, or the process dies trying. + + A socket probe rather than an HTTP poll: it needs no client machinery and + puts no connection pressure on a proxy that is still starting up. + + 240s, not 60: under Lustre load the CLI's Python imports alone can stall past + 60s (measured p90=46s with the distribution truncated at the old timeout, + 2026-08-01). A timeout kill also triggers a 16-rollout group retry — an + amplification loop where kills create the spawn herd that causes more stalls. + Dead-on-arrival proxies still fail fast via the process-exit check above. + """ + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if process.returncode is not None: + raise RuntimeError( + f"spawned Switchyard exited with code {process.returncode} before ready " + "(see the run's switchyard.log)" + ) + try: + _, writer = await asyncio.wait_for(asyncio.open_connection("127.0.0.1", port), timeout=1.0) + writer.close() + await writer.wait_closed() + return + except (ConnectionRefusedError, OSError, asyncio.TimeoutError): + pass + await asyncio.sleep(0.2) + raise RuntimeError(f"spawned Switchyard not ready after {timeout_s}s") + + +async def _teardown_switchyard_process(process: Process) -> None: + """Reap a spawned Switchyard. Idempotent.""" + if process.returncode is not None: + return + process.terminate() + try: + await asyncio.wait_for(process.wait(), timeout=10) + except asyncio.TimeoutError: + process.kill() + await process.wait() + + +def _switchyard_spawn_needed_for(params: SWEBenchWrapperInstanceConfig) -> bool: + """Whether this run gets its own Switchyard, decided from the run's own config.""" + if not params.switchyard_spawn_routing_profile or params.verify_golden_patch: + return False + if params.switchyard_session_id: + return True + return params.agent_framework == "opencode" and params.opencode_source == "opencode" + + +def _collect_switchyard_payload(params: SWEBenchWrapperInstanceConfig) -> Optional[Dict[str, Any]]: + """Read and reconstruct this rollout's captured trace, on the rollout's own node. + + Runs in the Ray task strictly AFTER proxy teardown, so the record set is + complete — no live writer exists. Centralizing this on the gym server starved + its event loop at scale (512+ rollouts x hundreds of Lustre reads each). + + Never raises: retrieval failure must not discard the agent/eval outcome. The + error travels in the payload and the server masks the sample (fail closed for + training, open for diagnostics). + """ + if not _switchyard_spawn_needed_for(params): + return None + rl_log_dir = params.persistent_dir / "switchyard_traces" + payload: Dict[str, Any] = { + "trace": None, + "root_id": None, + "subagent_trajectories": None, + "degraded": None, + "error": None, + } + t0 = time.monotonic() + try: + converter = VLLMConverter(return_token_id_information=True) + records_by_session = SWEBenchWrapper._read_switchyard_records(rl_log_dir) + if params.switchyard_session_id: + # OpenHands: Gym minted the session id; single linear chain. + trace = _retrieve_switchyard_trace_from_records( + records_by_session, params.switchyard_session_id, converter, allow_partial=True + ) + payload["trace"] = trace + payload["root_id"] = params.switchyard_session_id + if trace.partial_reason: + payload["degraded"] = f"partial trace: {trace.partial_reason}" + else: + # opencode mints its own ids: discover the session tree from records. + sessions = SWEBenchWrapper._sessions_from_records(records_by_session) + trace, root_id, subagents, degraded = _reconstruct_switchyard_sessions_from_records( + records_by_session, sessions, converter, filter_title_gen=True, allow_partial=True + ) + if trace.partial_reason: + degraded = degraded or f"partial trace: {trace.partial_reason}" + payload.update(trace=trace, root_id=root_id, subagent_trajectories=subagents, degraded=degraded) + except Exception as e: + payload["error"] = f"{type(e).__name__}: {e}" + print( + f"[switchyard-retrieval] root={payload['root_id']} degraded={bool(payload['degraded'])} " + f"error={payload['error']!r} took={time.monotonic() - t0:.1f}s", + flush=True, + ) + return payload + + +async def _run_agent_with_switchyard(params: SWEBenchWrapperInstanceConfig) -> Dict[str, Any]: + """Run one rollout, owning its Switchyard's whole lifecycle on this node. + + The proxy is started here rather than on the server so that its memory and + file descriptors land on the same node as the agent it serves; the server + reads the captured records back from shared storage afterwards. + """ + # A Ray worker is not a child of the Gym server, so it inherits neither the + # server's PATH nor its NEMO_GYM_CONFIG_DICT. Three things below need the latter: + # the spawned proxy's route bundle (the venv's switchyard wrapper derives + # SWITCHYARD_POLICY_MODEL from it), the opencode command builders, and the aiohttp + # client behind _fetch_switchyard_max_model_len -- the last two via + # get_global_config_dict(), which reads this variable directly (global_config.py:793) + # and otherwise falls through to parsing Hydra CLI args. It is stored shell-quoted; + # both consumers want raw YAML. + os.environ.setdefault("NEMO_GYM_CONFIG_DICT", shlex.split(params.ng_global_config_dict_str)[0]) + + process: Optional[Process] = None + if _switchyard_spawn_needed_for(params): + params.switchyard_spawned_base_url, process = await _spawn_switchyard_local(params) + if params.agent_framework == "opencode": + params.opencode_context_len = await SWEBenchWrapper._fetch_switchyard_max_model_len( + params.switchyard_spawned_base_url, params.model_server_name + ) + SWEBenchWrapper._build_agent_command(params) + try: + report_file = await RunOpenHandsAgent(config=params).process_single_datapoint() + finally: + if process is not None: + await _teardown_switchyard_process(process) + return { + "report_file": str(report_file) if report_file else None, + "switchyard": _collect_switchyard_payload(params), + } + + @ray.remote( scheduling_strategy="SPREAD", runtime_env={ @@ -2221,16 +2829,13 @@ def _classify_agent_error(err: Optional[str]) -> Optional[str]: }, num_cpus=0.1, ) -def runner_ray_remote(params_dict: dict[str, Any]) -> Optional[Path]: +def runner_ray_remote(params_dict: dict[str, Any]) -> Dict[str, Any]: # For some reason Ray may not pick up the proper model fields if we don't rebuild the model here. Very strange. SWEBenchWrapperInstanceConfig.model_rebuild(force=True) RunOpenHandsAgent.model_rebuild(force=True) params = SWEBenchWrapperInstanceConfig.model_validate(params_dict) - run_oh = RunOpenHandsAgent(config=params) - report_file = asyncio.run(run_oh.process_single_datapoint()) - - return report_file + return asyncio.run(_run_agent_with_switchyard(params)) def update_and_read_metrics(metrics_fpath: Path, update_dict: Dict[str, Any] | None = None) -> dict: @@ -2823,6 +3428,10 @@ class SWEBenchWrapper(SimpleResponsesAPIAgent): _sem: Optional[Semaphore] = None _vllm_converter: Optional[VLLMConverter] = None + # Switchyard payloads returned by Ray tasks, keyed by agent_run_id; set in + # _inner_responses, popped in run(). Direct responses() callers leak entries + # (same documented caveat as spawn-mode proxies) — growth is logged. + _pending_switchyard: Optional[Dict[str, Any]] = None _swe_bench_wrapper_server_config: Optional[SWEBenchWrapperServerConfig] = None model_config = ConfigDict(arbitrary_types_allowed=True) @@ -2832,9 +3441,11 @@ 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) + # Fail fast on a malformed Switchyard config rather than masking every sample. + _validate_switchyard_config(self.config) + # Round-robin cursor over the policy's generation backends; the spawned + # instances themselves live on their agents' nodes. + self._switchyard_spawn_count: int = 0 run_session_id = f"{int(time.time() * 1000)}_{str(uuid.uuid4())[:8]}" workspace_root = Path(__file__).parent @@ -2861,6 +3472,7 @@ def model_post_init(self, context: Any) -> None: self._sem = Semaphore(self.config.concurrency) self._vllm_converter = VLLMConverter(return_token_id_information=True) + self._pending_switchyard = {} return super().model_post_init(context) @@ -3083,8 +3695,9 @@ def _find_container(self, data_point: dict) -> str: f"Searched in paths: {tried_paths}." ) + @staticmethod def _build_apptainer_command( - self, params: SWEBenchWrapperInstanceConfig, command: ExecuteContainerCommandArgs + params: SWEBenchWrapperInstanceConfig, command: ExecuteContainerCommandArgs ) -> str: # Agent containers only ever see the redacted instance dict. dataset_path_to_mount = str( @@ -3185,20 +3798,32 @@ def _build_apptainer_command( opencode_dir = f"{params.opencode_setup_dir}/opencode" bun_dir = f"{params.opencode_setup_dir}/bun" (Path(opencode_dir) / "evaluation" / "oh").mkdir(parents=True, exist_ok=True) - # opencode reads SQLite migrations from `/../../migration` - # (packages/opencode/src/storage/db.ts) → /opencode_setup/migration. - mount_args.extend( - [ - f"--mount type=bind,src={opencode_dir},dst=/opencode_setup/opencode,ro", - f"--mount type=bind,src={opencode_dir},dst={opencode_dir},ro", - f"--mount type=bind,src={opencode_dir}/evaluation/oh,dst=/opencode_setup/opencode/evaluation/oh", - f"--mount type=bind,src={opencode_dir}/evaluation/oh,dst={opencode_dir}/evaluation/oh", - f"--mount type=bind,src={bun_dir},dst=/opencode_setup/bun,ro", - f"--mount type=bind,src={bun_dir},dst={bun_dir},ro", - f"--mount type=bind,src={dataset_path_to_mount},dst=/root/dataset/data.jsonl", - f"--mount type=bind,src={opencode_dir}/packages/opencode/migration,dst=/opencode_setup/migration,ro", - ] - ) + if params.opencode_source == "opencode": + # Upstream stock opencode: no fork repo structure, no migration mount. + mount_args.extend( + [ + f"--mount type=bind,src={opencode_dir},dst=/opencode_setup/opencode,ro", + f"--mount type=bind,src={opencode_dir}/evaluation/oh,dst=/opencode_setup/opencode/evaluation/oh", + f"--mount type=bind,src={bun_dir},dst=/opencode_setup/bun,ro", + f"--mount type=bind,src={dataset_path_to_mount},dst=/root/dataset/data.jsonl", + ] + ) + else: + # nv-opencode fork: existing mounts (unchanged). + # opencode reads SQLite migrations from `/../../migration` + # (packages/opencode/src/storage/db.ts) → /opencode_setup/migration. + mount_args.extend( + [ + f"--mount type=bind,src={opencode_dir},dst=/opencode_setup/opencode,ro", + f"--mount type=bind,src={opencode_dir},dst={opencode_dir},ro", + f"--mount type=bind,src={opencode_dir}/evaluation/oh,dst=/opencode_setup/opencode/evaluation/oh", + f"--mount type=bind,src={opencode_dir}/evaluation/oh,dst={opencode_dir}/evaluation/oh", + f"--mount type=bind,src={bun_dir},dst=/opencode_setup/bun,ro", + f"--mount type=bind,src={bun_dir},dst={bun_dir},ro", + f"--mount type=bind,src={dataset_path_to_mount},dst=/root/dataset/data.jsonl", + f"--mount type=bind,src={opencode_dir}/packages/opencode/migration,dst=/opencode_setup/migration,ro", + ] + ) user_message_host = params.persistent_dir / f"user_message_{params.agent_run_id}.txt" mount_args.append( f"--mount type=bind,src={user_message_host},dst=/opencode_setup/opencode/user_message.txt,ro" @@ -3479,7 +4104,7 @@ def _setup_params( # retries. Golden-patch verification makes no policy calls, so no session. switchyard_session_id = None if ( - self.config.switchyard_base_url + self.config.switchyard_spawn_routing_profile and self.config.agent_framework == "openhands" and not self.config.verify_golden_patch ): @@ -3600,18 +4225,48 @@ def _setup_params( params.eval_command = dataset_processor.get_run_command() params.eval_apptainer_command_str = self._build_apptainer_command(params, params.eval_command) - if self.config.agent_framework == "opencode": + # Switchyard-routed runs defer the build to responses(): the script + # embeds the Switchyard routing, and spawn mode's instance URL plus the + # served model's context length are only known there. + if not self._agent_command_deferred(params): + self._build_agent_command(params) + + return params, dataset_processor + + def _agent_command_deferred(self, params: SWEBenchWrapperInstanceConfig) -> bool: + """Whether ``responses`` builds the agent command instead of ``_setup_params``. + + Only spawn mode defers: the command embeds the spawned instance's URL, + which does not exist until the Ray task starts the proxy on its own node. + """ + return self._switchyard_spawn_needed(params) + + @staticmethod + def _build_agent_command(params: SWEBenchWrapperInstanceConfig) -> None: + """Build the agent script and command from the current params. + + Called at the end of ``_setup_params``, except for Switchyard-routed + runs, where ``responses`` calls it after the instance URL and the + served model's context length are known. + """ + if params.agent_framework == "opencode": params.agent_command = OpenCodeHarnessProcessor(config=params).get_run_command() else: params.agent_command = OpenHandsHarnessProcessor(config=params).get_run_command() - params.agent_apptainer_command_str = self._build_apptainer_command(params, params.agent_command) + params.agent_apptainer_command_str = SWEBenchWrapper._build_apptainer_command(params, params.agent_command) params.agent_script = params.agent_script_path.read_text() - return params, dataset_processor - async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: params, dataset_processor = self._setup_params(body) + if self._switchyard_spawn_needed(params): + # The instance itself is started by the rollout's own Ray task so it + # lands on the agent's node; only the backend choice is made here, + # where the round-robin counter lives. + params.switchyard_backend_url = self._next_switchyard_backend_url(params) + elif self._agent_command_deferred(params): + self._build_agent_command(params) + with (params.eval_private_dir / "params.json").open("w") as f: f.write(params.model_dump_json(indent=4)) @@ -3629,7 +4284,17 @@ async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body() async def _inner_responses( self, params: SWEBenchWrapperInstanceConfig, dataset_processor: BaseDatasetHarnessProcessor ) -> NeMoGymResponse: - maybe_report_file = await runner_ray_remote.remote(params.model_dump()) + task_result = await runner_ray_remote.remote(params.model_dump()) + maybe_report_file = (task_result or {}).get("report_file") + switchyard_payload = (task_result or {}).get("switchyard") + if switchyard_payload is not None: + self._pending_switchyard[params.agent_run_id] = switchyard_payload + if len(self._pending_switchyard) > 4096: + print( + f"WARNING: {len(self._pending_switchyard)} un-consumed switchyard payloads " + "(direct responses() callers leak them)", + flush=True, + ) metrics_to_update = dict() if maybe_report_file: @@ -3651,7 +4316,9 @@ async def _inner_responses( # 3) Agent itself timed out (wall-clock) — mask regardless of resolved. # 4) Memory watchdog killed the agent container (OOM). # 5) Memory watchdog killed the eval container. - persisted_metrics = SWEBenchMetrics.model_validate(update_and_read_metrics(params.metrics_fpath)) + persisted_metrics = SWEBenchMetrics.model_validate( + await asyncio.to_thread(update_and_read_metrics, params.metrics_fpath) + ) resolved_now = metrics_to_update.get("resolved", False) agent_error_kind = persisted_metrics.agent_error_kind eval_timed_out = bool(persisted_metrics.eval_timed_out) @@ -3726,7 +4393,7 @@ def _item_field(item, name: str): ) input_items, output_items = split_responses_input_output_items(responses_items) - updated_metrics = update_and_read_metrics(params.metrics_fpath, metrics_to_update) + updated_metrics = await asyncio.to_thread(update_and_read_metrics, params.metrics_fpath, metrics_to_update) # body.model can be None (replay JSONLs omit it; the openai_model proxy # picks the backend). NeMoGymResponse.model is a required non-None string, @@ -3775,29 +4442,58 @@ async def run(self, body: BaseRunRequest) -> SWEBenchVerifyResponse: 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: + # Captured records live under the run's own capture directory (the + # ``--rl-log-dir`` handed to its Switchyard), so retrieval reads the + # filesystem rather than the proxy, which may not have survived. + payload = self._pending_switchyard.pop(instance_config.agent_run_id, None) + if instance_config.switchyard_session_id: + # OpenHands: the Ray task retrieved on its own node; consume its payload. + if payload is None or payload.get("error") or payload.get("trace") is None: # 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}" - ) - + reason = (payload or {}).get("error") or "no switchyard payload returned by the rollout task" + switchyard_trace_error = f"session {instance_config.switchyard_session_id}: {reason}" + else: + try: + self._apply_switchyard_trace( + responses_create_params, response, payload["trace"], instance_config.switchyard_session_id + ) + except Exception as e: + instance_config.mask_sample = True + switchyard_trace_error = ( + f"session {instance_config.switchyard_session_id}: {type(e).__name__}: {e}" + ) + elif ( + instance_config.switchyard_spawn_routing_profile + and instance_config.agent_framework == "opencode" + and instance_config.opencode_source == "opencode" + ): + # Upstream opencode: root -> main rollout, subagents -> token-annotated + # subagent_trajectories. All reading/reconstruction already happened in + # the Ray task on the rollout's node; only in-memory application here. + if payload is None or payload.get("error") or payload.get("trace") is None: + instance_config.mask_sample = True + reason = (payload or {}).get("error") or "no switchyard payload returned by the rollout task" + switchyard_trace_error = f"opencode sessions: {reason}" + else: + try: + self._apply_switchyard_trace( + responses_create_params, response, payload["trace"], payload["root_id"] + ) + # Reconstructed (token-annotated) subagents replace the text + # ones on success — even when the reconstructed list is empty. + subagent_trajectories = payload["subagent_trajectories"] + except Exception as e: + instance_config.mask_sample = True + switchyard_trace_error = f"opencode sessions: {type(e).__name__}: {e}" + if not switchyard_trace_error and payload.get("degraded"): + # Real token data was recovered but the tree was ambiguous: + # emit it and mask, so the sample is excluded from the loss + # without aborting its whole prompt group. + instance_config.mask_sample = True + switchyard_trace_error = f"opencode sessions: {payload['degraded']}" return SWEBenchVerifyResponse( responses_create_params=responses_create_params, response=response, @@ -3809,21 +4505,171 @@ async def run(self, body: BaseRunRequest) -> SWEBenchVerifyResponse: 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. + def _switchyard_spawn_needed(self, params: SWEBenchWrapperInstanceConfig) -> bool: + """Whether this run gets its own Switchyard instance. + + Spawn for exactly the runs whose calls are captured: OpenHands runs + with a minted session id, and upstream-opencode runs (which mint their + own ids). Golden-patch verification makes no policy calls. + """ + if not self.config.switchyard_spawn_routing_profile or self.config.verify_golden_patch: + return False + if params.switchyard_session_id: + return True + return self.config.agent_framework == "opencode" and self.config.opencode_source == "opencode" + + def _next_switchyard_backend_url(self, params: SWEBenchWrapperInstanceConfig) -> Optional[str]: + """The generation endpoint this run's Switchyard should proxy to. - 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. + Round-robins across every backend the policy exposes so no single vLLM + server takes the whole rollout wave. Chosen here because the counter is + server-side state; the instance that uses it starts on the agent's node. """ - 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) + cfg = OmegaConf.create(shlex.split(params.ng_global_config_dict_str)[0]) + urls = [u for u in (cfg.get("policy_base_url") or []) if u] + if not urls: + return None + url = urls[self._switchyard_spawn_count % len(urls)] + self._switchyard_spawn_count += 1 + return str(url) + + + @staticmethod + async def _fetch_switchyard_max_model_len(base_url: str, route_id: str) -> Optional[int]: + """The served model's context length from Switchyard's ``/v1/models``, or ``None``. + + Switchyard surfaces the engine's ``max_model_len`` on token-capture + routes. Absence (older Switchyard, engine unreachable at proxy start) + degrades gracefully: the opencode provider entry ships without limits, + exactly today's behavior. + """ + try: + response = await request( + "GET", f"{base_url.rstrip('/')}/v1/models", timeout=ClientTimeout(total=30) + ) + await raise_for_status(response) + entries = (await get_response_json(response)).get("data") or [] + except Exception as e: + print(f"WARNING: could not fetch max_model_len from Switchyard: {e}", flush=True) + return None + for entry in entries: + if isinstance(entry, dict) and entry.get("id") == route_id: + value = entry.get("max_model_len") + if isinstance(value, int) and not isinstance(value, bool): + return value + print(f"WARNING: Switchyard /v1/models has no max_model_len for route {route_id!r}", flush=True) + return None + + @staticmethod + def _read_switchyard_records(rl_log_dir: Path) -> Dict[str, List[Dict[str, Any]]]: + """Every captured record under *rl_log_dir*, grouped by session id. + + Reads the capture directory Switchyard was given via ``--rl-log-dir`` + rather than querying the live proxy. Switchyard writes each record + atomically (tmp file + rename), so a finished rollout's records are + complete on disk whether or not its proxy process survived to serve + them — which it may not at training scale. + + Grouping keys off each record's own ``session_id`` field, so Switchyard's + on-disk directory naming stays its own business. + """ + by_session: Dict[str, List[Dict[str, Any]]] = {} + for record_path in sorted((rl_log_dir / "sessions").glob("*/*.json")): + try: + record = json.loads(record_path.read_text()) + except (OSError, ValueError): + # An unreadable record must not sink the whole rollout; the + # reconstruction below still validates what it does get. + continue + session_id = record.get("session_id") + if isinstance(session_id, str) and session_id: + by_session.setdefault(session_id, []).append(record) + for records in by_session.values(): + records.sort(key=lambda r: (r.get("captured_at") or "", r.get("uuid") or "")) + return by_session + + @staticmethod + def _sessions_from_records(records_by_session: Dict[str, List[Dict[str, Any]]]) -> List[Dict[str, Any]]: + """``[{"session_id", "parent_session_id"}, ...]`` for a run's captured records.""" + return [ + {"session_id": session_id, "parent_session_id": records[0].get("parent_session_id")} + for session_id, records in sorted(records_by_session.items()) + ] + + def _list_switchyard_sessions(self, rl_log_dir: Path) -> List[Dict[str, Any]]: + """Session ids (with parent links) Switchyard captured for this run. + + Harnesses like opencode mint their own session ids, so Gym discovers them + from the capture directory rather than any harness log. Scoped by the + per-run ``--rl-log-dir``, so the list covers exactly this rollout. + """ + return self._sessions_from_records(self._read_switchyard_records(rl_log_dir)) + + def _retrieve_switchyard_trace( + self, + records_by_session: Dict[str, List[Dict[str, Any]]], + session_id: str, + filter_title_gen: bool = False, + ) -> SwitchyardTrace: + """Rebuild one session's rollout from its captured records. + + Harness-neutral: keyed only on ``session_id`` so both the single-session + (OpenHands) and per-subagent (OpenCode) paths reuse it. Records come from + the capture directory rather than the proxy, so retrieval no longer + depends on the Switchyard process still being alive. + """ + return _retrieve_switchyard_trace_from_records( + records_by_session, session_id, self._vllm_converter, filter_title_gen=filter_title_gen + ) + + @staticmethod + def _apply_switchyard_trace( + responses_create_params: dict, + response: NeMoGymResponse, + trace: SwitchyardTrace, + session_id: str, + ) -> None: + """Replace the rollout's input/tools/output with a reconstructed Switchyard trace. + + Switchyard is the source of truth for the token-annotated messages; the + harness supplies only the outcome/reward. + """ + responses_create_params["input"] = [item.model_dump() for item in trace.input_items] + responses_create_params["tools"] = [tool.model_dump() for tool in trace.tools] + response.output = trace.output_items + response.metadata = { + "switchyard_source": "switchyard", + "switchyard_session_id": session_id, + "switchyard_record_uuids": json.dumps(trace.record_uuids), + "switchyard_model": trace.model, + } + + def _reconstruct_switchyard_sessions( + self, + records_by_session: Dict[str, List[Dict[str, Any]]], + sessions: List[Dict[str, Any]], + filter_title_gen: bool = False, + ) -> Tuple[SwitchyardTrace, str, List[Dict[str, Any]], Optional[str]]: + """Reconstruct every session in an OpenCode session tree. + + ``sessions`` is ``[{"session_id", "parent_session_id"}, ...]`` — OpenCode + mints a distinct session id per (sub)agent and Switchyard captures each + separately, so each is an independent linear chain reconstructed by + Layer 1. Returns ``(root_trace, root_session_id, subagent_entries, + degraded_reason)``: the parent-less root becomes the main rollout and each + subagent becomes a token-annotated entry tagged with its parent. + + Degrade rather than fail: when the tree is ambiguous — most commonly + opencode context compaction, which starts a fresh parent-less session + mid-rollout — the largest root is used as the main rollout and + ``degraded_reason`` is returned so the caller can mask the sample. Emitting + real token data plus a mask beats emitting nothing, because a rollout with + no generation data aborts the whole prompt group downstream rather than + just excluding itself from the loss. + """ + return _reconstruct_switchyard_sessions_from_records( + records_by_session, sessions, self._vllm_converter, filter_title_gen=filter_title_gen + ) if __name__ == "__main__": 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 4f833bb4db..ee107fe37f 100644 --- a/responses_api_agents/swe_agents/configs/swebench_multi_tools.yaml +++ b/responses_api_agents/swe_agents/configs/swebench_multi_tools.yaml @@ -15,7 +15,7 @@ swe_agents: 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://: + # switchyard_spawn_routing_profile: /path/to/routing_profiles.yaml # Container configuration container_formatter: ??? diff --git a/responses_api_agents/swe_agents/configs/swebench_nv_opencode.yaml b/responses_api_agents/swe_agents/configs/swebench_nv_opencode.yaml new file mode 100644 index 0000000000..de3eca2dc5 --- /dev/null +++ b/responses_api_agents/swe_agents/configs/swebench_nv_opencode.yaml @@ -0,0 +1,69 @@ +# SWE-bench wrapper configuration for the NeMo-Gym opencode fork (nv-opencode). +swe_agents: + responses_api_agents: + swe_agents: &swe_agents_config + entrypoint: app.py + domain: coding + description: SWE-bench driven by the opencode agent framework. + value: Eval software engineering capabilities on SWE-bench using opencode. + + # Agent framework configuration + agent_framework: opencode + opencode_source: nv-opencode # pinned NeMo-Gym fork (bench entry point) + agent_max_turns: 100 + # Pinned NeMo-Gym opencode fork (adds bench/cli.ts + the nemo-gym LanguageModelV3 provider). + agent_framework_repo: https://github.com/sdevare-nv/nv-opencode.git + agent_framework_commit: sdd/dev + + # Container configuration (same SIFs as the openhands path) + container_formatter: ??? + container_folder_path: null + swebench_agent_timeout: 1800 + swebench_tests_timeout: 900 + apptainer_memory_limit_mb: 32768 + command_exec_timeout: 300 + opencode_subagents_enabled: true + + dataset_path: ??? + + # Optional model server reference + model_server: + name: policy_model + type: responses_api_models + + datasets: + # Training dataset + - name: train + type: train + jsonl_fpath: responses_api_agents/swe_agents/data/swegym_for_sweagent_and_openhands.jsonl + gitlab_identifier: + dataset_name: swegym_for_sweagent_and_openhands + version: 0.0.2 + artifact_fpath: swegym-converted.jsonl + license: Apache 2.0 + # Validation dataset + - name: validation + type: validation + jsonl_fpath: responses_api_agents/swe_agents/data/swebench_verified_for_sweagent_and_openhands.jsonl + gitlab_identifier: + dataset_name: swebench_verified_for_sweagent_and_openhands + version: 0.0.1 + artifact_fpath: swebench_verified_for_sweagent_and_openhands.jsonl + license: TBD + # Example dataset for quick testing + - name: example + type: example + jsonl_fpath: responses_api_agents/swe_agents/data/example.jsonl + +# Alias copies so input rows whose `agent_ref.name` is `swe_agents_val` or +# `swe_agents_train` (baked in by upstream dataset prep, matching the keys +# `swebench_multi_tools.yaml` exposes) resolve against this same config. +swe_agents_val: + responses_api_agents: + swe_agents: + <<: *swe_agents_config + +swe_agents_train: + responses_api_agents: + swe_agents: + <<: *swe_agents_config diff --git a/responses_api_agents/swe_agents/configs/swebench_nv_opencode_training.yaml b/responses_api_agents/swe_agents/configs/swebench_nv_opencode_training.yaml new file mode 100644 index 0000000000..7e86aa1de4 --- /dev/null +++ b/responses_api_agents/swe_agents/configs/swebench_nv_opencode_training.yaml @@ -0,0 +1,57 @@ +# SWE-bench wrapper configuration for the NeMo-Gym opencode fork (nv-opencode), RL training. + +swe_agents_train: + responses_api_agents: + swe_agents: &swe_agents_config + entrypoint: app.py + domain: coding + description: SWE-bench driven by the opencode agent framework for RL training. + value: Train software engineering capabilities on SWE tasks using opencode rollouts. + + # Agent framework configuration + agent_framework: opencode + opencode_source: nv-opencode # pinned NeMo-Gym fork (bench entry point) + agent_max_turns: 100 + # Pinned NeMo-Gym opencode fork (adds bench/cli.ts + the nemo-gym LanguageModelV3 provider). + agent_framework_repo: https://github.com/sdevare-nv/nv-opencode.git + agent_framework_commit: sdd/dev + + # Container configuration (same SIFs as the openhands path) + container_formatter: ??? + container_folder_path: null + swebench_agent_timeout: 1800 + swebench_tests_timeout: 900 + apptainer_memory_limit_mb: 65536 + command_exec_timeout: 300 + + dataset_path: ??? + + agent_prompt_overrides: + - user_prompt_template: responses_api_agents/swe_agents/prompts/opencode_harness/user_prompt.txt + agent_cls: OpenCodeAgent + diversify_tool_names: false + + # Enable opencode's `task` tool (subagent sessions) + opencode_subagents_enabled: true + + model_server: + name: policy_model + type: responses_api_models + + datasets: + - name: train + type: train + jsonl_fpath: responses_api_agents/swe_agents/data/swegym_for_sweagent_and_openhands.jsonl + gitlab_identifier: + dataset_name: swegym_for_sweagent_and_openhands + version: 0.0.2 + artifact_fpath: swegym-converted.jsonl + license: Apache 2.0 + - name: example + type: example + jsonl_fpath: responses_api_agents/swe_agents/data/example.jsonl + +swe_agents_val: + responses_api_agents: + swe_agents: + <<: *swe_agents_config diff --git a/responses_api_agents/swe_agents/configs/swebench_opencode.yaml b/responses_api_agents/swe_agents/configs/swebench_opencode.yaml index c4298fb898..4d20f03ae3 100644 --- a/responses_api_agents/swe_agents/configs/swebench_opencode.yaml +++ b/responses_api_agents/swe_agents/configs/swebench_opencode.yaml @@ -1,18 +1,19 @@ -# SWE-bench wrapper configuration for opencode +# SWE-bench wrapper configuration for upstream opencode. +# Runs `opencode run` headless with an opencode.json provider pointed at Switchyard +# (no fork, no harness modification). Supply `switchyard_spawn_routing_profile` at run +# time — token capture needs it. swe_agents: responses_api_agents: swe_agents: &swe_agents_config entrypoint: app.py domain: coding - description: SWE-bench driven by the opencode agent framework. - value: Eval software engineering capabilities on SWE-bench using opencode. + description: SWE-bench driven by upstream opencode. + value: Eval software engineering capabilities on SWE-bench using upstream opencode. # Agent framework configuration agent_framework: opencode + opencode_source: opencode # upstream opencode-ai (installed via bun); no fork agent_max_turns: 100 - # Pinned NeMo-Gym opencode fork (adds bench/cli.ts + the nemo-gym LanguageModelV3 provider). - agent_framework_repo: https://github.com/sdevare-nv/nv-opencode.git - agent_framework_commit: sdd/dev # Container configuration (same SIFs as the openhands path) container_formatter: ??? diff --git a/responses_api_agents/swe_agents/configs/swebench_opencode_training.yaml b/responses_api_agents/swe_agents/configs/swebench_opencode_training.yaml index d724f97b6a..f803af6a54 100644 --- a/responses_api_agents/swe_agents/configs/swebench_opencode_training.yaml +++ b/responses_api_agents/swe_agents/configs/swebench_opencode_training.yaml @@ -1,19 +1,20 @@ -# SWE-bench wrapper configuration for opencode RL training. +# SWE-bench wrapper configuration for upstream opencode RL training. +# Runs `opencode run` headless with an opencode.json provider pointed at Switchyard +# (no fork, no harness modification). Supply `switchyard_spawn_routing_profile` at run +# time — token capture needs it. swe_agents_train: responses_api_agents: swe_agents: &swe_agents_config entrypoint: app.py domain: coding - description: SWE-bench driven by the opencode agent framework for RL training. - value: Train software engineering capabilities on SWE tasks using opencode rollouts. + description: SWE-bench driven by upstream opencode for RL training. + value: Train software engineering capabilities on SWE tasks using upstream opencode rollouts. # Agent framework configuration agent_framework: opencode + opencode_source: opencode # upstream opencode-ai (installed via bun); no fork agent_max_turns: 100 - # Pinned NeMo-Gym opencode fork (adds bench/cli.ts + the nemo-gym LanguageModelV3 provider). - agent_framework_repo: https://github.com/sdevare-nv/nv-opencode.git - agent_framework_commit: sdd/dev # Container configuration (same SIFs as the openhands path) 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 d0dae762e4..39bb07af7f 100644 --- a/responses_api_agents/swe_agents/configs/swebench_openhands.yaml +++ b/responses_api_agents/swe_agents/configs/swebench_openhands.yaml @@ -12,7 +12,7 @@ swe_agents: 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://: + # switchyard_spawn_routing_profile: /path/to/routing_profiles.yaml # 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 329f242ccb..3e6adf3aef 100644 --- a/responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml +++ b/responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml @@ -11,7 +11,7 @@ swe_agents_train: 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://: + # switchyard_spawn_routing_profile: /path/to/routing_profiles.yaml # Container configuration container_formatter: ??? container_folder_path: null @@ -66,7 +66,7 @@ swe_agents_val: 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://: + # switchyard_spawn_routing_profile: /path/to/routing_profiles.yaml # Container configuration container_formatter: ??? container_folder_path: null diff --git a/responses_api_agents/swe_agents/switchyard_profile.yaml b/responses_api_agents/swe_agents/switchyard_profile.yaml new file mode 100644 index 0000000000..ecc521f4ef --- /dev/null +++ b/responses_api_agents/swe_agents/switchyard_profile.yaml @@ -0,0 +1,34 @@ +# Default routing profile for per-rollout Switchyard spawn (`switchyard_spawn_routing_profile`). +# +# Deployment-independent: every env reference below is exported by +# _spawn_switchyard_local before the instance starts — endpoint and model from the +# run's config, parsers from the switchyard_tool_parser / switchyard_reasoning_parser +# config fields. A referenced variable that is unset fails the profile load with a +# named error. +# +# `default` and `policy_model` are aliases of one route: dispatch is an exact match on +# the request's model id and the harness paths send different ids (upstream opencode +# POSTs "default"; the OpenHands fork POSTs the model server's name, "policy_model" +# in training configs). +defaults: + api_key: dummy + base_url: ${SWITCHYARD_VLLM_BASE_URL} + format: openai + +routes: + default: &policy_route + type: model + target: ${SWITCHYARD_POLICY_MODEL} + token_capture_engine: vllm + token_injection: true + injection_transport: prefix_chat + # NeMo-RL asserts exact on-policy sampling params and upstream opencode omits + # them, so pin them on every upstream request. Numeric literals on purpose: + # env interpolation would turn them into strings. Recipes with different + # sampling params supply their own profile. + extra_body: + temperature: 1.0 + top_p: 1.0 + tool_parser: ${SWITCHYARD_TOOL_PARSER} + reasoning_parser: ${SWITCHYARD_REASONING_PARSER} + policy_model: *policy_route diff --git a/responses_api_agents/swe_agents/tests/test_app.py b/responses_api_agents/swe_agents/tests/test_app.py index fd39f56b5d..a8e7830119 100644 --- a/responses_api_agents/swe_agents/tests/test_app.py +++ b/responses_api_agents/swe_agents/tests/test_app.py @@ -56,6 +56,7 @@ SWEBenchWrapperServerConfig, SWERebenchDatasetProcessor, _extract_instance_dict, + _filter_title_gen_records, _parse_switchyard_base_url, _render_opencode_user_message, _resolve_opencode_workspace_path, @@ -954,7 +955,7 @@ 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_spawned_base_url="http://switchyard:4000", switchyard_session_id="0123456789abcdef0123456789abcdef", ng_global_config_dict_str=_ng_config_dict_str(), ) @@ -1001,15 +1002,6 @@ def test_config_dict_rewrite_points_model_server_at_switchyard(self) -> None: 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 ######################################## @@ -1105,6 +1097,7 @@ def _opencode_config(self, tmpdir, **overrides) -> SWEBenchWrapperInstanceConfig return _make_instance_config( tmpdir, agent_framework="opencode", + opencode_source="nv-opencode", # fork run-command path (default is now upstream 'opencode') opencode_setup_dir=opencode_setup_dir, agent_framework_repo="https://example.invalid/opencode.git", agent_framework_commit="deadbeef", @@ -2384,7 +2377,7 @@ def test_setup_params_switchyard_session(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.config.switchyard_base_url = "http://switchyard:4000" + wrapper.config.switchyard_spawn_routing_profile = "/tmp/profile.yaml" wrapper._swe_bench_wrapper_server_config.ng_global_config_dict_str = _ng_config_dict_str() params, _ = wrapper._setup_params(self._switchyard_body()) @@ -2393,7 +2386,10 @@ def test_setup_params_switchyard_session(self, monkeypatch) -> 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. + # NEMO_GYM_CONFIG_DICT — no OpenHands-side changes involved. The URL is + # only known once the Ray task spawns the proxy, so the script is rebuilt then. + params.switchyard_spawned_base_url = "http://switchyard:4000" + wrapper._build_agent_command(params) assert f'proxy_x_session_id = "{session_id}"' in params.agent_script assert "host: switchyard" in params.agent_script @@ -2420,7 +2416,7 @@ def test_setup_params_no_switchyard_for_golden_patch(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.config.switchyard_base_url = "http://switchyard:4000" + wrapper.config.switchyard_spawn_routing_profile = "/tmp/profile.yaml" wrapper.config.verify_golden_patch = True wrapper._swe_bench_wrapper_server_config.ng_global_config_dict_str = _ng_config_dict_str() @@ -2428,6 +2424,56 @@ def test_setup_params_no_switchyard_for_golden_patch(self, monkeypatch) -> None: assert params.switchyard_session_id is None assert "proxy_x_session_id" not in params.agent_script + def test_validate_switchyard_config_spawn_conflicts(self, tmp_path: Path) -> None: + profile = tmp_path / "profile.yaml" + profile.write_text("routes: {}\n") + config = _minimal_server_config() + + config.switchyard_spawn_routing_profile = str(profile) + swe_app._validate_switchyard_config(config) # valid: spawn only + + config.switchyard_spawn_routing_profile = str(tmp_path / "missing.yaml") + with pytest.raises(ValueError, match="does not exist"): + swe_app._validate_switchyard_config(config) + + def test_setup_params_spawn_mode_mints_session_and_rebuild(self, monkeypatch, tmp_path: Path) -> 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) + profile = tmp_path / "profile.yaml" + profile.write_text("routes: {}\n") + wrapper.config.switchyard_spawn_routing_profile = str(profile) + wrapper._swe_bench_wrapper_server_config.ng_global_config_dict_str = _ng_config_dict_str() + + params, _ = wrapper._setup_params(self._switchyard_body()) + # A session is minted even though the instance URL is not known yet. + assert params.switchyard_session_id is not None + assert params.switchyard_spawned_base_url is None + assert wrapper._switchyard_spawn_needed(params) + + # Spawn mode assigns the URL after setup and rebuilds the script. + params.switchyard_spawned_base_url = "http://spawnhost:12345" + wrapper._build_agent_command(params) + assert "host: spawnhost" in params.agent_script + assert f'proxy_x_session_id = "{params.switchyard_session_id}"' in params.agent_script + + def test_spawn_not_needed_for_golden_patch(self, monkeypatch, tmp_path: Path) -> 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) + profile = tmp_path / "profile.yaml" + profile.write_text("routes: {}\n") + wrapper.config.switchyard_spawn_routing_profile = str(profile) + 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 not wrapper._switchyard_spawn_needed(params) + class TestSWEBenchWrapperResponses: def _setup_oh_dirs(self, wrapper): @@ -2636,7 +2682,7 @@ def _switchyard_response(self, **instance_config_overrides) -> NeMoGymResponse: "metrics": json.dumps({"resolved": True, "patch_exists": True}), "instance_config": _make_instance_config( tempfile.mkdtemp(), - switchyard_base_url="http://switchyard:4000", + switchyard_spawned_base_url="http://switchyard:4000", switchyard_session_id="0123456789abcdef0123456789abcdef", **instance_config_overrides, ).model_dump_json(), @@ -2663,14 +2709,22 @@ async def test_run_switchyard_success_replaces_rollout(self, monkeypatch) -> Non model="Qwen/Qwen3-0.6B", ) + # run() consumes the payload the Ray task returned via _inner_responses. + wrapper._pending_switchyard["test_run_123"] = { + "trace": trace, + "root_id": "0123456789abcdef0123456789abcdef", + "subagent_trajectories": None, + "degraded": None, + "error": None, + } 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 wrapper._pending_switchyard == {} # popped exactly once assert result.switchyard_trace_error is None assert result.mask_sample is False assert result.instance_config.mask_sample is False @@ -2691,19 +2745,22 @@ async def test_run_switchyard_success_replaces_rollout(self, monkeypatch) -> Non async def test_run_switchyard_failure_masks_sample(self, monkeypatch) -> None: wrapper = _create_wrapper(monkeypatch) + wrapper._pending_switchyard["test_run_123"] = { + "trace": None, + "root_id": None, + "subagent_trajectories": None, + "degraded": None, + "error": "SwitchyardTraceError: record 1 prompt does not extend the reconstructed history", + } 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()) + assert wrapper._pending_switchyard == {} + # Fail closed for training, open for diagnostics. assert result.mask_sample is True assert result.instance_config.mask_sample is True @@ -2738,7 +2795,7 @@ async def test_run_without_switchyard_does_not_retrieve(self, monkeypatch) -> No with ( patch.object(SWEBenchWrapper, "responses", new_callable=AsyncMock, return_value=mock_response), - patch.object(SWEBenchWrapper, "_retrieve_switchyard_trace", new_callable=AsyncMock) as retrieve, + patch.object(SWEBenchWrapper, "_retrieve_switchyard_trace") as retrieve, ): result = await wrapper.run(self._run_body()) @@ -2747,31 +2804,211 @@ async def test_run_without_switchyard_does_not_retrieve(self, monkeypatch) -> No 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: + @staticmethod + def _write_record(rl_log_dir: Path, session_id: str, parent_id, uuid: str, **extra) -> None: + """Write one capture record the way Switchyard lays them out on disk.""" + session_dir = rl_log_dir / "sessions" / f"dir_{session_id}" + session_dir.mkdir(parents=True, exist_ok=True) + record = { + "schema_version": 1, + "session_id": session_id, + "parent_session_id": parent_id, + "uuid": uuid, + "captured_at": uuid, + **extra, + } + (session_dir / f"{uuid}.json").write_text(json.dumps(record)) + + def test_retrieve_switchyard_trace_builds_envelope_from_records(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + records = {"ses_1": [{"schema_version": 1, "session_id": "ses_1", "uuid": "u1"}]} + reconstruct_mock = MagicMock(return_value="trace") + monkeypatch.setattr(swe_app, "reconstruct_switchyard_rollout", reconstruct_mock) + + result = wrapper._retrieve_switchyard_trace(records, "ses_1") + + assert result == "trace" + reconstruct_mock.assert_called_once_with( + {"schema_version": 1, "session_id": "ses_1", "completions": records["ses_1"]}, + "ses_1", + wrapper._vllm_converter, + allow_partial=True, + ) + + def test_read_switchyard_records_groups_and_orders(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", + rl_log_dir = Path(tmpdir) / "switchyard_traces" + self._write_record(rl_log_dir, "root", None, "b") + self._write_record(rl_log_dir, "root", None, "a") + self._write_record(rl_log_dir, "sub", "root", "c") + + records = wrapper._read_switchyard_records(rl_log_dir) + + assert sorted(records) == ["root", "sub"] + # Ordered by (captured_at, uuid) so the chain reconstructs in call order. + assert [r["uuid"] for r in records["root"]] == ["a", "b"] + + def test_read_switchyard_records_skips_unreadable(self, monkeypatch) -> None: + """A torn record must not sink a rollout whose other records are intact.""" + wrapper = _create_wrapper(monkeypatch) + with tempfile.TemporaryDirectory() as tmpdir: + rl_log_dir = Path(tmpdir) / "switchyard_traces" + self._write_record(rl_log_dir, "root", None, "a") + (rl_log_dir / "sessions" / "dir_root" / "torn.json").write_text("{ not json") + + records = wrapper._read_switchyard_records(rl_log_dir) + + assert [r["uuid"] for r in records["root"]] == ["a"] + + def test_list_switchyard_sessions_reads_capture_dir(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + with tempfile.TemporaryDirectory() as tmpdir: + rl_log_dir = Path(tmpdir) / "switchyard_traces" + self._write_record(rl_log_dir, "root", None, "u1") + self._write_record(rl_log_dir, "sub", "root", "u2") + + result = wrapper._list_switchyard_sessions(rl_log_dir) + + assert result == [ + {"session_id": "root", "parent_session_id": None}, + {"session_id": "sub", "parent_session_id": "root"}, + ] + + @staticmethod + def _fake_trace(model: str) -> SwitchyardTrace: + item = MagicMock() + item.model_dump.return_value = {"model": model} + return SwitchyardTrace(input_items=[item], output_items=[item], tools=[], record_uuids=["u"], model=model) + + def test_reconstruct_switchyard_sessions_tree(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + # trace.model echoes the session id so we can assert per-session retrieval. + monkeypatch.setattr( + swe_app, + "_retrieve_switchyard_trace_from_records", + MagicMock(side_effect=lambda recs, sid, conv, **kw: self._fake_trace(sid)), + ) + sessions = [ + {"session_id": "root", "parent_session_id": None}, + {"session_id": "sub1", "parent_session_id": "root"}, + {"session_id": "sub2", "parent_session_id": "root"}, + ] + root_trace, root_id, subs, degraded = wrapper._reconstruct_switchyard_sessions({}, sessions) + + assert root_id == "root" and root_trace.model == "root" + assert degraded is None + assert [s["session_id"] for s in subs] == ["sub1", "sub2"] + assert all(s["parent_session_id"] == "root" for s in subs) + assert subs[0]["model"] == "sub1" and subs[0]["output"] == [{"model": "sub1"}] + + def test_reconstruct_switchyard_sessions_single_root(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + monkeypatch.setattr( + swe_app, + "_retrieve_switchyard_trace_from_records", + MagicMock(side_effect=lambda recs, sid, conv, **kw: self._fake_trace(sid)), + ) + root_trace, root_id, subs, degraded = wrapper._reconstruct_switchyard_sessions( + {}, [{"session_id": "only", "parent_session_id": None}] + ) + assert root_id == "only" and subs == [] and degraded is None + + def test_reconstruct_switchyard_sessions_degrades_on_multiple_roots(self, monkeypatch) -> None: + """Compaction starts a second parent-less session mid-rollout. + + Emit the larger chain and report it as degraded rather than raising: a + rollout with no generation data aborts its whole prompt group downstream, + while a masked one only removes itself from the loss. + """ + wrapper = _create_wrapper(monkeypatch) + monkeypatch.setattr( + swe_app, + "_retrieve_switchyard_trace_from_records", + MagicMock(side_effect=lambda recs, sid, conv, **kw: self._fake_trace(sid)), + ) + records = {"small": [{"uuid": "1"}], "big": [{"uuid": "1"}, {"uuid": "2"}]} + sessions = [ + {"session_id": "small", "parent_session_id": None}, + {"session_id": "big", "parent_session_id": None}, + ] + + root_trace, root_id, subs, degraded = wrapper._reconstruct_switchyard_sessions(records, sessions) + + assert root_id == "big" and root_trace.model == "big" + assert degraded and "2 root sessions" in degraded + + def test_reconstruct_switchyard_sessions_requires_a_root(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + monkeypatch.setattr(wrapper, "_retrieve_switchyard_trace", MagicMock()) + with pytest.raises(SwitchyardTraceError): + wrapper._reconstruct_switchyard_sessions( + {}, + [ + {"session_id": "a", "parent_session_id": "b"}, + ], ) - request_mock = AsyncMock(return_value=MagicMock()) - envelope = {"schema_version": 1} - monkeypatch.setattr(swe_app, "request", request_mock) + +def test_opencode_switchyard_config_points_at_switchyard() -> None: + cfg = swe_app._opencode_switchyard_config("http://switchyard:4000/", "policy-model") + assert cfg["model"] == "switchyard/policy-model" + provider_id = next(iter(cfg["provider"])) + # Load-bearing: a non-"opencode" provider id is what makes upstream opencode + # emit its native X-Session-Id correlation headers. + assert not provider_id.startswith("opencode") + provider = cfg["provider"][provider_id] + assert provider["npm"] == "@ai-sdk/openai-compatible" + assert provider["options"]["baseURL"] == "http://switchyard:4000/v1" + assert "policy-model" in provider["models"] + # Task-tool guidance is appended to opencode's system prompt via instructions. + assert cfg["instructions"] == [swe_app._OPENCODE_INSTRUCTIONS_PATH] + # Without a known context length, the model entry ships no limit metadata. + assert "limit" not in provider["models"]["policy-model"] + + +def test_opencode_switchyard_config_sets_model_limits() -> None: + cfg = swe_app._opencode_switchyard_config( + "http://switchyard:4000", "policy-model", context_len=32768 + ) + entry = cfg["provider"]["switchyard"]["models"]["policy-model"] + # context from the engine; output reserved as min(32000, context // 4) so + # opencode's compaction threshold (context - output) keeps most of the window. + assert entry["limit"] == {"context": 32768, "output": 8192} + + large = swe_app._opencode_switchyard_config( + "http://switchyard:4000", "policy-model", context_len=131072 + ) + assert large["provider"]["switchyard"]["models"]["policy-model"]["limit"]["output"] == 32000 + + +class TestFetchSwitchyardMaxModelLen: + def _wrapper(self, monkeypatch) -> SWEBenchWrapper: + return _create_wrapper(monkeypatch) + + @pytest.mark.asyncio + async def test_reads_route_entry(self, monkeypatch) -> None: + wrapper = self._wrapper(monkeypatch) + payload = {"data": [{"id": "other", "max_model_len": 1}, {"id": "policy-model", "max_model_len": 32768}]} + response = MagicMock(status=200) + monkeypatch.setattr(swe_app, "request", AsyncMock(return_value=response)) 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) + monkeypatch.setattr(swe_app, "get_response_json", AsyncMock(return_value=payload)) + assert await wrapper._fetch_switchyard_max_model_len("http://sy:4000", "policy-model") == 32768 - result = await wrapper._retrieve_switchyard_trace(instance_config) + @pytest.mark.asyncio + async def test_missing_route_returns_none(self, monkeypatch) -> None: + wrapper = self._wrapper(monkeypatch) + monkeypatch.setattr(swe_app, "request", AsyncMock(return_value=MagicMock(status=200))) + monkeypatch.setattr(swe_app, "raise_for_status", AsyncMock()) + monkeypatch.setattr(swe_app, "get_response_json", AsyncMock(return_value={"data": []})) + assert await wrapper._fetch_switchyard_max_model_len("http://sy:4000", "policy-model") is None - 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) + @pytest.mark.asyncio + async def test_fetch_failure_returns_none(self, monkeypatch) -> None: + wrapper = self._wrapper(monkeypatch) + monkeypatch.setattr(swe_app, "request", AsyncMock(side_effect=ConnectionError)) + assert await wrapper._fetch_switchyard_max_model_len("http://sy:4000", "policy-model") is None ######################################## @@ -2804,3 +3041,927 @@ def test_loads_from_lib_agent_dir(self) -> None: mod = _load_rebench_log_parsers(rebench_dir) assert "lib_test" in mod.NAME_TO_PARSER + + +######################################## +# _filter_title_gen_records tests +######################################## + + +class TestFilterTitleGenRecords: + def _record(self, system_content: str = "", user_content: str = "fix it") -> dict: + # Matches the Switchyard token_capture_response_processor record shape: + # messages = request_msgs + assistant_turn (flat list, no "request" wrapper) + msgs = [] + if system_content: + msgs.append({"role": "system", "content": system_content}) + msgs.append({"role": "user", "content": user_content}) + msgs.append({"role": "assistant", "content": "ok"}) + return {"messages": msgs} + + def test_strips_title_gen_record(self) -> None: + envelope = { + "completions": [ + self._record("You are a title generator"), + self._record("You are a coding assistant"), + ] + } + result = _filter_title_gen_records(envelope) + assert len(result["completions"]) == 1 + assert result["completions"][0]["messages"][0]["content"] == "You are a coding assistant" + + def test_case_insensitive(self) -> None: + envelope = {"completions": [self._record("YOU ARE A TITLE GENERATOR")]} + assert _filter_title_gen_records(envelope)["completions"] == [] + + def test_preserves_non_title_gen_records(self) -> None: + envelope = {"completions": [self._record("You are a helpful assistant"), self._record("", "task")]} + assert len(_filter_title_gen_records(envelope)["completions"]) == 2 + + def test_empty_completions(self) -> None: + assert _filter_title_gen_records({"completions": []}) == {"completions": []} + + def test_missing_completions_key(self) -> None: + assert _filter_title_gen_records({})["completions"] == [] + + def test_other_envelope_keys_preserved(self) -> None: + envelope = {"schema_version": 1, "completions": [self._record("You are a title generator")]} + result = _filter_title_gen_records(envelope) + assert result["schema_version"] == 1 + assert result["completions"] == [] + + +######################################## +# Upstream opencode setup + run-command tests +######################################## + + +class TestUpstreamOpenCodeSetup: + def _upstream_config(self, tmpdir, **overrides) -> SWEBenchWrapperInstanceConfig: + opencode_setup_dir = Path(tmpdir) / "upstream_opencode_setup" + opencode_setup_dir.mkdir(parents=True, exist_ok=True) + return _make_instance_config( + tmpdir, + agent_framework="opencode", + opencode_source="opencode", + opencode_setup_dir=opencode_setup_dir, + switchyard_spawned_base_url="http://switchyard:4000", + **overrides, + ) + + def test_setup_upstream_skips_if_already_installed(self, monkeypatch) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = self._upstream_config(tmpdir) + setup_dir = Path(tmpdir) / "parent" / "swe_upstream_opencode_setup" + bun_bin = setup_dir / "bun" / "bin" / "bun" + opencode_bin = setup_dir / "opencode" / "node_modules" / ".bin" / "opencode" + bun_bin.parent.mkdir(parents=True, exist_ok=True) + opencode_bin.parent.mkdir(parents=True, exist_ok=True) + bun_bin.touch() + opencode_bin.touch() + + run_cmd = MagicMock() + monkeypatch.setattr(BaseDatasetHarnessProcessor, "_run_setup_command", run_cmd) + + with patch.object( + BaseDatasetHarnessProcessor, + "parent_dir", + new_callable=lambda: property(lambda self: Path(tmpdir) / "parent"), + ): + processor = OpenCodeHarnessProcessor(config=config) + result = processor._setup_upstream() + + run_cmd.assert_not_called() + assert result == Path(tmpdir) / "parent" / "swe_upstream_opencode_setup" + + def test_setup_upstream_runs_install_commands(self, monkeypatch) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = self._upstream_config(tmpdir) + run_cmd = MagicMock() + monkeypatch.setattr(BaseDatasetHarnessProcessor, "_run_setup_command", run_cmd) + + with patch.object( + BaseDatasetHarnessProcessor, + "parent_dir", + new_callable=lambda: property(lambda self: Path(tmpdir) / "parent"), + ): + (Path(tmpdir) / "parent" / "swe_upstream_opencode_setup" / "opencode").mkdir( + parents=True, exist_ok=True + ) + processor = OpenCodeHarnessProcessor(config=config) + processor._setup_upstream() + + assert run_cmd.call_count == 2 + calls = [c.args[0] for c in run_cmd.call_args_list] + assert any("bun.sh/install" in c for c in calls) + assert any("bun add opencode-ai" in c for c in calls) + + def test_setup_dispatches_to_upstream_for_opencode_source(self, monkeypatch) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = self._upstream_config(tmpdir) + upstream_mock = MagicMock(return_value=Path(tmpdir) / "result") + monkeypatch.setattr(OpenCodeHarnessProcessor, "_setup_upstream", upstream_mock) + result = OpenCodeHarnessProcessor(config=config).setup() + upstream_mock.assert_called_once() + assert result == Path(tmpdir) / "result" + + +class TestUpstreamOpenCodeRunCommand: + @pytest.fixture + def _stub_model_server(self, monkeypatch): + def _fake(_global, name): + return type("Cfg", (), {"host": "sw-host", "port": 4000, "model": "policy-model"})() + + monkeypatch.setattr(swe_app, "get_first_server_config_dict", _fake) + monkeypatch.setattr(swe_app, "get_global_config_dict", MagicMock(return_value={})) + + def _upstream_config(self, tmpdir, **overrides) -> SWEBenchWrapperInstanceConfig: + opencode_setup_dir = Path(tmpdir) / "upstream_opencode_setup" + opencode_setup_dir.mkdir(parents=True, exist_ok=True) + return _make_instance_config( + tmpdir, + agent_framework="opencode", + opencode_source="opencode", + opencode_setup_dir=opencode_setup_dir, + switchyard_spawned_base_url="http://switchyard:4000", + ng_global_config_dict_str=_ng_config_dict_str(), + **overrides, + ) + + def _read_agent_script(self, config) -> str: + return (config.persistent_dir / f"agent_script_{config.agent_run_id}.sh").read_text() + + def test_get_run_command_dispatches_for_opencode_source(self, _stub_model_server) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = self._upstream_config(tmpdir) + config.persistent_dir.mkdir(parents=True, exist_ok=True) + upstream_mock = MagicMock( + return_value=ExecuteContainerCommandArgs( + command="echo ok", expected_file_pattern="/tmp/*.jsonl", mode="agent", timeout=100 + ) + ) + with patch.object(OpenCodeHarnessProcessor, "_get_upstream_run_command", upstream_mock): + result = OpenCodeHarnessProcessor(config=config).get_run_command() + upstream_mock.assert_called_once() + assert result.command == "echo ok" + + def test_upstream_run_command_basic(self, _stub_model_server) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = self._upstream_config(tmpdir) + config.persistent_dir.mkdir(parents=True, exist_ok=True) + result = OpenCodeHarnessProcessor(config=config)._get_upstream_run_command() + assert isinstance(result, ExecuteContainerCommandArgs) + assert result.mode == "agent" + assert "timeout" in result.command + assert "output.jsonl" in result.expected_file_pattern + assert str(config.opencode_setup_dir) in result.expected_file_pattern + + def test_upstream_agent_script_contents(self, _stub_model_server) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = self._upstream_config(tmpdir) + config.persistent_dir.mkdir(parents=True, exist_ok=True) + OpenCodeHarnessProcessor(config=config)._get_upstream_run_command() + script = self._read_agent_script(config) + assert "OPENCODE_DISABLE_MODELS_FETCH=1" in script + assert "SWITCHYARD_BASE_URL" in script + assert "switchyard:4000" in script + assert "opencode.json" in script + assert "opencode run" in script + assert "--model" in script + assert "switchyard/test-model" in script # body.model takes priority over server default + assert "_OC_ARGS" in script + assert "python3 -c" in script + assert "git_patch" in script + + def test_upstream_writes_user_message_file(self, _stub_model_server) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = self._upstream_config(tmpdir) + config.persistent_dir.mkdir(parents=True, exist_ok=True) + OpenCodeHarnessProcessor(config=config)._get_upstream_run_command() + user_msg_path = config.persistent_dir / f"user_message_{config.agent_run_id}.txt" + assert user_msg_path.exists() + assert "Fix bug" in user_msg_path.read_text() + + def test_upstream_opencode_json_is_valid(self, _stub_model_server) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = self._upstream_config(tmpdir) + config.persistent_dir.mkdir(parents=True, exist_ok=True) + OpenCodeHarnessProcessor(config=config)._get_upstream_run_command() + script = self._read_agent_script(config) + # Verify opencode.json is written with key Switchyard provider fields. + assert "opencode.json" in script + assert '"model": "switchyard/test-model"' in script + assert '"autoupdate": false' in script + assert '"baseURL": "http://switchyard:4000/v1"' in script + assert '"npm": "@ai-sdk/openai-compatible"' in script + # The task-tool guidance file is written and referenced as an instruction. + assert swe_app._OPENCODE_INSTRUCTIONS_PATH in script + assert "do not set the `task_id` parameter" in script + + def test_upstream_requires_spawned_switchyard(self, _stub_model_server) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = _make_instance_config( + tmpdir, + agent_framework="opencode", + opencode_source="opencode", + opencode_setup_dir=Path(tmpdir) / "setup", + switchyard_spawned_base_url=None, + ng_global_config_dict_str=_ng_config_dict_str(), + ) + config.persistent_dir.mkdir(parents=True, exist_ok=True) + (Path(tmpdir) / "setup").mkdir(parents=True, exist_ok=True) + with pytest.raises(AssertionError, match="requires a spawned Switchyard"): + OpenCodeHarnessProcessor(config=config)._get_upstream_run_command() + + def test_upstream_requires_opencode_setup_dir(self, _stub_model_server) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + config = _make_instance_config( + tmpdir, + agent_framework="opencode", + opencode_source="opencode", + opencode_setup_dir=None, + switchyard_spawned_base_url="http://switchyard:4000", + ng_global_config_dict_str=_ng_config_dict_str(), + ) + config.persistent_dir.mkdir(parents=True, exist_ok=True) + with pytest.raises(AssertionError, match="opencode setup directory"): + OpenCodeHarnessProcessor(config=config)._get_upstream_run_command() + + +######################################## +# _retrieve_switchyard_trace filter_title_gen tests +######################################## + + +class TestRetrieveSwitchyardTraceFilterTitleGen: + # Records use the real Switchyard shape: messages flat on each record. + def test_filter_false_passes_envelope_unchanged(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + records = {"ses_1": [{"messages": [{"role": "system", "content": "You are a title generator"}]}]} + reconstruct = MagicMock(return_value="trace") + monkeypatch.setattr(swe_app, "reconstruct_switchyard_rollout", reconstruct) + + wrapper._retrieve_switchyard_trace(records, "ses_1", filter_title_gen=False) + assert len(reconstruct.call_args[0][0]["completions"]) == 1 + + def test_filter_true_strips_title_gen(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + records = { + "ses_1": [ + {"messages": [{"role": "system", "content": "You are a title generator"}]}, + {"messages": [{"role": "system", "content": "You are a coding assistant"}]}, + ] + } + reconstruct = MagicMock(return_value="trace") + monkeypatch.setattr(swe_app, "reconstruct_switchyard_rollout", reconstruct) + + wrapper._retrieve_switchyard_trace(records, "ses_1", filter_title_gen=True) + called_env = reconstruct.call_args[0][0] + assert len(called_env["completions"]) == 1 + assert called_env["completions"][0]["messages"][0]["content"] == "You are a coding assistant" + + +######################################## +# _reconstruct_switchyard_sessions filter_title_gen threading +######################################## + + +class TestReconstructFilterTitleGenThreading: + @staticmethod + def _fake_trace(sid: str): + item = MagicMock() + item.model_dump.return_value = {"sid": sid} + from nemo_gym.switchyard_trace import SwitchyardTrace + + return SwitchyardTrace(input_items=[item], output_items=[item], tools=[], record_uuids=[], model=sid) + + @pytest.mark.asyncio + async def test_filter_true_threaded_to_retrieve(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + retrieve = MagicMock(side_effect=lambda recs, sid, conv, filter_title_gen=False, **kw: self._fake_trace(sid)) + monkeypatch.setattr(swe_app, "_retrieve_switchyard_trace_from_records", retrieve) + sessions = [ + {"session_id": "root", "parent_session_id": None}, + {"session_id": "sub", "parent_session_id": "root"}, + ] + wrapper._reconstruct_switchyard_sessions({}, sessions, filter_title_gen=True) + for call in retrieve.call_args_list: + assert call.kwargs.get("filter_title_gen") is True + + @pytest.mark.asyncio + async def test_filter_defaults_false(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + retrieve = MagicMock(side_effect=lambda recs, sid, conv, filter_title_gen=False, **kw: self._fake_trace(sid)) + monkeypatch.setattr(swe_app, "_retrieve_switchyard_trace_from_records", retrieve) + sessions = [{"session_id": "only", "parent_session_id": None}] + wrapper._reconstruct_switchyard_sessions({}, sessions) + for call in retrieve.call_args_list: + assert call.kwargs.get("filter_title_gen") is False + + +######################################## +# run() gate passes filter_title_gen=True +######################################## + + +class TestRunOpencodeGateFilterTitleGen: + @staticmethod + def _opencode_response(**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}), + "instance_config": _make_instance_config( + tempfile.mkdtemp(), + agent_framework="opencode", + opencode_source="opencode", + # Retrieval is gated on spawn mode: that is what assigns the + # per-run capture dir the records are read back from. + switchyard_spawn_routing_profile="/tmp/profile.yaml", + switchyard_spawned_base_url="http://switchyard:4000", + **overrides, + ).model_dump_json(), + }, + ) + + @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": "{}", + }, + ) + ) + + @pytest.mark.asyncio + async def test_run_applies_opencode_payload(self, monkeypatch) -> None: + """run() consumes the trace the Ray task reconstructed on its own node.""" + wrapper = _create_wrapper(monkeypatch) + from nemo_gym.switchyard_trace import SwitchyardTrace + + fake_trace = SwitchyardTrace(input_items=[], output_items=[], tools=[], record_uuids=[], model="m") + wrapper._pending_switchyard["test_run_123"] = { + "trace": fake_trace, + "root_id": "root", + "subagent_trajectories": [], + "degraded": None, + "error": None, + } + with patch.object( + SWEBenchWrapper, "responses", new_callable=AsyncMock, return_value=self._opencode_response() + ): + result = await wrapper.run(self._run_body()) + + assert wrapper._pending_switchyard == {} + assert result.switchyard_trace_error is None + assert result.mask_sample is False + # Reconstructed subagents replace the text ones even when empty. + assert result.subagent_trajectories == [] + assert result.response.metadata["switchyard_session_id"] == "root" + + @pytest.mark.asyncio + async def test_run_masks_on_payload_error(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + wrapper._pending_switchyard["test_run_123"] = { + "trace": None, + "root_id": None, + "subagent_trajectories": None, + "degraded": None, + "error": "SwitchyardTraceError: no root session captured", + } + with patch.object( + SWEBenchWrapper, "responses", new_callable=AsyncMock, return_value=self._opencode_response() + ): + result = await wrapper.run(self._run_body()) + assert result.mask_sample is True + assert "opencode sessions" in result.switchyard_trace_error + assert "no root session captured" in result.switchyard_trace_error + + @pytest.mark.asyncio + async def test_run_masks_and_emits_tokens_on_degraded_payload(self, monkeypatch) -> None: + """Degraded (compaction) payloads still emit real tokens, plus a mask.""" + wrapper = _create_wrapper(monkeypatch) + from nemo_gym.switchyard_trace import SwitchyardTrace + + fake_trace = SwitchyardTrace(input_items=[], output_items=[], tools=[], record_uuids=[], model="m") + wrapper._pending_switchyard["test_run_123"] = { + "trace": fake_trace, + "root_id": "root", + "subagent_trajectories": [], + "degraded": "ambiguous session tree: 2 root sessions (likely context compaction)", + "error": None, + } + with patch.object( + SWEBenchWrapper, "responses", new_callable=AsyncMock, return_value=self._opencode_response() + ): + result = await wrapper.run(self._run_body()) + assert result.mask_sample is True + assert "2 root sessions" in result.switchyard_trace_error + assert result.response.metadata["switchyard_session_id"] == "root" + + @pytest.mark.asyncio + async def test_run_masks_when_payload_missing(self, monkeypatch) -> None: + """A rollout that never returned a payload cannot be trained on.""" + wrapper = _create_wrapper(monkeypatch) + with patch.object( + SWEBenchWrapper, "responses", new_callable=AsyncMock, return_value=self._opencode_response() + ): + result = await wrapper.run(self._run_body()) + assert result.mask_sample is True + assert "no switchyard payload" in result.switchyard_trace_error + + +######################################## +# _build_apptainer_command upstream mount tests +######################################## + + +class TestBuildApptainerCommandUpstreamOpencode: + def _setup_upstream_dirs(self, params: SWEBenchWrapperInstanceConfig) -> None: + opencode_dir = Path(str(params.opencode_setup_dir)) / "opencode" + bun_dir = Path(str(params.opencode_setup_dir)) / "bun" + (opencode_dir / "evaluation" / "oh").mkdir(parents=True, exist_ok=True) + bun_dir.mkdir(parents=True, exist_ok=True) + (params.persistent_dir / f"user_message_{params.agent_run_id}.txt").write_text("task") + + def test_upstream_agent_omits_migration_mount(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + with tempfile.TemporaryDirectory() as tmpdir: + opencode_setup_dir = Path(tmpdir) / "upstream_setup" + params = _make_instance_config( + tmpdir, + agent_framework="opencode", + opencode_source="opencode", + opencode_setup_dir=opencode_setup_dir, + switchyard_spawned_base_url="http://switchyard:4000", + ) + params.persistent_dir.mkdir(parents=True, exist_ok=True) + self._setup_upstream_dirs(params) + cmd_args = ExecuteContainerCommandArgs( + command="x", expected_file_pattern="/tmp/*.jsonl", mode="agent", timeout=300 + ) + result = wrapper._build_apptainer_command(params, cmd_args) + assert "/opencode_setup/bun" in result + assert "/opencode_setup/opencode" in result + assert "migration" not in result + + def test_fork_path_retains_migration_mount(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + with tempfile.TemporaryDirectory() as tmpdir: + opencode_setup_dir = Path(tmpdir) / "fork_setup" + params = _make_instance_config( + tmpdir, + agent_framework="opencode", + opencode_source="nv-opencode", + opencode_setup_dir=opencode_setup_dir, + ) + params.persistent_dir.mkdir(parents=True, exist_ok=True) + opencode_dir = opencode_setup_dir / "opencode" + (opencode_dir / "evaluation" / "oh").mkdir(parents=True, exist_ok=True) + (opencode_dir / "packages" / "opencode" / "migration").mkdir(parents=True, exist_ok=True) + (opencode_setup_dir / "bun").mkdir(parents=True, exist_ok=True) + (params.persistent_dir / f"user_message_{params.agent_run_id}.txt").write_text("x") + cmd_args = ExecuteContainerCommandArgs( + command="x", expected_file_pattern="/tmp/*.jsonl", mode="agent", timeout=300 + ) + result = wrapper._build_apptainer_command(params, cmd_args) + assert "migration" in result + + +######################################## +# Node-local Switchyard spawn (proxies run beside their agent, not on the server) +######################################## + + +class TestSwitchyardSpawnsOnAgentNode: + def test_backend_url_round_robins_across_policy_endpoints(self, monkeypatch) -> None: + """The server hands each run a backend so one vLLM node does not take the whole wave.""" + wrapper = _create_wrapper(monkeypatch) + cfg = shlex.quote(OmegaConf.to_yaml(OmegaConf.create({"policy_base_url": ["http://a/v1", "http://b/v1"]}))) + with tempfile.TemporaryDirectory() as tmpdir: + params = _make_instance_config(tmpdir, ng_global_config_dict_str=cfg) + picked = [wrapper._next_switchyard_backend_url(params) for _ in range(4)] + + assert picked == ["http://a/v1", "http://b/v1", "http://a/v1", "http://b/v1"] + + def test_backend_url_none_when_policy_exposes_no_endpoint(self, monkeypatch) -> None: + wrapper = _create_wrapper(monkeypatch) + cfg = shlex.quote(OmegaConf.to_yaml(OmegaConf.create({}))) + with tempfile.TemporaryDirectory() as tmpdir: + params = _make_instance_config(tmpdir, ng_global_config_dict_str=cfg) + assert wrapper._next_switchyard_backend_url(params) is None + + def test_spawn_needed_decided_from_run_config(self) -> None: + """The Ray task has no server object, so the predicate reads the run's own config.""" + with tempfile.TemporaryDirectory() as tmpdir: + spawning = _make_instance_config( + tmpdir, + switchyard_spawn_routing_profile="/tmp/profile.yaml", + agent_framework="opencode", + opencode_source="opencode", + ) + assert swe_app._switchyard_spawn_needed_for(spawning) is True + + no_profile = _make_instance_config(tmpdir, agent_framework="opencode", opencode_source="opencode") + assert swe_app._switchyard_spawn_needed_for(no_profile) is False + + golden = _make_instance_config( + tmpdir, + switchyard_spawn_routing_profile="/tmp/profile.yaml", + agent_framework="opencode", + opencode_source="opencode", + verify_golden_patch=True, + ) + assert swe_app._switchyard_spawn_needed_for(golden) is False + + async def test_agent_run_reaps_its_switchyard(self, monkeypatch) -> None: + """The task that starts a proxy also stops it — nothing outlives its rollout.""" + with tempfile.TemporaryDirectory() as tmpdir: + params = _make_instance_config( + tmpdir, + switchyard_spawn_routing_profile="/tmp/profile.yaml", + agent_framework="opencode", + opencode_source="opencode", + ) + process = MagicMock() + monkeypatch.setattr( + swe_app, "_spawn_switchyard_local", AsyncMock(return_value=("http://node-7:1234", process)) + ) + monkeypatch.setattr(swe_app, "_teardown_switchyard_process", AsyncMock()) + monkeypatch.setattr( + SWEBenchWrapper, "_fetch_switchyard_max_model_len", AsyncMock(return_value=4096) + ) + monkeypatch.setattr(SWEBenchWrapper, "_build_agent_command", MagicMock()) + monkeypatch.setattr( + swe_app.RunOpenHandsAgent, "process_single_datapoint", AsyncMock(return_value=Path("/tmp/r.json")) + ) + + result = await swe_app._run_agent_with_switchyard(params) + + assert result["report_file"] == "/tmp/r.json" + assert result["switchyard"] is not None + assert params.switchyard_spawned_base_url == "http://node-7:1234" + assert params.opencode_context_len == 4096 + swe_app._teardown_switchyard_process.assert_awaited_once_with(process) + + async def test_switchyard_reaped_even_when_agent_raises(self, monkeypatch) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + params = _make_instance_config( + tmpdir, + switchyard_spawn_routing_profile="/tmp/profile.yaml", + agent_framework="opencode", + opencode_source="opencode", + ) + process = MagicMock() + monkeypatch.setattr( + swe_app, "_spawn_switchyard_local", AsyncMock(return_value=("http://node-7:1234", process)) + ) + monkeypatch.setattr(swe_app, "_teardown_switchyard_process", AsyncMock()) + monkeypatch.setattr( + SWEBenchWrapper, "_fetch_switchyard_max_model_len", AsyncMock(return_value=None) + ) + monkeypatch.setattr(SWEBenchWrapper, "_build_agent_command", MagicMock()) + monkeypatch.setattr( + swe_app.RunOpenHandsAgent, "process_single_datapoint", AsyncMock(side_effect=RuntimeError("boom")) + ) + + with pytest.raises(RuntimeError): + await swe_app._run_agent_with_switchyard(params) + + swe_app._teardown_switchyard_process.assert_awaited_once_with(process) + + def test_retrieval_gate_does_not_require_a_proxy_url(self, monkeypatch) -> None: + """Regression: the proxy URL is set on the agent's node and never returns to the + server, so gating retrieval on it silently skipped every rollout.""" + wrapper = _create_wrapper(monkeypatch) + with tempfile.TemporaryDirectory() as tmpdir: + rl_log_dir = Path(tmpdir) / "switchyard_traces" + session_dir = rl_log_dir / "sessions" / "dir_root" + session_dir.mkdir(parents=True, exist_ok=True) + (session_dir / "u1.json").write_text( + json.dumps({"schema_version": 1, "session_id": "root", "parent_session_id": None, "uuid": "u1"}) + ) + + # No proxy URL anywhere: retrieval is disk-only. + records = wrapper._read_switchyard_records(rl_log_dir) + sessions = wrapper._list_switchyard_sessions(rl_log_dir) + + assert records and sessions, "capture must be readable without any proxy URL" + + +class TestSpawnProfileEnv: + """The spawn env carries every ${VAR} the shipped routing profile references.""" + + @staticmethod + def _policy_cfg() -> str: + return shlex.quote( + OmegaConf.to_yaml( + OmegaConf.create( + {"policy_base_url": ["http://gen:8000/v1"], "policy_model_name": "/ckpts/policy/hf"} + ) + ) + ) + + async def _spawn_env(self, monkeypatch, tmpdir: str, **overrides) -> dict: + params = _make_instance_config( + tmpdir, + ng_global_config_dict_str=self._policy_cfg(), + switchyard_spawn_routing_profile="/tmp/profile.yaml", + agent_framework="opencode", + opencode_source="opencode", + **overrides, + ) + captured = {} + + async def fake_exec(*argv, **kwargs): + captured["env"] = kwargs["env"] + return MagicMock() + + monkeypatch.delenv("SWITCHYARD_VLLM_BASE_URL", raising=False) + monkeypatch.delenv("SWITCHYARD_POLICY_MODEL", raising=False) + monkeypatch.setattr(swe_app.asyncio, "create_subprocess_exec", fake_exec) + monkeypatch.setattr(swe_app, "_wait_switchyard_ready_local", AsyncMock()) + await swe_app._spawn_switchyard_local(params) + return captured["env"] + + async def test_endpoint_and_model_derived_from_run_config(self, monkeypatch) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + env = await self._spawn_env(monkeypatch, tmpdir) + assert env["SWITCHYARD_VLLM_BASE_URL"] == "http://gen:8000/v1" + assert env["SWITCHYARD_POLICY_MODEL"] == "/ckpts/policy/hf" + + async def test_round_robin_backend_wins_over_derived_url(self, monkeypatch) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + env = await self._spawn_env(monkeypatch, tmpdir, switchyard_backend_url="http://rr:1/v1") + assert env["SWITCHYARD_VLLM_BASE_URL"] == "http://rr:1/v1" + + async def test_parsers_and_pythonpath_from_config_fields(self, monkeypatch) -> None: + monkeypatch.setenv("PYTHONPATH", "/existing") + with tempfile.TemporaryDirectory() as tmpdir: + env = await self._spawn_env( + monkeypatch, + tmpdir, + switchyard_tool_parser="qwen3_coder", + switchyard_reasoning_parser="nano_v3", + switchyard_parser_pythonpath="/opt/vllm", + ) + assert env["SWITCHYARD_TOOL_PARSER"] == "qwen3_coder" + assert env["SWITCHYARD_REASONING_PARSER"] == "nano_v3" + assert env["PYTHONPATH"] == "/opt/vllm:/existing" + + async def test_unset_fields_leave_env_untouched(self, monkeypatch) -> None: + monkeypatch.setenv("PYTHONPATH", "/existing") + with tempfile.TemporaryDirectory() as tmpdir: + env = await self._spawn_env(monkeypatch, tmpdir) + assert "SWITCHYARD_TOOL_PARSER" not in env + assert "SWITCHYARD_REASONING_PARSER" not in env + assert env["PYTHONPATH"] == "/existing" + + +class TestShippedSwitchyardProfile: + """The default profile stays in lockstep with what the spawn env exports.""" + + _PROFILE_PATH = Path(swe_app.__file__).parent / "switchyard_profile.yaml" + + def test_env_refs_match_spawn_exports(self) -> None: + import re + + refs = set(re.findall(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}", self._PROFILE_PATH.read_text())) + assert refs == { + "SWITCHYARD_VLLM_BASE_URL", + "SWITCHYARD_POLICY_MODEL", + "SWITCHYARD_TOOL_PARSER", + "SWITCHYARD_REASONING_PARSER", + } + + def test_routes_are_aliases_with_injection_and_numeric_sampling(self) -> None: + import yaml + + loaded = yaml.safe_load(self._PROFILE_PATH.read_text()) + # Dispatch is exact-match on the request's model id; upstream opencode + # POSTs "default", the OpenHands fork POSTs the model server's name. + assert loaded["routes"]["default"] == loaded["routes"]["policy_model"] + route = loaded["routes"]["default"] + assert route["token_capture_engine"] == "vllm" + assert route["token_injection"] is True + assert route["injection_transport"] == "prefix_chat" + # Env interpolation is string-only: numeric sampling params must stay + # literals or the trainer's on-policy assert sees strings. + assert isinstance(route["extra_body"]["temperature"], float) + assert isinstance(route["extra_body"]["top_p"], float) + + +def _write_full_record(rl_log_dir: Path, session_id: str, parent_id, turn: int, history: list) -> list: + """One complete capture record (the shape Switchyard writes); returns new history.""" + generation = [100 + turn, 101 + turn] + session_dir = rl_log_dir / "sessions" / f"dir_{session_id}" + session_dir.mkdir(parents=True, exist_ok=True) + record = { + "schema_version": 1, + "session_id": session_id, + "parent_session_id": parent_id, + "uuid": f"{session_id}-{turn:03d}", + "captured_at": f"2026-01-01T00:00:{turn:02d}", + "model": "test-model", + "is_valid": True, + "finish_reason": "stop", + # Cumulative conversation: record N's prompt must extend record N-1's + # history (validator), so each record carries the whole chat so far. + "messages": [ + m for t in range(turn + 1) for m in ( + {"role": "user", "content": f"turn {t}"}, + {"role": "assistant", "content": f"reply {t}"}, + ) + ], + "tools": [], + "tool_choice": None, + "request_id": f"req-{turn}", + "token_count": len(history) + len(generation), + "prompt_token_ids": list(history), + "generation_token_ids": generation, + "generation_log_probs": [-0.1] * len(generation), + } + (session_dir / f"{record['uuid']}.json").write_text(json.dumps(record)) + return history + generation + + +class TestCollectSwitchyardPayload: + """Task-side retrieval: real records, real converter, no mocks. + + This is the code that runs inside runner_ray_remote on the rollout's node; + these tests exercise it standalone (Gate 2 exercises it inside a real Ray + task, including payload picklability through the object store). + """ + + def _params(self, tmpdir: str, **overrides): + defaults = dict( + switchyard_spawn_routing_profile="/tmp/profile.yaml", + agent_framework="opencode", + opencode_source="opencode", + ) + return _make_instance_config(tmpdir, **{**defaults, **overrides}) + + def _write_tree(self, rl_log_dir: Path) -> None: + history = [1, 2, 3] + for turn in range(3): + history = _write_full_record(rl_log_dir, "ses_root", None, turn, history) + _write_full_record(rl_log_dir, "ses_sub", "ses_root", 0, [1, 2, 3]) + + def test_opencode_tree_reconstructs(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + params = self._params(tmpdir) + self._write_tree(params.persistent_dir / "switchyard_traces") + payload = swe_app._collect_switchyard_payload(params) + + assert payload["error"] is None + assert payload["degraded"] is None + assert payload["root_id"] == "ses_root" + assert payload["trace"] is not None and payload["trace"].output_items + assert len(payload["subagent_trajectories"]) == 1 + assert payload["subagent_trajectories"][0]["session_id"] == "ses_sub" + + def test_two_roots_degrade_but_emit_tokens(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + params = self._params(tmpdir) + rl_log_dir = params.persistent_dir / "switchyard_traces" + self._write_tree(rl_log_dir) + _write_full_record(rl_log_dir, "ses_compacted", None, 0, [9, 9]) + payload = swe_app._collect_switchyard_payload(params) + + assert payload["error"] is None + assert payload["degraded"] and "2 root sessions" in payload["degraded"] + assert payload["root_id"] == "ses_root" # largest chain wins + assert payload["trace"] is not None and payload["trace"].output_items + + def test_no_records_is_error_not_raise(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + params = self._params(tmpdir) + payload = swe_app._collect_switchyard_payload(params) + assert payload["trace"] is None + assert "no root session captured" in payload["error"] + + def test_spawn_not_needed_returns_none(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + params = _make_instance_config(tmpdir) # no profile, openhands, no session + assert swe_app._collect_switchyard_payload(params) is None + + def test_openhands_session_path(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + params = self._params( + tmpdir, + agent_framework="openhands", + opencode_source="nv-opencode", + switchyard_session_id="ses_root", + ) + history = [1, 2, 3] + for turn in range(2): + history = _write_full_record( + params.persistent_dir / "switchyard_traces", "ses_root", None, turn, history + ) + payload = swe_app._collect_switchyard_payload(params) + + assert payload["error"] is None + assert payload["root_id"] == "ses_root" + assert payload["trace"] is not None and payload["trace"].output_items + + def test_payload_pickles(self) -> None: + """The payload crosses the Ray object store; plain pickle is the stricter proxy.""" + import pickle + + with tempfile.TemporaryDirectory() as tmpdir: + params = self._params(tmpdir) + self._write_tree(params.persistent_dir / "switchyard_traces") + payload = swe_app._collect_switchyard_payload(params) + restored = pickle.loads(pickle.dumps(payload)) + assert restored["trace"].output_items + assert restored["root_id"] == "ses_root" + + +class TestCollectSwitchyardPayloadPartial: + """A mid-chain validation break emits the longest valid prefix + mask, + so one bad rollout costs one masked sample instead of its whole group.""" + + def _params(self, tmpdir: str): + return _make_instance_config( + tmpdir, + switchyard_spawn_routing_profile="/tmp/profile.yaml", + agent_framework="opencode", + opencode_source="opencode", + ) + + def test_midchain_break_emits_prefix_and_degrades(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + params = self._params(tmpdir) + rl_log_dir = params.persistent_dir / "switchyard_traces" + history = [1, 2, 3] + for turn in range(2): + history = _write_full_record(rl_log_dir, "ses_root", None, turn, history) + # Turn 2 breaks token continuity: prompt does not extend history. + _write_full_record(rl_log_dir, "ses_root", None, 2, [999, 998]) + payload = swe_app._collect_switchyard_payload(params) + + assert payload["error"] is None + assert payload["trace"] is not None and payload["trace"].output_items + assert payload["degraded"] and "partial trace" in payload["degraded"] + assert "record 2" in payload["degraded"] + # Only the two valid records made it in. + assert payload["trace"].record_uuids == ["ses_root-000", "ses_root-001"] + + def test_record_zero_break_is_still_an_error(self) -> None: + """No valid prefix exists -> error path (empty output, masked, group cost). + Deliberately unchanged: emitting fabricated tokens would be worse.""" + with tempfile.TemporaryDirectory() as tmpdir: + params = self._params(tmpdir) + rl_log_dir = params.persistent_dir / "switchyard_traces" + sd = rl_log_dir / "sessions" / "dir_ses_root" + sd.mkdir(parents=True) + (sd / "r0.json").write_text(json.dumps({ + "schema_version": 1, "session_id": "ses_root", "parent_session_id": None, + "uuid": "r0", "captured_at": "t0", "model": "test-model", "is_valid": True, + "finish_reason": "stop", + "messages": [{"role": "user", "content": "only a user message"}], + "tools": [], "tool_choice": None, "request_id": "r", + "token_count": 2, "prompt_token_ids": [1], "generation_token_ids": [2], + "generation_log_probs": [-0.1], + })) + payload = swe_app._collect_switchyard_payload(params) + assert payload["trace"] is None + assert "assistant message" in payload["error"] + + def test_default_reconstruction_is_partial(self) -> None: + """Partial is the DEFAULT (training is the primary consumer); strict + fail-fast is the explicit opt-in for analysis/tests.""" + from nemo_gym.switchyard_trace import SwitchyardTraceError, reconstruct_switchyard_rollout + from responses_api_models.vllm_model.app import VLLMConverter + + with tempfile.TemporaryDirectory() as tmpdir: + rl_log_dir = Path(tmpdir) / "switchyard_traces" + history = [1, 2, 3] + history = _write_full_record(rl_log_dir, "s", None, 0, history) + _write_full_record(rl_log_dir, "s", None, 1, [777]) + records = SWEBenchWrapper._read_switchyard_records(rl_log_dir) + envelope = {"schema_version": 1, "session_id": "s", "completions": records["s"]} + conv = VLLMConverter(return_token_id_information=True) + trace = reconstruct_switchyard_rollout(envelope, "s", conv) + assert trace.partial_reason and "do not extend" in trace.partial_reason + assert trace.record_uuids == ["s-000"] + with pytest.raises(SwitchyardTraceError, match="do not extend"): + reconstruct_switchyard_rollout(envelope, "s", conv, allow_partial=False) diff --git a/responses_api_agents/swe_agents/tests/test_switchyard_spawn_real.py b/responses_api_agents/swe_agents/tests/test_switchyard_spawn_real.py new file mode 100644 index 0000000000..3616715249 --- /dev/null +++ b/responses_api_agents/swe_agents/tests/test_switchyard_spawn_real.py @@ -0,0 +1,195 @@ +"""Real-process regression test for the per-rollout Switchyard lifecycle. + +The core of this test runs INSIDE a real Ray task with the production +decorator options — deliberately unmocked. The rest of the suite mocks +``asyncio.create_subprocess_exec``, which verifies the argv we build but can +never fail on what actually broke in production three separate times: state +the Gym server process has that a Ray worker does not (PATH, environment, +``sys.path``). Two 32-node training allocations were lost to bugs this file's +checks catch in under a minute on one CPU. + +Covered in one task round-trip: +- the ``switchyard`` CLI resolves without PATH (interpreter-adjacent fallback); +- the spawned proxy reaches READY (binds and accepts TCP) — a proxy that + spawns but exits pre-bind fails every rollout; +- params cross the boundary exactly as production sends them (``model_dump`` + -> ``model_validate``); +- task-side retrieval reconstructs a real trace from seeded records, and the + payload (pydantic models included) survives the Ray object store; +- teardown is clean. + +The venv wrapper's env derivation (NEMO_GYM_CONFIG_DICT -> +SWITCHYARD_POLICY_MODEL) is deployment-specific and exercised by the +deployment's own gate; the routing profile here uses literal values so the +test runs on any install with the real CLI. + +Skips when ray or the switchyard CLI is unavailable. If the Ray worker +unexpectedly has the CLI on PATH, the environment does not reproduce +production and the test reports itself inconclusive via skip rather than +asserting a vacuous pass. +""" + +import asyncio +import json +import shutil +import socket +import sys +import tempfile +import time +from pathlib import Path + +import pytest + + +ray = pytest.importorskip("ray") + + +_CLI_AVAILABLE = shutil.which("switchyard") is not None or (Path(sys.executable).parent / "switchyard").exists() + +pytestmark = pytest.mark.skipif(not _CLI_AVAILABLE, reason="switchyard CLI not installed") + +_PROFILE_YAML = """\ +defaults: + api_key: dummy + base_url: http://127.0.0.1:9999/v1 + format: openai +routes: + default: + type: model + target: test-model + format: openai +""" + + +def _seed_records(rl_log_dir: Path, session_id: str = "ses_root", turns: int = 3) -> None: + """Records in Switchyard's on-disk shape: cumulative messages, extending tokens.""" + history = [1, 2, 3] + d = rl_log_dir / "sessions" / f"dir_{session_id}" + d.mkdir(parents=True, exist_ok=True) + for turn in range(turns): + generation = [100 + turn, 101 + turn] + record = { + "schema_version": 1, + "session_id": session_id, + "parent_session_id": None, + "uuid": f"{session_id}-{turn:03d}", + "captured_at": f"2026-01-01T00:00:{turn:02d}", + "model": "test-model", + "is_valid": True, + "finish_reason": "stop", + "messages": [ + m + for t in range(turn + 1) + for m in ( + {"role": "user", "content": f"turn {t}"}, + {"role": "assistant", "content": f"reply {t}"}, + ) + ], + "tools": [], + "tool_choice": None, + "request_id": f"req-{turn}", + "token_count": len(history) + len(generation), + "prompt_token_ids": list(history), + "generation_token_ids": generation, + "generation_log_probs": [-0.1] * len(generation), + } + (d / f"{record['uuid']}.json").write_text(json.dumps(record)) + history = history + generation + + +def _task_body(gym_dir: str, params_dict: dict) -> dict: + """Runs inside the Ray worker. + + Under pytest this module has a real dotted name, so Ray exports this + function BY REFERENCE: the worker imports the module before running the + body, which is why the decorator passes PYTHONPATH in runtime_env — the + driver's sys.path does not cross the boundary. The inserts below are what + the body itself needs for its bare `import app`.""" + sys.path.insert(0, gym_dir) + sys.path.insert(0, str(Path(gym_dir) / "responses_api_agents" / "swe_agents")) + + import app as swe_app + + swe_app.SWEBenchWrapperInstanceConfig.model_rebuild(force=True) + params = swe_app.SWEBenchWrapperInstanceConfig.model_validate(params_dict) + out: dict = {} + + async def _run() -> None: + # Negative control: the pre-fix form must fail here, else this worker + # has the venv on PATH and the resolution fallback goes unexercised. + try: + proc = await asyncio.create_subprocess_exec( + "switchyard", "--help", stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.DEVNULL + ) + await proc.wait() + out["control_raised"] = False + except FileNotFoundError: + out["control_raised"] = True + + process = None + try: + base_url, process = await swe_app._spawn_switchyard_local(params) + out["spawn_ok"] = True + host, port = base_url.removeprefix("http://").split(":") + deadline = time.monotonic() + 5 + ok = False + while time.monotonic() < deadline and not ok: + with socket.socket() as s: + s.settimeout(1) + ok = s.connect_ex((host, int(port))) == 0 + out["tcp_accepting"] = ok + out["process_alive"] = process.returncode is None + finally: + if process is not None: + await swe_app._teardown_switchyard_process(process) + out["teardown_ok"] = True + + out["payload"] = swe_app._collect_switchyard_payload(params) + + asyncio.run(_run()) + return out + + +def test_spawn_ready_and_retrieval_in_real_ray_worker(tmp_path: Path) -> None: + from swe_agents.tests.test_app import _make_instance_config + + gym_dir = str(Path(__file__).resolve().parents[3]) + profile = tmp_path / "profile.yaml" + profile.write_text(_PROFILE_YAML) + + with tempfile.TemporaryDirectory() as tmpdir: + params = _make_instance_config( + tmpdir, + switchyard_spawn_routing_profile=str(profile), + agent_framework="opencode", + opencode_source="opencode", + ) + _seed_records(params.persistent_dir / "switchyard_traces") + + ray.init(ignore_reinit_error=True, include_dashboard=False) + try: + task = ray.remote( + runtime_env={ + "py_executable": sys.executable, + # by-reference export: the worker must be able to import this + # module's package (responses_api_agents/) itself + "env_vars": {"PYTHONPATH": str(Path(gym_dir) / "responses_api_agents")}, + }, + num_cpus=0.1, + )(_task_body) + out = ray.get(task.remote(gym_dir, params.model_dump()), timeout=540) + finally: + ray.shutdown() + + if out.get("control_raised") is False: + pytest.skip("Ray worker has switchyard on PATH — environment does not reproduce production") + + assert out["spawn_ok"], "spawn failed in the Ray worker" + assert out["tcp_accepting"], "proxy never reached READY (bound port)" + assert out["process_alive"] + assert out["teardown_ok"] + + payload = out["payload"] # crossed the object store via ray.get + assert payload is not None and payload.get("error") is None + assert payload["root_id"] == "ses_root" + assert payload["trace"].output_items, "reconstructed trace lost its tokens crossing the object store" diff --git a/tests/unit_tests/test_switchyard_trace.py b/tests/unit_tests/test_switchyard_trace.py index 0057a879a2..b49bcbb3c7 100644 --- a/tests/unit_tests/test_switchyard_trace.py +++ b/tests/unit_tests/test_switchyard_trace.py @@ -91,7 +91,9 @@ def _envelope(records: list, **overrides) -> dict: def _reconstruct(envelope: dict): converter = ResponsesConverter(return_token_id_information=True) - return reconstruct_switchyard_rollout(envelope, SESSION_ID, converter) + # These tests certify the strict validator; the default is allow_partial=True + # for the training consumer, so strictness is requested explicitly here. + return reconstruct_switchyard_rollout(envelope, SESSION_ID, converter, allow_partial=False) class TestReconstruction: @@ -126,6 +128,48 @@ def test_two_call_tool_trajectory(self) -> None: assert trace.tools[0].name == "str_replace_editor" assert trace.tools[0].description == "Edit files" assert trace.tools[0].parameters == TOOL_SCHEMA + + def test_reformatted_tool_args_tolerated_and_original_kept(self) -> None: + # The model emits pretty-printed JSON args; the client re-sends them compact + # in the next turn's history. The extension check must tolerate that + # (canonicalized comparison), and the reconstructed item must keep the + # ORIGINAL arguments (canonicalization is comparison-only). + pretty = '{\n "path": "a.py"\n}' + compact = '{"path": "a.py"}' + first = _record( + 0, + [ + {"role": "system", "content": "You are a SWE agent."}, + {"role": "user", "content": "Fix the bug."}, + { + "role": "assistant", + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "str_replace_editor", "arguments": pretty}} + ], + }, + ], + TRIPLE_0, + ) + second = _record( + 1, + [ + {"role": "system", "content": "You are a SWE agent."}, + {"role": "user", "content": "Fix the bug."}, + { + "role": "assistant", + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "str_replace_editor", "arguments": compact}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "edited ok"}, + {"role": "assistant", "content": "All fixed."}, + ], + TRIPLE_1, + ) + + trace = _reconstruct(_envelope([first, second])) # tolerated — does not raise + + assert trace.output_items[0].arguments == pretty # original kept, not canonicalized assert trace.tools[0].type == "function" def test_single_record_session(self) -> None: