diff --git a/benchmark/README.md b/benchmark/README.md index dfaa36a52..74fa8ec13 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -13,6 +13,9 @@ Both paths use the same generated dataset, task proxy, pinned agent versions, an layout. Passing `--server-config` starts the Rust server; omitting it disables Switchyard and points Harbor directly at the upstream provider. +For a small automated MMLU-Redux example using NeMo Gym instead of Harbor, see +[Evaluate Switchyard routing with NeMo Gym](nemo_gym/README.md). + ## Prerequisites From the repo root: diff --git a/benchmark/nemo_gym/README.md b/benchmark/nemo_gym/README.md new file mode 100644 index 000000000..9e00ee48f --- /dev/null +++ b/benchmark/nemo_gym/README.md @@ -0,0 +1,161 @@ +# Evaluate Switchyard routing with NeMo Gym + +[NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) is a library for evaluating and improving models and agents, combining infrastructure for developing environments and running evaluation and training at scale with popular benchmarks and training environments. + +Gym provides the evaluation substrate: it supplies benchmark tasks, runs the model interaction, verifies answers, and reports rewards. Switchyard sits in the request path and selects the upstream model. + +This tutorial uses [MMLU-Redux 2.0](https://huggingface.co/datasets/edinburgh-dawg/mmlu-redux-2.0) to compare a fixed model with Switchyard routing. + +We compare a fixed Nemotron 3 Ultra baseline with seeded Random routing between Ultra and Nemotron 3.5 Lightning. Both conditions use the same questions and generation settings. + +![Gym evaluation through LiteLLM and Switchyard Random routing](architecture.svg) + +## 1. Set up + +You need: + +- Bash on Linux/macOS +- Git and curl +- [uv](https://docs.astral.sh/uv/) +- The [Rust toolchain prerequisites](../../docs/getting_started.md#prerequisites) for the current checkout bindings +- An NVIDIA API key from [build.nvidia.com](https://build.nvidia.com/) with access to `nvidia/nemotron-3-ultra-550b-a55b` and `nvidia/nemotron-3.5-lightning-30b-a3b` + +Run these commands in Bash from the Switchyard repository root. These one-time commands create a Gym checkout under `scratch/` at the tested `v0.6.0` release. Choose an unused `GYM_DIR` without spaces or shell metacharacters. + +```bash +export GYM_DIR="$PWD/scratch/nemo-gym-litellm/Gym" +mkdir -p "$(dirname "$GYM_DIR")" && +git clone https://github.com/NVIDIA-NeMo/Gym.git "$GYM_DIR" && +git -C "$GYM_DIR" checkout v0.6.0 && +uv tool run --from uv==0.11.29 uv sync \ + --directory "$GYM_DIR" --frozen --no-dev --python 3.13.14 && +uv tool run --from uv==0.11.29 uv pip install --no-deps \ + --python "$GYM_DIR/.venv/bin/python" uv==0.11.29 && +"$GYM_DIR/.venv/bin/uv" sync --project examples/litellm --locked --python 3.12 +``` + +Gym and LiteLLM use separate Python environments. The proxy builds Switchyard bindings from this checkout; the native CLI workflow does not need Docker. + +## 2. Run both conditions + +The default is five tasks per condition with an 8,192-token output limit, normally ten upstream calls when every request succeeds on its first attempt. + +Replace the placeholder below with your NVIDIA API key, then run. + +```bash +export NVIDIA_API_KEY="your-api-key" +bash benchmark/nemo_gym/run.sh +``` + +The [runner](./run.sh) prepares MMLU-Redux, starts the local LiteLLM proxy, evaluates fixed then routed, stops the proxy, and prints the comparison. It saves `comparison.txt` and other artifacts under `benchmark/nemo_gym/results//`. + +You do not need to run these commands separately, but at a high level, the runner does: + +```text +gym eval prepare --benchmark mmlu-redux +gym eval run --benchmark mmlu-redux --model-type litellm_model --model ... +``` + +Runtime varies with endpoint load, model latency, and retries. + +### How the requests are wired + +Gym's `litellm_model` adapter connects to the local proxy through `policy_base_url`. The runner requests `--model fixed`, then `--model routed`; these are the model-group names in [litellm.yaml](litellm.yaml): + +| Group | Available models | +|---|---| +| `fixed` | Nemotron 3 Ultra | +| `routed` | Nemotron 3 Ultra and Nemotron 3.5 Lightning | + +These are LiteLLM group names, not special Gym modes. The same Ultra configuration is reused in both groups, and thinking is disabled for both models for this short multiple-choice workload. + +The existing [Switchyard–LiteLLM integration](../../examples/litellm/README.md) chooses from each group using [routes.toml](routes.toml), which sets `algorithm = "random"` and `seed = 6`. Switchyard runs as a library inside LiteLLM, not as another server. + +A small tutorial callback fixes response fields needed by Gym `v0.6.0` and records which deployment served each response. It does not choose models or count tokens. Gym supplies the captured-call usage and errors. + +This tutorial demonstrates the LiteLLM path. See Gym's [full Switchyard integration documentation](https://docs.nvidia.com/nemo/gym/main/model-server/switchyard/) for its other deployment modes and configuration options. + +Gym components can outlive the command briefly; let them finish shutting down before an immediate rerun. + +## 3. Understand the result + +Start with task coverage and serving models, then compare rewards, tokens, and latency. This is the output from a live run of the default five tasks per condition: + +```text +fixed: expected=5, completed=5, missing=0, unexpected=0, failures=0 +routed: expected=5, completed=5, missing=0, unexpected=0, failures=0 +Pairing: matched=5, fixed-only=0, routed-only=0 + +Metric fixed routed +Paired rollouts 5 5 +Incomplete responses 0 0 +Mean reward 0.400 1.000 +Captured input tokens 598 598 +Captured output tokens 30 754 +Mean rollout latency (ms) 6645.761 9763.644 +Captured calls 5 5 +Captured failed calls 0 0 +Calls with unknown usage 0 0 + +fixed serving models: {"nvidia_nim/nvidia/nemotron-3-ultra-550b-a55b": 5} + +routed serving models: {"nvidia_nim/nvidia/nemotron-3-ultra-550b-a55b": 2, "nvidia_nim/nvidia/nemotron-3.5-lightning-30b-a3b": 3} + +Gym-scored incomplete responses remain in the comparison. Gym-captured tokens include extra attempts; +unknown means some usage was not reported. +Random makes no classifier calls. A small run demonstrates the integration, not a routing advantage. +``` + +- **Pairing:** the comparator matches fixed and routed rollouts by task and repeat, then confirms their inputs and settings are identical. Missing results or terminal Gym failures stop the comparison so averages are not calculated from different workloads. +- **Incomplete responses:** a model can hit its output limit while Gym still records and scores the rollout. These responses remain in the comparison and are reported with a warning. +- **Models:** fixed stayed on Ultra, while routed used Ultra and Lightning. The counts describe the final serving deployment, not every attempted model. A small Random run need not split evenly. +- **Reward:** Gym's multiple-choice verifier scores the boxed answer letter (correct = 1, wrong = 0). +- **Tokens:** input and output counts include all calls Gym captured, including extra attempts. + +Calls retried inside Gym's model adapter, LiteLLM, or the upstream provider may not appear separately. These are Gym-captured measurements, not a complete provider bill or fallback trace. + +Tokens are not dollar costs. The default five-task prefix is a smoke test, not a representative MMLU-Redux score; Random is not capability-based routing. + +## 4. (Optional) Try a small change + +- **Workload:** use a fresh results directory and adjust the task count: + + ```bash + RESULTS_DIR="$PWD/benchmark/nemo_gym/results/my-run" \ + LIMIT=2 bash benchmark/nemo_gym/run.sh + ``` + +- **Models:** edit [litellm.yaml](litellm.yaml), keeping one fixed candidate and two distinct routed candidates, including the fixed model. Keep per-model settings identical between conditions. +- **Routing:** change the seed in [routes.toml](routes.toml), keeping `algorithm = "random"`. A seed repeats assignments only for an identical request sequence; retries or concurrency can change them. +- **Another benchmark:** use `--benchmark NAME` in your own Gym calls against LiteLLM, with a compatible agent and verifier. The included runner is configured for MMLU-Redux; adapt the script and comparison logic for other benchmarks. + +See `bash benchmark/nemo_gym/run.sh --help` for output, port, repeat, and concurrency options. + +
+Saved files and token counts + +Each `fixed/` and `routed/` folder contains: + +- `rollouts.jsonl`: model responses and rewards. +- `rollouts_materialized_inputs.jsonl`: the tasks and settings used. +- `rollouts_failures.jsonl`: failed tasks, if any. +- `model-calls/`: Gym's captured model exchanges. +- `models.jsonl`: response IDs and serving models. + +The result root also contains `comparison.txt` and `versions.txt`, recorded for reference. + +If something fails, start with that folder's `gym.log` or the result folder's `litellm.log`. + +To view the comparison again without calling the models, replace `my-run` with your results folder: + +```bash +"$GYM_DIR/.venv/bin/python" benchmark/nemo_gym/compare.py \ + benchmark/nemo_gym/results/my-run +``` + +Older result folders without `models.jsonl` do not support this report; keep their saved `comparison.txt`. + +Keep the saved inputs because the source dataset can change. Request logs contain prompts and responses, so review them before sharing. + +
+ +**Validation:** tested with Gym `v0.6.0`, and as a live run of five tasks per condition. Ultra and Lightning were both exercised, and all ten captured calls succeeded. diff --git a/benchmark/nemo_gym/architecture.svg b/benchmark/nemo_gym/architecture.svg new file mode 100644 index 000000000..4bc1f9fb1 --- /dev/null +++ b/benchmark/nemo_gym/architecture.svg @@ -0,0 +1,88 @@ + + Evaluating Switchyard routing with NeMo Gym + MMLU-Redux 2.0 supplies the multiple-choice questions. Gym's simple_agent initializes each task session, obtains a model answer, and sends it to the mcqa resources server, which grades the boxed answer letter against the expected answer. Gym's litellm_model calls a runner-managed LiteLLM proxy through the OpenAI Responses API. Switchyard Random routing runs as a library inside LiteLLM, not as a separate Switchyard server. Compare fixed and routed runs on identical inputs using rewards, serving models, Gym-captured input and output tokens, latency and captured errors. + SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + + + + + + + + + Evaluating Switchyard routing with NeMo Gym + Run the same tasks with a fixed model or a router, then compare rewards, tokens and latency + + + NeMo Gym + + + Dataset + (tasks) + MMLU-Redux 2.0 + + + Initialize + Task session + + + Agent + Executes + Agent server + simple_agent + + + Verify + Resources server + Multiple-choice verifier + + + + + + + + Model request / response + + Gym model server + litellm_model + + OpenAI Responses API + + LiteLLM proxy + + Switchyard library + + + + Upstream models + + Nemotron 3.5 Lightning + + Nemotron 3 Ultra + + + Serving + model + + + Output: + Compare fixed vs. routed runs + Per-task rewards + Serving models + Input + output tokens + Latency and captured errors + diff --git a/benchmark/nemo_gym/compare.py b/benchmark/nemo_gym/compare.py new file mode 100644 index 000000000..584634779 --- /dev/null +++ b/benchmark/nemo_gym/compare.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compare one runner-produced experiment using Gym-captured usage and model attribution.""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from collections import Counter +from pathlib import Path +from statistics import mean +from typing import Any, cast + + +def require(condition: bool, message: str) -> None: + """Reject incomplete or incompatible evidence before calculating metrics.""" + if not condition: + raise ValueError(message) + + +def number(value: Any, name: str) -> int | float: + """Read a finite, nonnegative measurement without treating missing values as zero.""" + require( + type(value) in (int, float) and math.isfinite(value) and value >= 0, + f"{name} must be a finite, nonnegative number", + ) + return cast(int | float, value) + + +def read_jsonl(path: Path) -> list[dict[str, Any]]: + """Read the JSONL objects written by Gym and the tutorial callback.""" + return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line] + + +def index_rows(path: Path) -> dict[tuple[int, int], dict[str, Any]]: + """Index Gym rows by task and repeat, rejecting duplicate identities.""" + indexed = {} + for row in read_jsonl(path): + key = (row["_ng_task_index"], row["_ng_rollout_index"]) + require(key not in indexed, f"{path}: duplicate task/repeat {key}") + indexed[key] = row + return indexed + + +def has_answer_text(response: dict[str, Any]) -> bool: + """Find answer text without mistaking reasoning or tool calls for the answer.""" + return any( + part["type"] == "output_text" and bool(part["text"].strip()) + for item in response["output"] + if item["type"] == "message" and item["role"] == "assistant" + for part in item["content"] + ) + + +def load_run(path: Path) -> dict[str, Any]: + """Load Gym's expected inputs, completed rollouts, and terminal failure sidecar.""" + failure_path, rollout_path = path / "rollouts_failures.jsonl", path / "rollouts.jsonl" + return { + "path": path, + "inputs": index_rows(path / "rollouts_materialized_inputs.jsonl"), + "rows": index_rows(rollout_path) if rollout_path.exists() else {}, + "failures": read_jsonl(failure_path) if failure_path.exists() else [], + } + + +def token_total(calls: list[dict[str, Any]], field: str) -> int | str: + """Sum reported token counts while explicitly retaining unknown contributions.""" + values = [call.get(field) for call in calls] + known = [value for value in values if value is not None] + require(all(type(value) is int and value >= 0 for value in known), f"Invalid captured {field}") + if not known: + return "unknown" + total = sum(known) + return total if len(known) == len(values) else f"{total}+unknown" + + +def summarize(run: dict[str, Any], route: str) -> tuple[dict[str, int | float | str], Counter[str]]: + """Attribute final answers by response ID and count all Gym-captured calls.""" + model_path = run["path"] / "models.jsonl" + require(model_path.is_file(), f"Missing {model_path}; use a fresh run with the updated runner") + attribution: dict[str, str] = {} + for record in read_jsonl(model_path): + response_id, model = record["response_id"], record["model"] + require(model not in ("fixed", "routed"), f"{route}: invalid model attribution") + require(response_id not in attribution, f"{route}: ambiguous model attribution") + attribution[response_id] = model + + calls, rewards, latencies, models, statuses = [], [], [], [], [] + response_ids: set[str] = set() + for key, row in sorted(run["rows"].items()): + label = f"{route} {key}" + require(not row.get("_ng_failure_class"), f"{label}: failed rollout") + response = row["response"] + require(has_answer_text(response), f"{label}: missing final answer text") + require( + response["status"] in ("completed", "incomplete"), f"{label}: invalid response status" + ) + response_id = response["id"] + require(response_id in attribution, f"{label}: missing final-answer model attribution") + require(response_id not in response_ids, f"{label}: response reused across rollouts") + response_ids.add(response_id) + capture = row["ng_model_call_capture"] + require(not capture.get("gaps"), f"{label}: incomplete model-call capture") + records = capture["calls"] + terminal = [call for call in records if call.get("response_id") == response_id] + require(len(terminal) == 1, f"{label}: missing or ambiguous terminal capture") + call = terminal[0] + require( + call["status_code"] == 200 and not call.get("error_category"), + f"{label}: terminal call failed", + ) + require( + call["response_status"] == response["status"], + f"{label}: response and capture status differ", + ) + require(call["model"] == route, f"{label}: captured response has the wrong model group") + reward = number(row["reward"], f"{label}: reward") + require(reward <= 1, f"{label}: MCQA reward must be between zero and one") + rewards.append(reward) + latencies.append(number(row["ng_perf"]["total_latency_ms"], f"{label}: rollout latency")) + models.append(attribution[response_id]) + statuses.append(response["status"]) + calls.extend(records) + return { + "Paired rollouts": len(rewards), + "Incomplete responses": statuses.count("incomplete"), + "Mean reward": mean(rewards), + "Captured input tokens": token_total(calls, "tokens_in"), + "Captured output tokens": token_total(calls, "tokens_out"), + "Mean rollout latency (ms)": mean(latencies), + "Captured calls": len(calls), + "Captured failed calls": sum( + call.get("status_code") != 200 or bool(call.get("error_category")) for call in calls + ), + "Calls with unknown usage": sum( + call.get("tokens_in") is None or call.get("tokens_out") is None for call in calls + ), + }, Counter(models) + + +def compare(results: Path) -> None: + """Report coverage first, then compare both conditions from one experiment directory.""" + runs = {route: load_run(results / route) for route in ("fixed", "routed")} + complete = True + for name, run in runs.items(): + expected, actual = set(run["inputs"]), set(run["rows"]) + missing, unexpected = expected - actual, actual - expected + print( + f"{name}: expected={len(expected)}, completed={len(actual)}, missing={len(missing)}, unexpected={len(unexpected)}, failures={len(run['failures'])}" + ) + complete &= bool(expected) and not missing and not unexpected and not run["failures"] + fixed_keys, routed_keys = set(runs["fixed"]["rows"]), set(runs["routed"]["rows"]) + print( + f"Pairing: matched={len(fixed_keys & routed_keys)}, fixed-only={len(fixed_keys - routed_keys)}, routed-only={len(routed_keys - fixed_keys)}" + ) + require(complete, "Incomplete runs: inspect the failure artifacts; no averages calculated") + require( + runs["fixed"]["inputs"] == runs["routed"]["inputs"], + "Task inputs, verifier metadata, or generation settings differ", + ) + summaries = {name: summarize(run, name) for name, run in runs.items()} + print(f"\n{'Metric':<29} {'fixed':>14} {'routed':>14}") + for metric in summaries["fixed"][0]: + values = [summaries[name][0][metric] for name in ("fixed", "routed")] + formatted = [f"{value:.3f}" if isinstance(value, float) else str(value) for value in values] + print(f"{metric:<29} {formatted[0]:>14} {formatted[1]:>14}") + for name, (summary, models) in summaries.items(): + print(f"\n{name} serving models: {json.dumps(models, sort_keys=True)}") + if ( + summary["Incomplete responses"] + or summary["Captured failed calls"] + or summary["Calls with unknown usage"] + or summary["Captured calls"] != summary["Paired rollouts"] + ): + print( + f"WARNING: {name} has incomplete responses, recovered errors, additional calls, or unknown usage." + ) + print( + "\nGym-scored incomplete responses remain in the comparison. " + "Gym-captured tokens include extra attempts; +unknown means some usage was not reported." + ) + print( + "Retries inside Gym's model adapter, LiteLLM, or the provider may not appear separately. Tokens are not dollar costs." + ) + print( + "Random makes no classifier calls. A small run demonstrates the integration, not a routing advantage." + ) + + +def main(argv: list[str] | None = None) -> int: + """Run the comparison without importing Gym or Switchyard.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "results", type=Path, help="Experiment directory containing fixed/ and routed/" + ) + args = parser.parse_args(argv) + try: + compare(args.results) + except (OSError, ValueError, KeyError, TypeError) as error: + print(f"Cannot compare: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmark/nemo_gym/gym_routing_plugin.py b/benchmark/nemo_gym/gym_routing_plugin.py new file mode 100644 index 000000000..333853d88 --- /dev/null +++ b/benchmark/nemo_gym/gym_routing_plugin.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize Gym responses and record the serving model before LiteLLM applies its alias.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +from litellm.integrations.custom_logger import CustomLogger + + +def response_payload(response: Any) -> dict[str, Any]: + """Keep missing detail counts unknown while satisfying Gym's object-shaped fields.""" + payload = response.model_dump() + for item in payload.get("output") or []: + if not isinstance(item, dict) or item.get("type") != "reasoning": + continue + if item.get("summary") is None: + item["summary"] = [] + for content in item.get("content") or []: + if isinstance(content, dict) and content.get("type") == "output_text": + content["type"] = "reasoning_text" + usage = payload.get("usage") + if isinstance(usage, dict): + for key, leaf in ( + ("input_tokens_details", "cached_tokens"), + ("output_tokens_details", "reasoning_tokens"), + ): + if usage.get(key) is None: + usage[key] = {leaf: None} + elif isinstance(usage[key], dict): + usage[key].setdefault(leaf, None) + hidden = getattr(response, "_hidden_params", None) + if hidden is not None: + payload["_hidden_params"] = hidden + return payload + + +class GymRoutingPlugin(CustomLogger): + """Provide response compatibility and final-model attribution, not routing or accounting.""" + + def __init__(self, results: Path) -> None: + super().__init__() + self.results = results + + async def async_post_call_success_hook( + self, data: dict[str, Any], user_api_key_dict: Any, response: Any + ) -> dict[str, Any]: + """Save only response identity and the serving deployment; never prompts or credentials.""" + payload = response_payload(response) + route = data.get("model") + metadata = data.get("litellm_metadata") or {} + record = {"response_id": payload.get("id"), "model": metadata.get("deployment")} + if ( + route not in ("fixed", "routed") + or not all(isinstance(value, str) and value.strip() for value in record.values()) + or record["model"] in ("fixed", "routed") + ): + raise ValueError("Missing response ID or serving deployment for Gym model attribution") + directory = self.results / route + directory.mkdir(parents=True, exist_ok=True) + with (directory / "models.jsonl").open("a", encoding="utf-8") as stream: + stream.write(json.dumps(record) + "\n") + return payload + + +PLUGIN = ( + GymRoutingPlugin(Path(os.environ["NEMO_GYM_LITELLM_RESULTS"])) + if os.environ.get("NEMO_GYM_LITELLM_RESULTS") + else None +) diff --git a/benchmark/nemo_gym/litellm.yaml b/benchmark/nemo_gym/litellm.yaml new file mode 100644 index 000000000..6ff745c7d --- /dev/null +++ b/benchmark/nemo_gym/litellm.yaml @@ -0,0 +1,30 @@ +model_list: + - model_name: fixed + litellm_params: &ultra + model: nvidia_nim/nvidia/nemotron-3-ultra-550b-a55b + api_base: os.environ/NVIDIA_BASE_URL + api_key: os.environ/NVIDIA_API_KEY + max_retries: 0 + extra_body: + chat_template_kwargs: + enable_thinking: false + - model_name: routed + litellm_params: *ultra + - model_name: routed + litellm_params: + model: nvidia_nim/nvidia/nemotron-3.5-lightning-30b-a3b + api_base: os.environ/NVIDIA_BASE_URL + api_key: os.environ/NVIDIA_API_KEY + max_retries: 0 + extra_body: + chat_template_kwargs: + enable_thinking: false +router_settings: + num_retries: 0 + plugins: + - switchyard_litellm.configuration.configured_plugin.ROUTING_PLUGIN +litellm_settings: + num_retries: 0 + callbacks: + - switchyard_litellm.configuration.configured_plugin.ROUTING_PLUGIN + - gym_routing_plugin.PLUGIN diff --git a/benchmark/nemo_gym/routes.toml b/benchmark/nemo_gym/routes.toml new file mode 100644 index 000000000..00dcf73a6 --- /dev/null +++ b/benchmark/nemo_gym/routes.toml @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +algorithm = "random" +seed = 6 diff --git a/benchmark/nemo_gym/run.sh b/benchmark/nemo_gym/run.sh new file mode 100644 index 000000000..1beff744b --- /dev/null +++ b/benchmark/nemo_gym/run.sh @@ -0,0 +1,180 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Run fixed-model and Random-routed conditions over the same NeMo Gym rollouts. + +set +x +set -euo pipefail + +# Step 1: Choose settings and check the setup before starting any services. +# Environment variables override these defaults; each run needs a fresh results directory. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SWITCHYARD_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +PROFILE="$SCRIPT_DIR/litellm.yaml" +ROUTING="$SCRIPT_DIR/routes.toml" +RESULTS_DIR="${RESULTS_DIR:-$SCRIPT_DIR/results/$(date -u +%Y%m%dT%H%M%SZ)}" +PORT="${LITELLM_PORT:-4000}" +LIMIT="${LIMIT:-5}" +REPEATS="${REPEATS:-1}" +CONCURRENCY="${CONCURRENCY:-1}" + +if [[ $# -gt 0 ]]; then + cat <<'EOF' +Usage: bash benchmark/nemo_gym/run.sh +Set GYM_DIR to the Gym checkout from the README setup. + +Optional environment variables (defaults): + LIMIT=5, REPEATS=1, CONCURRENCY=1, LITELLM_PORT=4000 + RESULTS_DIR Fresh output directory (results/) + NVIDIA_BASE_URL Provider endpoint (NVIDIA API Catalog /v1) +Provider credentials come from the environment variables named in litellm.yaml. +EOF + [[ $# -eq 1 && ( "$1" == "-h" || "$1" == "--help" ) ]] && exit 0 + exit 2 +fi + +die() { echo "error: $*" >&2; exit 1; } +[[ -n "${GYM_DIR:-}" ]] || die "set GYM_DIR; see the README setup" +GYM_DIR="$(cd "$GYM_DIR" && pwd)" +[[ "$RESULTS_DIR" = /* ]] || RESULTS_DIR="$PWD/$RESULTS_DIR" +[[ -f "$PROFILE" && -f "$ROUTING" ]] || die "LiteLLM profile or routing TOML does not exist" +[[ ! -e "$RESULTS_DIR" && ! -L "$RESULTS_DIR" ]] || die "results path already exists: $RESULTS_DIR" +[[ "$GYM_DIR" =~ ^[a-zA-Z0-9_./-]+$ && "$RESULTS_DIR" =~ ^[a-zA-Z0-9_./-]+$ ]] || die "Gym workspace/output paths must not contain spaces or shell metacharacters" +for value in "$LIMIT" "$REPEATS" "$CONCURRENCY" "$PORT"; do + [[ "$value" =~ ^[1-9][0-9]*$ ]] || die "limits, repeats, concurrency and port must be positive integers" +done +GYM="$GYM_DIR/.venv/bin/gym" +PYTHON="$GYM_DIR/.venv/bin/python" +[[ -x "$GYM" && -x "$PYTHON" ]] || die "complete the README's Gym setup" +export PATH="$GYM_DIR/.venv/bin:$PATH" +for tool in git uv curl; do command -v "$tool" >/dev/null || die "missing required tool: $tool"; done +GYM_REVISION="$(git -C "$GYM_DIR" describe --always --dirty --abbrev=40)" +SWITCHYARD_REVISION="$(git -C "$SWITCHYARD_ROOT" describe --always --dirty --abbrev=40)" +"$PYTHON" - "$PORT" <<'PY' +import socket +import sys + +port = int(sys.argv[1]) +if not 1 <= port <= 65535 or port == 11000: + raise SystemExit('error: choose a LiteLLM port in 1..65535 other than Gym port 11000') +for value in (port, 11000): + try: + with socket.socket() as sock: + sock.bind(('127.0.0.1', value)) + except OSError: + raise SystemExit(f'error: port {value} is occupied; do not share a running Gym/proxy') from None +PY +umask 077 +mkdir -p "$(dirname "$RESULTS_DIR")" +mkdir "$RESULTS_DIR" +RESULTS_DIR="$(cd "$RESULTS_DIR" && pwd)" +BENCHMARK_DATA="$GYM_DIR/benchmarks/mmlu-redux/data/mmlu-redux_benchmark.jsonl" +ROOT_URL="http://127.0.0.1:$PORT" +GYM_PID="" +PROXY_PID="" + +# Track the processes started by this script so exits and interrupts can stop them. +stop_process() { + local pid="$1" + [[ -n "$pid" ]] || return 0 + if kill -0 "$pid" 2>/dev/null; then + kill -INT "$pid" 2>/dev/null || true + for _ in {1..100}; do kill -0 "$pid" 2>/dev/null || break; sleep 0.1; done + if kill -0 "$pid" 2>/dev/null; then + echo "warning: escalating shutdown of run-owned process $pid" >&2 + kill -TERM "$pid" 2>/dev/null || true + sleep 1 + kill -KILL "$pid" 2>/dev/null || true + fi + fi + wait "$pid" 2>/dev/null || true +} +cleanup() { + local status=$? + trap - EXIT + trap '' INT TERM + stop_process "$GYM_PID" + stop_process "$PROXY_PID" + exit "$status" +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM +run_gym() { + local log_file="$1" status=0 + shift + (cd "$GYM_DIR"; exec "$PYTHON" -c 'import os, signal, sys; signal.signal(signal.SIGINT, signal.SIG_DFL); os.execv(sys.argv[1], sys.argv[1:])' "$GYM" "$@") >"$log_file" 2>&1 & + GYM_PID=$! + wait "$GYM_PID" || status=$? + GYM_PID="" + return "$status" +} + +# Step 2: Prepare MMLU-Redux, reusing the prepared dataset if it already exists. +echo "Results: $RESULTS_DIR" +echo "Preparing up to $LIMIT tasks; repeats: $REPEATS; concurrency: $CONCURRENCY" +if [[ ! -s "$BENCHMARK_DATA" ]]; then + run_gym "$RESULTS_DIR/prepare.log" eval prepare --benchmark mmlu-redux "hydra.run.dir=$RESULTS_DIR/hydra-prepare" || die "preparation failed; see $RESULTS_DIR/prepare.log" +fi +[[ -s "$BENCHMARK_DATA" ]] || die "MMLU-Redux preparation produced no data" + +# Step 3: Start a local LiteLLM server and wait until it is ready. +# litellm.yaml defines the model groups; routes.toml configures Switchyard Random routing. +# Switchyard runs as a library inside LiteLLM, not as a separate server. +export NVIDIA_BASE_URL="${NVIDIA_BASE_URL:-https://integrate.api.nvidia.com/v1}" +export SWITCHYARD_LITELLM_CONFIG="$ROUTING" +export NEMO_GYM_LITELLM_RESULTS="$RESULTS_DIR" +echo "Starting LiteLLM at $ROOT_URL; see $RESULTS_DIR/litellm.log" +PYTHONPATH="$SCRIPT_DIR:$SWITCHYARD_ROOT/examples/litellm/src${PYTHONPATH:+:$PYTHONPATH}" \ + uv run --project "$SWITCHYARD_ROOT/examples/litellm" --locked \ + --with 'litellm[proxy]==1.97.0' --with 'fastapi==0.136.3' --with 'starlette==1.3.1' \ + litellm --config "$PROFILE" --host 127.0.0.1 --port "$PORT" --num_workers 1 \ + >"$RESULTS_DIR/litellm.log" 2>&1 & +PROXY_PID=$! +ready=false +for _ in {1..240}; do + kill -0 "$PROXY_PID" 2>/dev/null || die "LiteLLM exited; see $RESULTS_DIR/litellm.log" + if curl -fsS --max-time 2 "$ROOT_URL/health/readiness/details" >/dev/null 2>&1; then + ready=true + break + fi + sleep 1 +done +[[ "$ready" == true ]] || die "LiteLLM readiness timed out; see $RESULTS_DIR/litellm.log" + +# Step 4: Save the checkout versions once for reference and prepare the output folders. +# Pairing uses Gym's materialized inputs, not configuration fingerprints. +printf 'Gym: %s\nSwitchyard: %s\n' "$GYM_REVISION" "$SWITCHYARD_REVISION" >"$RESULTS_DIR/versions.txt" +mkdir -p "$RESULTS_DIR"/{fixed,routed}/model-calls + +# Step 5: Evaluate the fixed baseline, then Random routing, on the same tasks. +# Gym's litellm_model adapter calls the proxy; fixed and routed name its model groups. +# Only the group and output paths change; the evaluation settings stay the same. +# Both conditions use one dedicated LiteLLM instance and separate Gym capture paths. +for route in fixed routed; do + run_dir="$RESULTS_DIR/$route" + echo "Gym evaluation: gym eval run --benchmark mmlu-redux --model-type litellm_model --model $route" + echo "Running $route; see $run_dir/gym.log" + run_gym "$run_dir/gym.log" eval run \ + --benchmark mmlu-redux --model-type litellm_model --model "$route" \ + --output "$run_dir/rollouts.jsonl" --split benchmark \ + --limit "$LIMIT" --num-repeats "$REPEATS" --concurrency "$CONCURRENCY" \ + --temperature 0 --max-output-tokens 8192 \ + "++policy_base_url=$ROOT_URL/v1" ++policy_api_key=unused \ + ++route_failures_to_sidecar=true ++observability_enabled=true \ + "++model_call_capture_dir=$run_dir/model-calls" \ + "++nemo_gym_log_dir=$run_dir/server-logs" \ + ++mcqa_simple_agent.responses_api_agents.simple_agent.max_steps=1 \ + "hydra.run.dir=$run_dir/hydra" || die "$route evaluation failed; see $run_dir/gym.log" + # Stop before the next condition if Gym reports terminal rollout failures. + [[ ! -s "$run_dir/rollouts_failures.jsonl" ]] || die "$route has terminal rollout failures; see $run_dir/rollouts_failures.jsonl" +done + +# Step 6: Stop LiteLLM, then compare the saved fixed and routed runs. +# compare.py checks complete, paired evidence before reporting rewards, tokens, and latency. +stop_process "$PROXY_PID" +PROXY_PID="" +echo "Comparing fixed and routed results" +"$PYTHON" "$SCRIPT_DIR/compare.py" "$RESULTS_DIR" 2>&1 | tee "$RESULTS_DIR/comparison.txt" +echo "Comparison written to $RESULTS_DIR/comparison.txt" diff --git a/tests/test_nemo_gym_compare.py b/tests/test_nemo_gym_compare.py new file mode 100644 index 000000000..26b770c6a --- /dev/null +++ b/tests/test_nemo_gym_compare.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib.util +import json +from copy import deepcopy +from pathlib import Path +from types import ModuleType +from typing import Any + +import pytest + +BIG = "nvidia_nim/nvidia/nemotron-3-ultra-550b-a55b" +SMALL = "nvidia_nim/nvidia/nemotron-3.5-lightning-30b-a3b" + + +@pytest.fixture +def comparator() -> ModuleType: + """Load the offline comparator without the model-serving dependencies.""" + path = Path(__file__).resolve().parents[1] / "benchmark/nemo_gym/compare.py" + spec = importlib.util.spec_from_file_location("switchyard_nemo_gym_compare", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def artifacts() -> dict[str, dict[str, Any]]: + """Build paired runs with one correct and one wrong answer per condition.""" + runs = {} + for route in ("fixed", "routed"): + inputs, rows, models = [], [], [] + for index in range(2): + task = { + "_ng_task_index": index, + "_ng_rollout_index": 0, + "expected_answer": "B", + "responses_create_params": { + "input": [{"role": "user", "content": f"Question {index}"}], + "temperature": 0, + "max_output_tokens": 8192, + }, + } + response_id = f"{route}-{index}" + inputs.append(task) + rows.append( + { + **deepcopy(task), + "reward": 1 - index, + "response": { + "id": response_id, + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "\\boxed{B}" if index == 0 else "\\boxed{A}", + } + ], + } + ], + }, + "ng_perf": {"total_latency_ms": 100 + 200 * index}, + "ng_model_call_capture": { + "calls": [ + { + "model_call_id": response_id, + "response_id": response_id, + "model": route, + "status_code": 200, + "error_category": None, + "response_status": "completed", + "tokens_in": 10 + 10 * index, + "tokens_out": 20, + } + ], + }, + } + ) + models.append( + { + "response_id": response_id, + "model": SMALL if route == "routed" and index == 0 else BIG, + } + ) + runs[route] = { + "rollouts_materialized_inputs.jsonl": inputs, + "rollouts.jsonl": rows, + "rollouts_failures.jsonl": [], + "models.jsonl": models, + } + return runs + + +def _write_runs(tmp_path: Path, artifacts: dict[str, dict[str, Any]]) -> list[str]: + for route, files in artifacts.items(): + directory = tmp_path / route + directory.mkdir() + for name, rows in files.items(): + (directory / name).write_text( + "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" + ) + return [str(tmp_path)] + + +def _assert_metric(output: str, metric: str, values: list[str]) -> None: + line = next(line for line in output.splitlines() if line[:29].rstrip() == metric) + assert line[29:].split() == values + + +def test_complete_reordered_pair( + tmp_path: Path, + comparator: ModuleType, + artifacts: dict[str, dict[str, Any]], + capsys: pytest.CaptureFixture[str], +) -> None: + """Pair by identity, not file order, and report Gym usage and serving models.""" + for rows in artifacts["routed"].values(): + rows.reverse() + assert comparator.main(_write_runs(tmp_path, artifacts)) == 0 + output = capsys.readouterr() + assert output.err == "" + assert "Pairing: matched=2, fixed-only=0, routed-only=0" in output.out + assert f'fixed serving models: {{"{BIG}": 2}}' in output.out + assert f"routed serving models: {json.dumps({SMALL: 1, BIG: 1}, sort_keys=True)}" in output.out + for metric, values in { + "Paired rollouts": ["2", "2"], + "Incomplete responses": ["0", "0"], + "Mean reward": ["0.500", "0.500"], + "Captured input tokens": ["30", "30"], + "Captured output tokens": ["40", "40"], + "Mean rollout latency (ms)": ["200", "200"], + "Captured calls": ["2", "2"], + "Captured failed calls": ["0", "0"], + "Calls with unknown usage": ["0", "0"], + }.items(): + _assert_metric(output.out, metric, values) + assert "WARNING" not in output.out + assert "Gym-captured" in output.out + assert "Gateway-reported" not in output.out + + +@pytest.mark.parametrize( + ("problem", "expected_error"), + [ + ("incomplete", "Incomplete runs"), + ("mismatched", "Task inputs, verifier metadata, or generation settings differ"), + ("capture", "incomplete model-call capture"), + ], +) +def test_invalid_evidence_never_prints_averages( + tmp_path: Path, + comparator: ModuleType, + artifacts: dict[str, dict[str, Any]], + capsys: pytest.CaptureFixture[str], + problem: str, + expected_error: str, +) -> None: + """Reject misleading comparisons with a useful diagnostic, not partial averages.""" + routed = artifacts["routed"] + if problem == "incomplete": + for run in artifacts.values(): + run["rollouts.jsonl"].pop() + elif problem == "mismatched": + routed["rollouts_materialized_inputs.jsonl"][0]["expected_answer"] = "A" + else: + routed["rollouts.jsonl"][0]["ng_model_call_capture"]["gaps"] = ["missing exchange"] + assert comparator.main(_write_runs(tmp_path, artifacts)) == 1 + output = capsys.readouterr() + assert "Cannot compare:" in output.err + assert expected_error in output.err + assert "Mean reward" not in output.out + + +def test_incomplete_model_response_is_reported( + tmp_path: Path, + comparator: ModuleType, + artifacts: dict[str, dict[str, Any]], + capsys: pytest.CaptureFixture[str], +) -> None: + """Keep a Gym-scored truncation in the comparison and make it visible.""" + row = artifacts["routed"]["rollouts.jsonl"][0] + row["response"]["status"] = "incomplete" + row["reward"] = 0 + row["ng_model_call_capture"]["calls"][0]["response_status"] = "incomplete" + + assert comparator.main(_write_runs(tmp_path, artifacts)) == 0 + output = capsys.readouterr().out + _assert_metric(output, "Incomplete responses", ["0", "1"]) + assert "WARNING: routed" in output + assert "WARNING: fixed" not in output + + +def test_recovery_keeps_extra_work_and_terminal_attribution( + tmp_path: Path, + comparator: ModuleType, + artifacts: dict[str, dict[str, Any]], + capsys: pytest.CaptureFixture[str], +) -> None: + """Count extra captured work without treating unreported failed usage as zero.""" + calls = artifacts["routed"]["rollouts.jsonl"][0]["ng_model_call_capture"]["calls"] + calls.insert( + 0, + { + "model_call_id": "failed", + "status_code": 503, + "error_category": "upstream", + "tokens_in": None, + "tokens_out": None, + }, + ) + calls.append( + { + **calls[1], + "model_call_id": "extra", + "response_id": "superseded", + "tokens_in": 7, + "tokens_out": 13, + } + ) + artifacts["routed"]["models.jsonl"].append({"response_id": "superseded", "model": BIG}) + assert comparator.main(_write_runs(tmp_path, artifacts)) == 0 + output = capsys.readouterr().out + for metric, values in { + "Captured input tokens": ["30", "37+unknown"], + "Captured output tokens": ["40", "53+unknown"], + "Captured calls": ["2", "4"], + "Captured failed calls": ["0", "1"], + "Calls with unknown usage": ["0", "1"], + }.items(): + _assert_metric(output, metric, values) + assert f"routed serving models: {json.dumps({SMALL: 1, BIG: 1}, sort_keys=True)}" in output + assert "WARNING: routed" in output + assert "WARNING: fixed" not in output diff --git a/tests/test_nemo_gym_litellm.py b/tests/test_nemo_gym_litellm.py new file mode 100644 index 000000000..ffd718fb8 --- /dev/null +++ b/tests/test_nemo_gym_litellm.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from types import ModuleType +from typing import Any + +import pytest + +litellm = pytest.importorskip("litellm") + +from litellm import ResponsesAPIResponse # noqa: E402 + +BIG = "nvidia_nim/nvidia/nemotron-3-ultra-550b-a55b" +SMALL = "nvidia_nim/nvidia/nemotron-3.5-lightning-30b-a3b" +ROOT = Path(__file__).resolve().parents[1] +CALLBACK = ROOT / "benchmark/nemo_gym/gym_routing_plugin.py" + + +@pytest.fixture +def callback(monkeypatch: pytest.MonkeyPatch) -> ModuleType: + monkeypatch.delenv("NEMO_GYM_LITELLM_RESULTS", raising=False) + spec = importlib.util.spec_from_file_location("switchyard_gym_routing_plugin_test", CALLBACK) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _response( + *, + response_id: str = "response-1", + status: str = "incomplete", + usage: dict[str, Any] | None = None, +) -> ResponsesAPIResponse: + response = ResponsesAPIResponse( + id=response_id, + created_at=1700000000, + object="response", + model=SMALL, + output=[], + status=status, + incomplete_details={"reason": "max_output_tokens"} if status == "incomplete" else None, + usage=usage, + ) + response._hidden_params = {"model_id": "deployment-1", "custom_llm_provider": "nvidia_nim"} + return response + + +def test_response_payload_preserves_wire_identity_and_unknown_usage_details( + callback: ModuleType, +) -> None: + response = _response( + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "input_tokens_details": None, + "output_tokens_details": None, + } + ) + + payload = callback.response_payload(response) + + assert payload["id"] == "response-1" + assert payload["status"] == "incomplete" + assert payload["incomplete_details"] == {"reason": "max_output_tokens"} + assert payload["usage"]["input_tokens"] == 10 + assert payload["usage"]["output_tokens"] == 5 + assert payload["usage"]["total_tokens"] == 15 + assert payload["usage"]["input_tokens_details"] == {"cached_tokens": None} + assert payload["usage"]["output_tokens_details"] == {"reasoning_tokens": None} + assert payload["_hidden_params"] == response._hidden_params + + reported = _response( + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "input_tokens_details": {"cached_tokens": 3}, + "output_tokens_details": {"reasoning_tokens": 4}, + } + ) + reported_payload = callback.response_payload(reported) + assert reported_payload["usage"]["input_tokens_details"]["cached_tokens"] == 3 + assert reported_payload["usage"]["output_tokens_details"]["reasoning_tokens"] == 4 + assert callback.response_payload(_response(usage=None))["usage"] is None + + +@pytest.mark.parametrize("finish_reason", ["stop", "length"]) +def test_response_payload_normalizes_litellm_reasoning( + callback: ModuleType, finish_reason: str +) -> None: + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + from openai.types.responses.response_reasoning_item import ResponseReasoningItem + + response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Synthetic schema test", + responses_api_request={}, + chat_completion_response={ + "id": "chatcmpl-reasoning-test", + "created": 1700000000, + "object": "chat.completion", + "model": SMALL, + "choices": [ + { + "index": 0, + "finish_reason": finish_reason, + "message": { + "role": "assistant", + "reasoning_content": "Synthetic reasoning for schema compatibility.", + "content": "\\boxed{B}", + }, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + before = response.model_dump() + payload = callback.response_payload(response) + reasoning, answer = payload["output"] + validated = ResponseReasoningItem.model_validate(reasoning) + + assert validated.summary == [] + assert validated.content is not None + assert validated.content[0].type == "reasoning_text" + assert validated.content[0].text == "Synthetic reasoning for schema compatibility." + assert reasoning["id"] == before["output"][0]["id"] + assert reasoning["status"] == before["output"][0]["status"] + assert answer == before["output"][1] + assert answer["content"][0]["type"] == "output_text" + assert answer["content"][0]["text"] == "\\boxed{B}" + assert payload["id"] == before["id"] + assert payload["status"] == ("completed" if finish_reason == "stop" else "incomplete") + assert payload.get("incomplete_details") == before.get("incomplete_details") + for key in ("input_tokens", "output_tokens", "total_tokens"): + assert payload["usage"][key] == before["usage"][key] + assert response.model_dump() == before + + +async def test_callback_records_serving_model_without_request_content( + callback: ModuleType, tmp_path: Path +) -> None: + plugin = callback.GymRoutingPlugin(tmp_path) + response = _response( + status="completed", usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15} + ) + data = { + "model": "routed", + "messages": [{"role": "user", "content": "secret-prompt-marker"}], + "headers": {"Authorization": "secret-header-marker"}, + "litellm_metadata": { + "deployment": SMALL, + "routing_plugin_signals": {"switchyard": {"selected_model_id": BIG}}, + }, + } + payload = await plugin.async_post_call_success_hook(data, None, response) + records = [ + json.loads(line) for line in (tmp_path / "routed/models.jsonl").read_text().splitlines() + ] + assert records == [{"response_id": "response-1", "model": SMALL}] + assert payload == callback.response_payload(response) + assert sorted(path.name for path in (tmp_path / "routed").iterdir()) == ["models.jsonl"] + + +async def test_callback_rejects_missing_serving_model(callback: ModuleType, tmp_path: Path) -> None: + plugin = callback.GymRoutingPlugin(tmp_path) + with pytest.raises(ValueError, match="serving deployment"): + await plugin.async_post_call_success_hook( + {"model": "routed"}, None, _response(status="completed") + ) + assert not list(tmp_path.iterdir()) diff --git a/tests/test_nemo_gym_run.py b/tests/test_nemo_gym_run.py new file mode 100644 index 000000000..49c240887 --- /dev/null +++ b/tests/test_nemo_gym_run.py @@ -0,0 +1,325 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import signal +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +RUNNER = ROOT / "benchmark/nemo_gym/run.sh" + + +def _write_executable(path: Path, content: str) -> None: + path.write_text(content, encoding="utf-8") + path.chmod(0o755) + + +@pytest.fixture +def fake_runner_env(tmp_path: Path) -> tuple[dict[str, str], Path, Path]: + """Use local tool doubles and an event log to check subprocess lifecycle ordering.""" + gym_dir = tmp_path / "Gym" + bin_dir = tmp_path / "bin" + results_dir = tmp_path / "results" + event_log = tmp_path / "events.log" + gym_bin = gym_dir / ".venv/bin" + gym_bin.mkdir(parents=True) + bin_dir.mkdir() + + _write_executable( + gym_bin / "python", + """#!/bin/bash +if [[ "${1:-}" == "$FAKE_COMPARE_PATH" ]]; then + printf 'compare:%s\\n' "$*" >> "$EVENT_LOG" + printf 'fixture comparison\\n' + exit 0 +fi +exec "$REAL_PYTHON" "$@" +""", + ) + _write_executable( + gym_bin / "gym", + """#!/bin/bash +set -u +if [[ "${1:-}" == "eval" && "${2:-}" == "prepare" ]]; then + printf 'prepare:%s\\n' "$*" >> "$EVENT_LOG" + mkdir -p "$(dirname "$FAKE_BENCHMARK_DATA")" + printf '{"task": 1}\\n' > "$FAKE_BENCHMARK_DATA" + exit 0 +fi +route="" +output="" +previous="" +for argument in "$@"; do + if [[ "$previous" == "--model" ]]; then route="$argument"; fi + if [[ "$previous" == "--output" ]]; then output="$argument"; fi + previous="$argument" +done +if [[ "${FAKE_GYM_BEHAVIOR:-}" == "wait" ]]; then + exec "$REAL_PYTHON" -c ' +import os +import signal +import sys + +route = sys.argv[1] +arguments = " ".join(sys.argv[2:]) + +def stop(number, _frame): + with open(os.environ["EVENT_LOG"], "a", encoding="utf-8") as stream: + stream.write(f"gym_stop:{route}:{os.getpid()}\\n") + raise SystemExit(128 + number) + +signal.signal(signal.SIGINT, stop) +signal.signal(signal.SIGTERM, stop) +with open(os.environ["EVENT_LOG"], "a", encoding="utf-8") as stream: + stream.write(f"gym_start:{route}:{os.getpid()}:{arguments}\\n") +signal.pause() +' "$route" "$@" +fi +printf 'gym_start:%s:%s:%s\\n' "$route" "$$" "$*" >> "$EVENT_LOG" +run_dir="$(dirname "$output")" +mkdir -p "$run_dir" +printf '{"route": "%s"}\\n' "$route" > "$output" +printf 'gym_done:%s\\n' "$route" >> "$EVENT_LOG" +if [[ "$route" == "fixed" && "${FAKE_GYM_BEHAVIOR:-}" == "exit7" ]]; then exit 7; fi +if [[ "$route" == "fixed" && "${FAKE_GYM_BEHAVIOR:-}" == "sidecar" ]]; then + printf '{"failure": true}\\n' > "$run_dir/rollouts_failures.jsonl" +fi +""", + ) + _write_executable( + bin_dir / "uv", + """#!/bin/bash +exec "$REAL_PYTHON" -c ' +import os +import signal +from pathlib import Path + +def stop(_number, _frame): + with open(os.environ["EVENT_LOG"], "a", encoding="utf-8") as stream: + stream.write(f"proxy_stop:{os.getpid()}\\n") + raise SystemExit(0) + +signal.signal(signal.SIGINT, stop) +signal.signal(signal.SIGTERM, stop) +with open(os.environ["EVENT_LOG"], "a", encoding="utf-8") as stream: + stream.write(f"proxy_start:{os.getpid()}\\n") +Path(os.environ["NEMO_GYM_LITELLM_RESULTS"], "proxy-ready").touch() +signal.pause() +' +""", + ) + _write_executable( + bin_dir / "curl", + """#!/bin/bash +[[ -f "$NEMO_GYM_LITELLM_RESULTS/proxy-ready" ]] +""", + ) + _write_executable( + bin_dir / "git", + """#!/bin/bash +printf '%040d-dirty\\n' 2 +""", + ) + env = os.environ.copy() + for key in ( + "NVIDIA_API_KEY", + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "OPENROUTER_API_KEY", + "NVIDIA_BASE_URL", + "LITELLM_PORT", + ): + env.pop(key, None) + env.update( + { + "PATH": f"{bin_dir}{os.pathsep}{env['PATH']}", + "GYM_DIR": str(gym_dir), + "RESULTS_DIR": str(results_dir), + "EVENT_LOG": str(event_log), + "FAKE_BENCHMARK_DATA": str( + gym_dir / "benchmarks/mmlu-redux/data/mmlu-redux_benchmark.jsonl" + ), + "FAKE_COMPARE_PATH": str(ROOT / "benchmark/nemo_gym/compare.py"), + "REAL_PYTHON": sys.executable, + } + ) + return env, results_dir, event_log + + +def _run(env: dict[str, str]) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["/bin/bash", str(RUNNER)], + cwd=ROOT, + env=env, + text=True, + capture_output=True, + timeout=30, + check=False, + ) + + +def _events(path: Path) -> list[str]: + return path.read_text(encoding="utf-8").splitlines() if path.exists() else [] + + +def _pid_is_alive(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + return True + + +def test_help_and_shell_syntax_need_no_setup(tmp_path: Path) -> None: + before = list(tmp_path.iterdir()) + result = subprocess.run( + ["/bin/bash", str(RUNNER), "--help"], + cwd=tmp_path, + env={"PATH": os.defpath}, + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 0 + assert result.stderr == "" + for name in ( + "GYM_DIR", + "NVIDIA_BASE_URL", + "LITELLM_PORT", + "LIMIT", + "REPEATS", + "CONCURRENCY", + ): + assert name in result.stdout + assert list(tmp_path.iterdir()) == before + assert subprocess.run(["/bin/bash", "-n", str(RUNNER)], check=False).returncode == 0 + + +def test_existing_results_are_not_overwritten( + fake_runner_env: tuple[dict[str, str], Path, Path], +) -> None: + env, results_dir, event_log = fake_runner_env + results_dir.mkdir() + sentinel = results_dir / "sentinel" + sentinel.write_text("keep", encoding="utf-8") + result = _run(env) + assert result.returncode != 0 + assert _events(event_log) == [] + assert sentinel.read_text(encoding="utf-8") == "keep" + + +def test_success_uses_one_stock_proxy_for_fixed_and_routed( + fake_runner_env: tuple[dict[str, str], Path, Path], +) -> None: + env, results_dir, event_log = fake_runner_env + env.update({"LIMIT": "2", "REPEATS": "3", "CONCURRENCY": "1"}) + result = _run(env) + assert result.returncode == 0, result.stderr + events = _events(event_log) + assert sum(line.startswith("prepare:") for line in events) == 1 + starts = [line for line in events if line.startswith("proxy_start:")] + assert len(starts) == 1 + proxy_pid = int(starts[0].split(":", 2)[1]) + gym_runs = [line for line in events if line.startswith("gym_start:")] + assert [line.split(":", 2)[1] for line in gym_runs] == ["fixed", "routed"] + for invocation in gym_runs: + for flag in ( + "--benchmark mmlu-redux", + "--model-type litellm_model", + "--max-output-tokens 8192", + "--limit 2", + "--num-repeats 3", + "--concurrency 1", + ): + assert flag in invocation + comparison = next(i for i, line in enumerate(events) if line.startswith("compare:")) + stop = next(i for i, line in enumerate(events) if line.startswith("proxy_stop:")) + assert all(events.index(f"gym_done:{route}") < stop for route in ("fixed", "routed")) + assert stop < comparison + assert not _pid_is_alive(proxy_pid) + assert (results_dir / "versions.txt").is_file() + assert (results_dir / "comparison.txt").read_text(encoding="utf-8") == "fixture comparison\n" + + +@pytest.mark.parametrize("behavior", ["exit7", "sidecar"]) +def test_fixed_failure_stops_before_routed_and_keeps_logs( + fake_runner_env: tuple[dict[str, str], Path, Path], behavior: str +) -> None: + env, results_dir, event_log = fake_runner_env + env["FAKE_GYM_BEHAVIOR"] = behavior + result = _run(env) + events = _events(event_log) + assert result.returncode != 0 + assert any(line.startswith("gym_start:fixed:") for line in events) + assert not any(line.startswith("gym_start:routed:") for line in events) + assert not any(line.startswith("compare:") for line in events) + assert sum(line.startswith("proxy_stop:") for line in events) == 1 + assert (results_dir / "fixed/gym.log").exists() + assert (results_dir / "litellm.log").exists() + if behavior == "sidecar": + assert (results_dir / "fixed/rollouts_failures.jsonl").stat().st_size > 0 + + +@pytest.mark.parametrize( + ("signal_number", "expected_status"), + [(signal.SIGINT, 130), (signal.SIGTERM, 143)], +) +def test_interrupt_stops_only_owned_children( + fake_runner_env: tuple[dict[str, str], Path, Path], + signal_number: signal.Signals, + expected_status: int, +) -> None: + env, _, event_log = fake_runner_env + env["FAKE_GYM_BEHAVIOR"] = "wait" + sentinel = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"]) + runner = subprocess.Popen( + ["/bin/bash", str(RUNNER)], + cwd=ROOT, + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + start_new_session=True, + ) + owned_pids: list[int] = [] + try: + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + starts = [ + line + for line in _events(event_log) + if line.startswith("proxy_start:") or line.startswith("gym_start:fixed:") + ] + if len(starts) == 2: + break + time.sleep(0.05) + else: + pytest.fail("fake Gym did not enter its wait state") + owned_pids = [ + int(line.split(":", 2)[1] if line.startswith("proxy_start:") else line.split(":", 3)[2]) + for line in starts + ] + runner.send_signal(signal_number) + stdout, stderr = runner.communicate(timeout=20) + assert runner.returncode == expected_status, (stdout, stderr) + events = _events(event_log) + assert any(line.startswith("gym_stop:fixed:") for line in events) + assert sum(line.startswith("proxy_stop:") for line in events) == 1 + assert all(not _pid_is_alive(pid) for pid in owned_pids) + assert sentinel.poll() is None + finally: + if runner.poll() is None: + os.killpg(runner.pid, signal.SIGKILL) + runner.wait(timeout=5) + for pid in owned_pids: + if _pid_is_alive(pid): + os.kill(pid, signal.SIGKILL) + sentinel.terminate() + sentinel.wait(timeout=5)