From d85d0719e1a31a48305bf3abb5fbeeebf398d62e Mon Sep 17 00:00:00 2001 From: Shashank Verma Date: Thu, 10 Sep 2026 13:32:11 -0700 Subject: [PATCH 1/6] feat(benchmark): add NeMo Gym fixed-vs-routed evaluation tutorial Signed-off-by: Shashank Verma --- benchmark/README.md | 3 + benchmark/nemo_gym/README.md | 250 ++++++++++++++++++ benchmark/nemo_gym/architecture.svg | 86 ++++++ benchmark/nemo_gym/compare.py | 265 +++++++++++++++++++ benchmark/nemo_gym/routes.toml | 35 +++ tests/test_nemo_gym_compare.py | 393 ++++++++++++++++++++++++++++ 6 files changed, 1032 insertions(+) create mode 100644 benchmark/nemo_gym/README.md create mode 100644 benchmark/nemo_gym/architecture.svg create mode 100644 benchmark/nemo_gym/compare.py create mode 100644 benchmark/nemo_gym/routes.toml create mode 100644 tests/test_nemo_gym_compare.py diff --git a/benchmark/README.md b/benchmark/README.md index dfaa36a52..6440362b7 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 example using Gym's included questions 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..3794d0390 --- /dev/null +++ b/benchmark/nemo_gym/README.md @@ -0,0 +1,250 @@ +# Evaluate Switchyard routing with NeMo Gym + +[NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) is a library for evaluating models and agents using tasks with verifiable outcomes. +This tutorial uses its five included multiple-choice examples to compare a fixed model with Switchyard routing. + +![Gym evaluation with a hosted Switchyard model server](architecture.svg) + +## What changes between the runs? + +Both routes are defined in [routes.toml](routes.toml): + +| Route | Behavior | +|---|---| +| `fixed` | Always use Nemotron 3 Super. | +| `routed` | Ask GPT-OSS 20B to classify the task, then use GPT-OSS 20B or Nemotron 3 Super. | + +The dataset, agent, verifier, temperature, and answer-token limit stay the same. +The classifier is an extra model call: its tokens count even when the router selects GPT-OSS 20B. +Super uses `enable_thinking=false`; GPT-OSS 20B uses `reasoning_effort=low`, including for +classification. These per-model settings stay unchanged between conditions. This is not a +benchmark of either model's maximum reasoning capability. + +## 1. Install the pinned Gym checkout + +You need Git, [uv](https://docs.astral.sh/uv/), two Bash terminals, and an NVIDIA API key for +the public endpoint `https://integrate.api.nvidia.com/v1`. Use these exact model IDs: +`openai/gpt-oss-20b` and `nvidia/nemotron-3-super-120b-a12b`. +Inference may consume credits. The classifier requires strict JSON Schema responses from +its selected model. + +**Smoke-tested:** A one-question paired run completed on this endpoint with correct answers +in both conditions, GPT-OSS 20B serving the routed answer, and no reported model or classifier +errors. The full five-question comparison below has not yet been validated for this pair. + +Run from the **Switchyard repository root**: + +```bash +WORK="$PWD/scratch/nemo-gym-tutorial" +mkdir -p "$WORK" && +git clone https://github.com/NVIDIA-NeMo/Gym.git "$WORK/Gym" && +git -C "$WORK/Gym" checkout 3a26c35fa90c243427378569511f7b06f503e0fd && +uv tool run --from uv==0.11.29 uv venv --python 3.13.14 "$WORK/.venv" && +uv tool run --from uv==0.11.29 uv pip install \ + --python "$WORK/.venv/bin/python" uv==0.11.29 -e "$WORK/Gym" +``` + +The pinned uv can download Python 3.13.14 even if your existing uv is older. It is also installed +inside the tutorial environment for Gym's component setup; your global uv is left unchanged. + +This keeps the checkout and environment under Switchyard's ignored `scratch/` directory, +without changing another Gym checkout. Use a fresh directory for this one-time setup. +The editable install lets Gym's component environments use the same pinned source. + +Gym installs **`nemo-switchyard==0.2.0`** into its model-server environment and hosts the native +proxy in-process. You do **not** need Docker or a separately running `switchyard-server`. +The Switchyard checkout contains this tutorial; its current `main` is **not** the proxy being +executed. Do not copy newer routing options into this version-pinned example. + +## 2. Start the fixed condition — Terminal 1 + +Keep this terminal at the Switchyard repository root. Set your API key in the terminal, +replacing the placeholder with your key: + +```bash +export NVIDIA_API_KEY='' +``` + +Start Gym's resources, agent, and model servers. Run only one Gym environment at a time. +The first start also installs their dependencies. + +```bash +EXAMPLE="$PWD/benchmark/nemo_gym" +WORK="$PWD/scratch/nemo-gym-tutorial" +source "$WORK/.venv/bin/activate" +RUN_DIR="$EXAMPLE/results/first-run" +OUT="$RUN_DIR/fixed" +mkdir -p "$RUN_DIR" + +mkdir "$OUT" && +git -C "$WORK/Gym" rev-parse HEAD > "$OUT/gym-commit.txt" && +gym env start --resources-server mcqa --model-type switchyard_model --model fixed \ + "++policy_model.responses_api_models.switchyard_model.deployment=$EXAMPLE/routes.toml" \ + ++policy_model.responses_api_models.switchyard_model.switchyard_base_url=null \ + "++policy_model.responses_api_models.switchyard_model.condition_dir=$OUT" \ + ++mcqa_simple_agent.responses_api_agents.simple_agent.max_steps=1 \ + ++observability_enabled=true \ + "++model_call_capture_dir=$OUT/model-calls" \ + "++nemo_gym_log_dir=$OUT/server-logs" \ + "hydra.run.dir=$OUT/hydra-start" +``` + +Wait for **`All 3 / 3 servers ready!`**. Leave this terminal running. +`mkdir "$OUT"` deliberately refuses to reuse an existing condition directory. +For a new comparison, change `RUN_DIR` to the same fresh path in **both terminals**. + +## 3. Run the five questions — Terminal 2 + +Open another Bash terminal at the **same Switchyard repository root**: + +```bash +EXAMPLE="$PWD/benchmark/nemo_gym" +WORK="$PWD/scratch/nemo-gym-tutorial" +source "$WORK/.venv/bin/activate" +RUN_DIR="$EXAMPLE/results/first-run" +OUT="$RUN_DIR/fixed" + +gym eval run --no-serve --agent mcqa_simple_agent \ + --input "$WORK/Gym/resources_servers/mcqa/data/example.jsonl" \ + --output "$OUT/rollouts.jsonl" \ + --limit 5 --num-repeats 1 --concurrency 1 \ + --temperature 0 --max-output-tokens 4096 \ + ++route_failures_to_sidecar=true \ + ++observability_enabled=true \ + "++model_call_capture_dir=$OUT/model-calls" \ + "hydra.run.dir=$OUT/hydra-eval" +``` + +These questions are included with Gym. The MCQA resources server checks the answer letter +against the expected answer; it does not call an LLM judge. A wrong answer earns zero reward. +An infrastructure failure is a different outcome, recorded separately. + +The two-terminal flow is intentional: Gym's one-command evaluation mode does not accept +`--split example`. `--no-serve` collects against the servers you already started. + +**After collection finishes, press Ctrl-C in Terminal 1 and wait for shutdown to finish.** +Gym writes `switchyard-stats.json` during shutdown, before stopping its hosted proxy. +Do not compare the runs before that file has been written. If Gym reports that a worker +exceeded its shutdown timeout, still wait for shutdown and check the statistics file. +Missing statistics are a failed run, not zero usage. + +## 4. Repeat with routing + +Use the same two terminals, environment, API key, and `RUN_DIR` as the fixed run. +Do not change the TOML or generation settings. Run each **complete block** below: both +terminals must use the `routed` directory, not the earlier `fixed` directory. + +**Terminal 1 — start the routed servers, after stopping the fixed servers:** + +```bash +OUT="$RUN_DIR/routed" +mkdir "$OUT" && +git -C "$WORK/Gym" rev-parse HEAD > "$OUT/gym-commit.txt" && +gym env start --resources-server mcqa --model-type switchyard_model --model routed \ + "++policy_model.responses_api_models.switchyard_model.deployment=$EXAMPLE/routes.toml" \ + ++policy_model.responses_api_models.switchyard_model.switchyard_base_url=null \ + "++policy_model.responses_api_models.switchyard_model.condition_dir=$OUT" \ + ++mcqa_simple_agent.responses_api_agents.simple_agent.max_steps=1 \ + ++observability_enabled=true \ + "++model_call_capture_dir=$OUT/model-calls" \ + "++nemo_gym_log_dir=$OUT/server-logs" \ + "hydra.run.dir=$OUT/hydra-start" +``` + +Wait for **`All 3 / 3 servers ready!`**. + +**Terminal 2 — collect the routed results:** + +```bash +OUT="$RUN_DIR/routed" +gym eval run --no-serve --agent mcqa_simple_agent \ + --input "$WORK/Gym/resources_servers/mcqa/data/example.jsonl" \ + --output "$OUT/rollouts.jsonl" \ + --limit 5 --num-repeats 1 --concurrency 1 \ + --temperature 0 --max-output-tokens 4096 \ + ++route_failures_to_sidecar=true \ + ++observability_enabled=true \ + "++model_call_capture_dir=$OUT/model-calls" \ + "hydra.run.dir=$OUT/hydra-eval" +``` + +**After collection finishes, press Ctrl-C in Terminal 1 and wait for shutdown.** +The comparison needs `routed/switchyard-stats.json`, which does not exist while those +servers are still running. Only then continue to Step 5. + +The route is selected when **starting the servers**. Changing a model flag only on +`gym eval run --no-serve` does not change a running server's route. + +On the clean path, these five tasks use 15 upstream calls in total: five fixed answers, +five classifier calls, and five routed answers. `routes.toml` disables Switchyard HTTP +retries, but this is not a hard spending cap: Gym retries and routing fallbacks can add calls. +The answer limit is 4,096 tokens; the separate classifier limit is 512 tokens. These are limits, +not fixed usage. Keep the same answer limit and per-model reasoning settings in both conditions. + +## 5. Compare the runs + +After **both server runs have shut down**, run this in Terminal 2 with the same variables +and environment. This prints the comparison table; routed collection alone does not. +The command only reads saved files and makes no inference calls. + +```bash +python "$EXAMPLE/compare.py" "$RUN_DIR/fixed" "$RUN_DIR/routed" +``` + +The script first reports expected, completed, missing, unmatched, and failed rollouts. It +requires identical materialized tasks and generation settings, matching deployment hashes, +the pinned versions, one successful captured model call per task, and usable final answers. +It refuses partial comparisons: classifier statistics cover the whole run and cannot be +fairly combined with only a successful subset of tasks. + +Then it prints paired mean reward, selected-model tokens, classifier tokens, their combined +reported total, mean rollout latency, routing overhead, request/error counts, and selected +models. Model names come from `ng_model_call_capture.calls[].model`, not the agent's top-level +`response.model`, which may contain only the route name. Capture token totals must agree with +proxy totals. The capture summaries retain model/status/usage fields; the script does not +expect raw response payloads inside them. + +- **Rollout latency** includes the agent, model request, and verification. +- **Routing overhead** is Switchyard's reported routing time, including classifier work. + It is already part of end-to-end latency; do not add it again. +- **Tokens are not dollars.** Model prices differ, and provider-reported usage may be incomplete. +- **Zero errors does not prove every classifier decision was valid.** Switchyard 0.2.0 can + fall back to the strong model after an unusable verdict without incrementing classifier + errors. Inspect `server-logs/policy_model.log`; the report labels this limitation. + +### Interpreting your result + +The report describes your own runs: there are no fixed expected scores or routing proportions. +Fewer answer-model tokens do not necessarily mean fewer combined tokens once classifier +usage is included. Routing decisions and service latency can vary. Five questions demonstrate +the workflow, not a statistically meaningful routing advantage. + +## What was saved? + +Each condition has its own directory under `RUN_DIR`: + +| Artifact | What it tells you | +|---|---| +| `gym-commit.txt` | Exact Gym source revision. | +| `switchyard-condition.json` | Route, hosted Switchyard version, deployment hash, and redacted configuration. | +| `switchyard-stats.json` | Shutdown snapshot: selected-model and classifier usage, latency, errors, and routing statistics. | +| `rollouts_materialized_inputs.jsonl` | The exact tasks and generation settings, with task/repeat indexes for pairing. | +| `rollouts.jsonl` / `rollouts_failures.jsonl` | Completed results and separately recorded infrastructure failures. | +| `rollouts_aggregate_metrics.json` | Gym's aggregate evaluation metrics. | +| `model-calls/` | Raw per-rollout requests and responses. Normalized summaries appear in the rollout's `ng_model_call_capture`. | +| `server-logs/` | Gym component logs, including Switchyard routing decisions and classifier warnings. | +| `hydra-start/` / `hydra-eval/` | Resolved configuration snapshots from the two commands, kept out of the repository root. | + +Keep keys in environment variables, not in TOML or committed files. Captures contain task +prompts and model responses; review them before sharing. + +## Next experiment + +To change models or providers, edit `llm_clients` and `targets` in `routes.toml` **before** +running both conditions into a fresh result directory. The classifier provider must support +the pinned version's strict JSON Schema output. To make the smoke test smaller, use +`--limit 1` in both collections; the clean path then has three upstream calls. + +For other benchmarks, agent harnesses, and externally managed proxies, see the +[Gym Switchyard integration reference](https://docs.nvidia.com/nemo/gym/main/model-server/switchyard/). +This small comparison script deliberately supports only the hosted, single-turn setup above. diff --git a/benchmark/nemo_gym/architecture.svg b/benchmark/nemo_gym/architecture.svg new file mode 100644 index 000000000..d0d7259b9 --- /dev/null +++ b/benchmark/nemo_gym/architecture.svg @@ -0,0 +1,86 @@ + + NeMo Gym and Switchyard for Evaluation + The evaluation flow follows four steps: dataset, initialize environment, agent executes, and verify. The simple agent runs in the agent server; the MCQA verifier runs in the resources server. Below the agent, Gym's model server hosts Switchyard in-process and exchanges requests and responses with upstream models. Run the same tasks twice: the fixed route always uses the stronger model, while the routed condition uses a classifier to choose the weaker or stronger model. Compare paired rewards, selected models, model and classifier tokens, latency, and routing statistics. + SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + + + + + + + + + NeMo Gym + Switchyard for Evaluation + Run the same tasks with a fixed model or a router, then compare rewards, tokens and latency + + + NeMo Gym + + + Dataset + (tasks) + MCQA examples + + + Initialize + Environment + + + Agent + Executes + Agent server + simple_agent + + + Verify + Resources server + MCQA verifier + + + + + + + + Model request / response + + + Model server + + Switchyard (hosted by Gym) + Fixed: always use the stronger model + Routed: classifier chooses weak or strong + + + + Upstream models + + Weaker model + + Stronger model + + + Usage / + routing + + + Output: + Compare fixed vs. routed runs + Per-task rewards + Selected models + Model + classifier tokens + Latency and routing statistics + diff --git a/benchmark/nemo_gym/compare.py b/benchmark/nemo_gym/compare.py new file mode 100644 index 000000000..3854aaa5c --- /dev/null +++ b/benchmark/nemo_gym/compare.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compare complete, paired, single-turn MCQA runs from the hosted Gym tutorial.""" + +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 + +GYM_COMMIT = "3a26c35fa90c243427378569511f7b06f503e0fd" +SWITCHYARD_VERSION = "0.2.0" +HOSTED_SCOPE = "this run (proxy hosted for exactly this run)" + + +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_object(path: Path) -> dict[str, Any]: + """Read a JSON artifact with an object at its root.""" + value = json.loads(path.read_text(encoding="utf-8")) + require(isinstance(value, dict), f"{path}: expected a JSON object") + return cast(dict[str, Any], value) + + +def read_jsonl(path: Path) -> list[dict[str, Any]]: + """Read JSONL objects, ignoring empty lines.""" + rows = [ + json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip() + ] + require(all(isinstance(row, dict) for row in rows), f"{path}: expected JSON objects") + return rows + + +def index_rows(path: Path) -> dict[tuple[int, int], dict[str, Any]]: + """Index by Gym task and repeat, rejecting duplicate or malformed identities.""" + indexed: dict[tuple[int, int], dict[str, Any]] = {} + for row in read_jsonl(path): + key = (row.get("_ng_task_index"), row.get("_ng_rollout_index")) + require( + all(type(part) is int and part >= 0 for part in key), + f"{path}: invalid task/repeat index", + ) + key = cast(tuple[int, int], key) + require(key not in indexed, f"{path}: duplicate task/repeat {key}") + indexed[key] = row + return indexed + + +def has_answer(response: Any) -> bool: + """Require a completed Responses API message, not just reasoning or a tool call.""" + return ( + isinstance(response, dict) + and response.get("status") == "completed" + and any( + item.get("type") == "message" + and item.get("role") == "assistant" + and any( + part.get("type") == "output_text" + and isinstance(part.get("text"), str) + and part["text"].strip() + for part in (item.get("content") or []) + if isinstance(part, dict) + ) + for item in (response.get("output") or []) + if isinstance(item, dict) + ) + ) + + +def load_run(path: Path) -> dict[str, Any]: + """Load expected tasks, completed rollouts, failures, and hosted-proxy provenance.""" + failure_path = path / "rollouts_failures.jsonl" + rollout_path = path / "rollouts.jsonl" + condition_path = path / "switchyard-condition.json" + stats_path = path / "switchyard-stats.json" + require( + condition_path.is_file(), + f"Missing {condition_path}. This file is written when the model server starts. " + "Set Terminal 1's OUT and condition_dir to the same run directory as Terminal 2. " + "If another condition's metadata was overwritten, rerun that condition in a fresh directory.", + ) + require( + stats_path.is_file(), + f"Missing {stats_path}. After evaluation finishes, press Ctrl-C in the terminal " + "running gym env start and wait for shutdown to save the statistics. " + "If the servers have already stopped, inspect server-logs/policy_model.log for snapshot errors.", + ) + return { + "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 [], + "condition": read_object(condition_path), + "snapshot": read_object(stats_path), + "gym_commit": (path / "gym-commit.txt").read_text(encoding="utf-8").strip(), + } + + +def summarize(run: dict[str, Any], route: str) -> tuple[dict[str, int | float], Counter[str]]: + """Calculate metrics only for this tutorial's complete, one-call-per-task runs.""" + condition, snapshot = run["condition"], run["snapshot"] + require( + condition["route"] == route, + f"{route}: manifest records route {condition['route']!r}. " + "The server's route and output directory must agree; rerun this condition in a fresh directory.", + ) + require(condition["mode"] == snapshot["mode"] == "hosted", f"{route}: expected hosted mode") + require(snapshot["scope"] == HOSTED_SCOPE, f"{route}: statistics are not run-scoped") + require( + condition["nemo_switchyard_version"] == SWITCHYARD_VERSION, + f"{route}: unexpected Switchyard version", + ) + require(run["gym_commit"] == GYM_COMMIT, f"{route}: unexpected Gym commit") + stats = snapshot["stats"] + classifier = stats["classifier"] + model_errors = number(stats["total_errors"], "model errors") + classifier_errors = number(classifier["total_errors"], "classifier errors") + require( + model_errors == classifier_errors == 0, + f"{route}: model errors={model_errors}, classifier errors={classifier_errors}", + ) + + calls, rewards, latencies = [], [], [] + for key, row in sorted(run["rows"].items()): + label = f"{route} {key}" + require(not row.get("_ng_failure_class"), f"{label}: failed rollout") + require(has_answer(row["response"]), f"{label}: missing or incomplete final answer") + reward = number(row["reward"], f"{label}: reward") + require(reward <= 1, f"{label}: MCQA reward must be between zero and one") + capture = row["ng_model_call_capture"] + require(isinstance(capture, dict), f"{label}: missing model-call capture") + require(not capture.get("gaps"), f"{label}: incomplete model-call capture") + require(len(capture["calls"]) == 1, f"{label}: expected exactly one captured model call") + call = capture["calls"][0] + require( + call["status_code"] == 200 and not call.get("error_category"), + f"{label}: model call failed", + ) + require(call["response_status"] == "completed", f"{label}: captured call is incomplete") + model = call.get("model") + require( + isinstance(model, str) and bool(model.strip()) and model not in {"fixed", "routed"}, + f"{label}: missing selected model", + ) + calls.append(call) + rewards.append(reward) + latencies.append(number(row["ng_perf"]["total_latency_ms"], f"{label}: rollout latency")) + + selected_tokens = sum(number(call["tokens_total"], "captured model tokens") for call in calls) + require( + number(stats["total_requests"], "model requests") == len(calls), + f"{route}: proxy/capture request counts differ", + ) + require( + number(stats["total_tokens"]["total"], "proxy model tokens") == selected_tokens, + f"{route}: proxy/capture token totals differ", + ) + classifier_calls = number(classifier["total_requests"], "classifier requests") + require( + classifier_calls == (0 if route == "fixed" else len(calls)), + f"{route}: unexpected classifier request count", + ) + classifier_tokens = number(classifier["total_tokens"]["total"], "classifier tokens") + return { + "Paired rollouts": len(rewards), + "Mean reward": mean(rewards), + "Selected-model tokens": selected_tokens, + "Classifier tokens": classifier_tokens, + "Combined reported tokens": selected_tokens + classifier_tokens, + "Mean rollout latency (ms)": mean(latencies), + "Mean routing overhead (ms)": number( + stats["routing_overhead"]["avg_ms"], "routing overhead" + ), + "Classifier requests": classifier_calls, + "Model errors": model_errors, + "Classifier errors": classifier_errors, + }, Counter(call["model"] for call in calls) + + +def compare(fixed: Path, routed: Path) -> None: + """Report coverage first, then compare like-for-like complete runs.""" + runs = {"fixed": load_run(fixed), "routed": load_run(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)}, " + f"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)}, " + f"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", + ) + hashes = [run["condition"].get("deployment_sha256") for run in runs.values()] + require( + all( + isinstance(value, str) + and len(value) == 64 + and all(char in "0123456789abcdef" for char in value) + for value in hashes + ) + and hashes[0] == hashes[1], + "Expected the same recorded deployment hash in both runs", + ) + 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 = [str(value) if type(value) is int else f"{value:.3f}" for value in values] + print(f"{metric:<29} {formatted[0]:>14} {formatted[1]:>14}") + for name, (_, models) in summaries.items(): + fallbacks = runs[name]["snapshot"]["stats"].get("routing_fallbacks", "unavailable") + print(f"\n{name} selected models: {json.dumps(models, sort_keys=True)}") + print(f"{name} reported fallbacks: {json.dumps(fallbacks, sort_keys=True)}") + print( + "\nClassifier fail-open decisions are not counted in these v0.2.0 statistics; " + "inspect Gym's model-server log." + ) + print( + "Reported tokens are not dollar costs. A small run demonstrates the workflow, 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("fixed", type=Path, help="Fixed run directory") + parser.add_argument("routed", type=Path, help="Routed run directory") + args = parser.parse_args(argv) + try: + compare(args.fixed, args.routed) + 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/routes.toml b/benchmark/nemo_gym/routes.toml new file mode 100644 index 000000000..1d1ae8914 --- /dev/null +++ b/benchmark/nemo_gym/routes.toml @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +schema_version = 1 + +[llm_clients.nvidia] +format = "openai_chat" +base_url = "https://integrate.api.nvidia.com/v1" +api_key_env = "NVIDIA_API_KEY" +max_retries = 0 + +[targets.strong] +id = "nvidia/nemotron-3-super-120b-a12b" +llm_client = "nvidia" +extra_body = { chat_template_kwargs = { enable_thinking = false } } + +[targets.weak] +id = "openai/gpt-oss-20b" +llm_client = "nvidia" +extra_body = { reasoning_effort = "low" } + +[routes.fixed] +id = "fixed" +type = "passthrough" +target = "strong" + +[routes.routed] +id = "routed" +type = "llm_classifier" +mode = "capability" +classifier_target = "weak" +strong_target = "strong" +weak_target = "weak" +base_threshold = 0.5 +max_output_tokens = 512 diff --git a/tests/test_nemo_gym_compare.py b/tests/test_nemo_gym_compare.py new file mode 100644 index 000000000..103d3334f --- /dev/null +++ b/tests/test_nemo_gym_compare.py @@ -0,0 +1,393 @@ +# 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 +import os +import re +import shutil +import subprocess +from copy import deepcopy +from pathlib import Path +from types import ModuleType +from typing import Any + +import pytest + +MISSING = object() +BIG = "nvidia/nemotron-3-super-120b-a12b" +SMALL = "openai/gpt-oss-20b" +FILES = { + "inputs": "rollouts_materialized_inputs.jsonl", + "rows": "rollouts.jsonl", + "failures": "rollouts_failures.jsonl", + "condition": "switchyard-condition.json", + "snapshot": "switchyard-stats.json", + "commit": "gym-commit.txt", +} + + +@pytest.fixture +def comparator() -> ModuleType: + 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 + assert 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]]: + runs = {} + for route in ("fixed", "routed"): + inputs, rows = [], [] + for index in range(2): + task = { + "_ng_task_index": index, + "_ng_rollout_index": 0, + "expected_answer": "B", + "grading_mode": "strict_single_letter_boxed", + "agent_ref": {"name": "mcqa_simple_agent"}, + "responses_create_params": { + "input": [{"role": "user", "content": f"Question {index}"}], + "temperature": 0, + "max_output_tokens": 512, + }, + } + inputs.append(task) + response = { + "status": "completed", + "model": SMALL if route == "routed" and index == 0 else BIG, + "output": [ + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "\\boxed{B}" if index == 0 else "\\boxed{A}", + } + ], + } + ], + } + rows.append( + { + **deepcopy(task), + "response": dict(deepcopy(response), model=route), + "reward": 1 - index, + "ng_perf": {"total_latency_ms": 100 + 200 * index}, + "ng_model_call_capture": { + "gaps": [], + "calls": [ + { + "status_code": 200, + "error_category": None, + "response_status": "completed", + "model": response["model"], + "tokens_total": 30 + 10 * index, + } + ], + }, + } + ) + runs[route] = { + "inputs": inputs, + "rows": rows, + "failures": [], + "commit": "3a26c35fa90c243427378569511f7b06f503e0fd", + "condition": { + "route": route, + "mode": "hosted", + "nemo_switchyard_version": "0.2.0", + "deployment_sha256": "a" * 64, + }, + "snapshot": { + "mode": "hosted", + "scope": "this run (proxy hosted for exactly this run)", + "stats": { + "total_errors": 0, + "total_requests": 2, + "total_tokens": {"total": 70}, + "routing_overhead": {"avg_ms": 5}, + "classifier": { + "total_errors": 0, + "total_requests": 2 if route == "routed" else 0, + "total_tokens": {"total": 11 if route == "routed" else 0}, + }, + }, + }, + } + return runs + + +def _write_runs(tmp_path: Path, artifacts: dict[str, dict[str, Any]]) -> list[str]: + for route, run in artifacts.items(): + directory = tmp_path / route + directory.mkdir() + for name, filename in FILES.items(): + value = run[name] + if filename.endswith(".jsonl"): + text = "".join(json.dumps(row) + "\n" for row in value) + else: + text = value + "\n" if name == "commit" else json.dumps(value) + (directory / filename).write_text(text, encoding="utf-8") + return [str(tmp_path / route) for route in ("fixed", "routed")] + + +def test_reordered_pairing_uses_captured_models_and_separate_classifier_tokens( + tmp_path: Path, + comparator: ModuleType, + artifacts: dict, + capsys: pytest.CaptureFixture[str], +) -> None: + artifacts["routed"]["rows"].reverse() + artifacts["routed"]["inputs"].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 selected models: {{"{BIG}": 2}}' in output.out + assert f"routed selected models: {json.dumps({SMALL: 1, BIG: 1}, sort_keys=True)}" in output.out + expected = { + "Mean reward": ["0.500", "0.500"], + "Selected-model tokens": ["70", "70"], + "Classifier tokens": ["0", "11"], + "Combined reported tokens": ["70", "81"], + "Mean rollout latency (ms)": ["200", "200"], + "Mean routing overhead (ms)": ["5", "5"], + "Classifier requests": ["0", "2"], + } + for metric, values in expected.items(): + line = next(line for line in output.out.splitlines() if line.startswith(metric)) + assert line[len(metric) :].split() == values + + +@pytest.mark.parametrize( + ("path", "value"), + [ + ("routed.rows", []), + ( + "routed.failures", + [{"_ng_task_index": 0, "_ng_rollout_index": 0, "_ng_failure_class": "timeout"}], + ), + ("routed.rows.1._ng_task_index", 0), + ("routed.rows.0._ng_task_index", -1), + ("routed.rows.0._ng_task_index", True), + ("routed.rows.0._ng_rollout_index", "0"), + ("routed.rows.0._ng_rollout_index", MISSING), + ("routed.rows.0._ng_task_index", 99), + ("routed.inputs.1._ng_task_index", 0), + ("routed.inputs.0.expected_answer", "A"), + ("routed.inputs.0.responses_create_params.input.0.content", "Different question"), + ("routed.inputs.0.responses_create_params.temperature", 1), + ("routed.inputs.0.responses_create_params.max_output_tokens", 256), + ("routed.condition.deployment_sha256", "b" * 64), + ("routed.condition.deployment_sha256", MISSING), + ("routed.condition.nemo_switchyard_version", "0.3.0"), + ("routed.commit", "different-commit"), + ("routed.condition.route", "fixed"), + ("routed.condition.mode", "external"), + ("routed.snapshot.mode", "external"), + ("routed.snapshot.scope", "all runs"), + ("routed.condition", None), + ("routed.snapshot", []), + ("routed.rows.0._ng_failure_class", "timeout"), + ("routed.rows.0.reward", None), + ("routed.rows.0.reward", 1.1), + ("routed.rows.0.response", MISSING), + ("routed.rows.0.response", None), + ("routed.rows.0.response.status", "incomplete"), + ("routed.rows.0.response.output", []), + ("routed.rows.0.response.output", None), + ("routed.rows.0.response.output", [None]), + ("routed.rows.0.response.output.0.role", "user"), + ("routed.rows.0.response.output.0.content", None), + ("routed.rows.0.response.output.0.content", [None]), + ("routed.rows.0.response.output.0.content.0.text", " "), + ("routed.rows.0.ng_model_call_capture", MISSING), + ("routed.rows.0.ng_model_call_capture", None), + ("routed.rows.0.ng_model_call_capture.gaps", ["missing call"]), + ("routed.rows.0.ng_model_call_capture.calls", []), + ("routed.rows.0.ng_model_call_capture.calls", [{}, {}]), + ("routed.rows.0.ng_model_call_capture.calls", [None]), + ("routed.rows.0.ng_model_call_capture.calls.0.status_code", 500), + ("routed.rows.0.ng_model_call_capture.calls.0.error_category", "timeout"), + ("routed.rows.0.ng_model_call_capture.calls.0.response_status", None), + ("routed.rows.0.ng_model_call_capture.calls.0.response_status", "incomplete"), + ("routed.rows.0.ng_model_call_capture.calls.0.model", "routed"), + ("routed.rows.0.ng_model_call_capture.calls.0.model", MISSING), + ("routed.rows.0.ng_model_call_capture.calls.0.tokens_total", None), + ("routed.rows.0.ng_perf.total_latency_ms", None), + ("routed.snapshot.stats.total_requests", 3), + ("routed.snapshot.stats.total_tokens.total", 81), + ("routed.snapshot.stats.classifier.total_requests", 0), + ("fixed.snapshot.stats.classifier.total_requests", 2), + ("routed.snapshot.stats.classifier.total_tokens.total", None), + ], +) +def test_rejects_incompatible_or_malformed_artifacts( + tmp_path: Path, + comparator: ModuleType, + artifacts: dict, + capsys: pytest.CaptureFixture[str], + path: str, + value: Any, +) -> None: + keys = [int(key) if key.isdecimal() else key for key in path.split(".")] + target = artifacts + for key in keys[:-1]: + target = target[key] + if value is MISSING: + del target[keys[-1]] + else: + target[keys[-1]] = value + assert comparator.main(_write_runs(tmp_path, artifacts)) == 1 + output = capsys.readouterr() + assert "Cannot compare:" in output.err + assert "Mean reward" not in output.out + + +@pytest.mark.parametrize("side", ["fixed", "routed"]) +@pytest.mark.parametrize("classifier", [False, True]) +@pytest.mark.parametrize("value", [MISSING, None, 1]) +def test_error_counts_must_be_present_and_zero( + tmp_path: Path, + comparator: ModuleType, + artifacts: dict, + side: str, + classifier: bool, + value: Any, +) -> None: + stats = artifacts[side]["snapshot"]["stats"] + target = stats["classifier"] if classifier else stats + if value is MISSING: + del target["total_errors"] + else: + target["total_errors"] = value + assert comparator.main(_write_runs(tmp_path, artifacts)) == 1 + + +@pytest.mark.parametrize("sides", [("fixed",), ("routed",), ("fixed", "routed")]) +def test_missing_expected_row_rejects_even_when_both_sides_match( + tmp_path: Path, + comparator: ModuleType, + artifacts: dict, + capsys: pytest.CaptureFixture[str], + sides: tuple, +) -> None: + for side in sides: + artifacts[side]["rows"].pop() + assert comparator.main(_write_runs(tmp_path, artifacts)) == 1 + output = capsys.readouterr() + assert "missing=1" in output.out + assert "Incomplete runs" in output.err + assert "Mean reward" not in output.out + + +@pytest.mark.parametrize( + ("filename", "hint"), + [ + ("switchyard-condition.json", "when the model server starts"), + ("switchyard-stats.json", "press Ctrl-C"), + ], +) +def test_missing_server_artifacts_explain_the_required_step( + tmp_path: Path, + comparator: ModuleType, + artifacts: dict, + capsys: pytest.CaptureFixture[str], + filename: str, + hint: str, +) -> None: + args = _write_runs(tmp_path, artifacts) + missing = Path(args[1]) / filename + missing.unlink() + assert comparator.main(args) == 1 + output = capsys.readouterr() + assert str(missing) in output.err + assert hint in output.err + assert "Mean reward" not in output.out + assert not missing.exists() + + +def test_misfiled_manifest_explains_route_directory_mismatch( + tmp_path: Path, + comparator: ModuleType, + artifacts: dict, + capsys: pytest.CaptureFixture[str], +) -> None: + artifacts["fixed"]["condition"]["route"] = "routed" + args = _write_runs(tmp_path, artifacts) + assert comparator.main(args) == 1 + output = capsys.readouterr() + assert "fixed: manifest records route 'routed'" in output.err + assert "fresh directory" in output.err + assert "Mean reward" not in output.out + manifest = json.loads((Path(args[0]) / "switchyard-condition.json").read_text()) + assert manifest["route"] == "routed" + + +@pytest.mark.skipif(shutil.which("bash") is None, reason="Bash is not installed") +@pytest.mark.parametrize("block_index", [0, 1]) +def test_routed_commands_reset_a_stale_fixed_output_directory( + tmp_path: Path, + block_index: int, +) -> None: + readme = Path(__file__).resolve().parents[1] / "benchmark/nemo_gym/README.md" + section = readme.read_text(encoding="utf-8").split("## 4. Repeat with routing", 1)[1] + section = section.split("## 5.", 1)[0] + blocks = re.findall(r"```bash\n(.*?)\n```", section, re.DOTALL) + assert len(blocks) == 2 + run_dir = tmp_path / "run" + run_dir.mkdir() + result = subprocess.run( + [ + "bash", + "-c", + "gym() { printf '%s\\n' \"$@\"; }; git() { printf 'test-revision\\n'; };\n" + + blocks[block_index], + ], + env={ + "PATH": os.defpath, + "EXAMPLE": str(tmp_path / "example"), + "WORK": str(tmp_path / "work"), + "RUN_DIR": str(run_dir), + "OUT": str(run_dir / "fixed"), + "ROUTE": "fixed", + }, + text=True, + capture_output=True, + check=False, + timeout=10, + ) + assert result.returncode == 0, result.stderr + args = result.stdout.splitlines() + assert f"++model_call_capture_dir={run_dir}/routed/model-calls" in args + if block_index == 0: + assert args[args.index("--model") + 1] == "routed" + assert ( + f"++policy_model.responses_api_models.switchyard_model.condition_dir={run_dir}/routed" + in args + ) + assert (run_dir / "routed/gym-commit.txt").read_text() == "test-revision\n" + else: + assert args[args.index("--output") + 1] == str(run_dir / "routed/rollouts.jsonl") + assert not (run_dir / "fixed").exists() + + +@pytest.mark.skipif(shutil.which("bash") is None, reason="Bash is not installed") +def test_readme_shell_blocks_have_valid_bash_syntax() -> None: + readme = Path(__file__).resolve().parents[1] / "benchmark/nemo_gym/README.md" + blocks = re.findall(r"```bash\n(.*?)\n```", readme.read_text(encoding="utf-8"), re.DOTALL) + assert blocks, "The tutorial must contain executable Bash examples" + for index, block in enumerate(blocks, start=1): + result = subprocess.run( + ["bash", "-n"], input=block, text=True, capture_output=True, check=False, timeout=10 + ) + assert result.returncode == 0, f"Bash example {index}: {result.stderr}" From 892ee8d67f04f4547c2b5b7c198c374402128b31 Mon Sep 17 00:00:00 2001 From: Shashank Verma Date: Mon, 14 Sep 2026 15:11:55 -0700 Subject: [PATCH 2/6] refactor(tutorial): evaluate checkout routing through LiteLLM Refactor the Switchyard eval with Gym tutorial to use the LiteLLM proxy. Additional changes per review comments. Signed-off-by: Shashank Verma --- benchmark/README.md | 2 +- benchmark/nemo_gym/README.md | 311 ++++++---------- benchmark/nemo_gym/architecture.svg | 34 +- benchmark/nemo_gym/compare.py | 274 +++++++++----- benchmark/nemo_gym/gym_routing_plugin.py | 196 ++++++++++ benchmark/nemo_gym/litellm.yaml | 28 ++ benchmark/nemo_gym/routes.toml | 34 +- benchmark/nemo_gym/run.sh | 200 ++++++++++ tests/test_nemo_gym_compare.py | 442 +++++++++------------- tests/test_nemo_gym_litellm.py | 356 ++++++++++++++++++ tests/test_nemo_gym_run.py | 454 +++++++++++++++++++++++ 11 files changed, 1720 insertions(+), 611 deletions(-) create mode 100644 benchmark/nemo_gym/gym_routing_plugin.py create mode 100644 benchmark/nemo_gym/litellm.yaml create mode 100644 benchmark/nemo_gym/run.sh create mode 100644 tests/test_nemo_gym_litellm.py create mode 100644 tests/test_nemo_gym_run.py diff --git a/benchmark/README.md b/benchmark/README.md index 6440362b7..74fa8ec13 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -13,7 +13,7 @@ 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 example using Gym's included questions instead of Harbor, see +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 diff --git a/benchmark/nemo_gym/README.md b/benchmark/nemo_gym/README.md index 3794d0390..4466bbb1a 100644 --- a/benchmark/nemo_gym/README.md +++ b/benchmark/nemo_gym/README.md @@ -1,250 +1,159 @@ # Evaluate Switchyard routing with NeMo Gym -[NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) is a library for evaluating models and agents using tasks with verifiable outcomes. -This tutorial uses its five included multiple-choice examples to compare a fixed model with Switchyard routing. +[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 evaluation with a hosted Switchyard model server](architecture.svg) +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. -## What changes between the runs? +Gym provides the evaluation substrate: it supplies tasks, runs the agent, and verifies answers to report rewards. +Switchyard sits in the model-request path, selecting which upstream model serves each request. -Both routes are defined in [routes.toml](routes.toml): +![Gym evaluation through LiteLLM and Switchyard Random routing](architecture.svg) -| Route | Behavior | -|---|---| -| `fixed` | Always use Nemotron 3 Super. | -| `routed` | Ask GPT-OSS 20B to classify the task, then use GPT-OSS 20B or Nemotron 3 Super. | +## Understand the wiring -The dataset, agent, verifier, temperature, and answer-token limit stay the same. -The classifier is an extra model call: its tokens count even when the router selects GPT-OSS 20B. -Super uses `enable_thinking=false`; GPT-OSS 20B uses `reasoning_effort=low`, including for -classification. These per-model settings stay unchanged between conditions. This is not a -benchmark of either model's maximum reasoning capability. +Gym calls a LiteLLM endpoint through its `litellm_model` adapter. The Switchyard library is integrated in LiteLLM, which is the path we'll use in this example. As such, Switchyard does not run as a separate server here. -## 1. Install the pinned Gym checkout +**LiteLLM defines the model groups.** This excerpt from [litellm.yaml](litellm.yaml) shows the candidates; provider settings are omitted here: -You need Git, [uv](https://docs.astral.sh/uv/), two Bash terminals, and an NVIDIA API key for -the public endpoint `https://integrate.api.nvidia.com/v1`. Use these exact model IDs: -`openai/gpt-oss-20b` and `nvidia/nemotron-3-super-120b-a12b`. -Inference may consume credits. The classifier requires strict JSON Schema responses from -its selected model. +```yaml +model_list: + - model_name: fixed + litellm_params: {model: nvidia_nim/nvidia/nemotron-3-super-120b-a12b} + - model_name: routed + litellm_params: {model: nvidia_nim/nvidia/nemotron-3-super-120b-a12b} + - model_name: routed + litellm_params: {model: nvidia_nim/openai/gpt-oss-20b} +``` -**Smoke-tested:** A one-question paired run completed on this endpoint with correct answers -in both conditions, GPT-OSS 20B serving the routed answer, and no reported model or classifier -errors. The full five-question comparison below has not yet been validated for this pair. +`fixed` and `routed` are names we chose, not special Gym modes. The fixed group has one candidate, so it always uses Super. The routed group has two candidates for Switchyard to choose from. -Run from the **Switchyard repository root**: +**Switchyard defines the selection policy.** [routes.toml](routes.toml) contains: -```bash -WORK="$PWD/scratch/nemo-gym-tutorial" -mkdir -p "$WORK" && -git clone https://github.com/NVIDIA-NeMo/Gym.git "$WORK/Gym" && -git -C "$WORK/Gym" checkout 3a26c35fa90c243427378569511f7b06f503e0fd && -uv tool run --from uv==0.11.29 uv venv --python 3.13.14 "$WORK/.venv" && -uv tool run --from uv==0.11.29 uv pip install \ - --python "$WORK/.venv/bin/python" uv==0.11.29 -e "$WORK/Gym" +```toml +algorithm = "random" +seed = 6 ``` -The pinned uv can download Python 3.13.14 even if your existing uv is older. It is also installed -inside the tutorial environment for Gym's component setup; your global uv is left unchanged. +LiteLLM uses the [Switchyard integration](../../examples/litellm/README.md) to choose a model using `routes.toml`. A small adapter records each choice and makes the response compatible with Gym. -This keeps the checkout and environment under Switchyard's ignored `scratch/` directory, -without changing another Gym checkout. Use a fresh directory for this one-time setup. -The editable install lets Gym's component environments use the same pinned source. +**Gym requests a group, not a concrete model.** These excerpts show the model wiring inside the runner; **these are wrapped in a [run.sh](./run.sh) script we'll run later, and are not additional steps to execute in this tutorial**: -Gym installs **`nemo-switchyard==0.2.0`** into its model-server environment and hosts the native -proxy in-process. You do **not** need Docker or a separately running `switchyard-server`. -The Switchyard checkout contains this tutorial; its current `main` is **not** the proxy being -executed. Do not copy newer routing options into this version-pinned example. +```text +gym eval run --benchmark mmlu-redux --model-type litellm_model --model fixed \ + ++policy_base_url=http://127.0.0.1:4000/v1 ++policy_api_key=unused -## 2. Start the fixed condition — Terminal 1 +gym eval run --benchmark mmlu-redux --model-type litellm_model --model routed \ + ++policy_base_url=http://127.0.0.1:4000/v1 ++policy_api_key=unused +``` +Note the following +- `--model-type` selects the adapter +- `policy_base_url` points it at LiteLLM +- `--model` selects the group. The local proxy uses provider credentials from the environment, not Gym's placeholder key. +- The runner adds identical task limits and separate output/capture paths. The dataset, agent, verifier, temperature, and 4,096-token answer limit stay unchanged. -Keep this terminal at the Switchyard repository root. Set your API key in the terminal, -replacing the placeholder with your key: +## 1. Set up -```bash -export NVIDIA_API_KEY='' -``` +You need: -Start Gym's resources, agent, and model servers. Run only one Gym environment at a time. -The first start also installs their dependencies. +- 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 `openai/gpt-oss-20b` and `nvidia/nemotron-3-super-120b-a12b` + +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 -EXAMPLE="$PWD/benchmark/nemo_gym" -WORK="$PWD/scratch/nemo-gym-tutorial" -source "$WORK/.venv/bin/activate" -RUN_DIR="$EXAMPLE/results/first-run" -OUT="$RUN_DIR/fixed" -mkdir -p "$RUN_DIR" - -mkdir "$OUT" && -git -C "$WORK/Gym" rev-parse HEAD > "$OUT/gym-commit.txt" && -gym env start --resources-server mcqa --model-type switchyard_model --model fixed \ - "++policy_model.responses_api_models.switchyard_model.deployment=$EXAMPLE/routes.toml" \ - ++policy_model.responses_api_models.switchyard_model.switchyard_base_url=null \ - "++policy_model.responses_api_models.switchyard_model.condition_dir=$OUT" \ - ++mcqa_simple_agent.responses_api_agents.simple_agent.max_steps=1 \ - ++observability_enabled=true \ - "++model_call_capture_dir=$OUT/model-calls" \ - "++nemo_gym_log_dir=$OUT/server-logs" \ - "hydra.run.dir=$OUT/hydra-start" +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 ``` -Wait for **`All 3 / 3 servers ready!`**. Leave this terminal running. -`mkdir "$OUT"` deliberately refuses to reuse an existing condition directory. -For a new comparison, change `RUN_DIR` to the same fresh path in **both terminals**. +Gym and LiteLLM use separate Python environments. The proxy builds Switchyard bindings from this checkout; the native CLI workflow does not need Docker. -## 3. Run the five questions — Terminal 2 +## 2. Run both conditions -Open another Bash terminal at the **same Switchyard repository root**: +The default is five tasks per condition, normally ten upstream calls. Inference can consume credits, and retries can add calls. Replace the placeholder below with your NVIDIA API key, then run. Pasting a key into this command may save it in shell history. ```bash -EXAMPLE="$PWD/benchmark/nemo_gym" -WORK="$PWD/scratch/nemo-gym-tutorial" -source "$WORK/.venv/bin/activate" -RUN_DIR="$EXAMPLE/results/first-run" -OUT="$RUN_DIR/fixed" - -gym eval run --no-serve --agent mcqa_simple_agent \ - --input "$WORK/Gym/resources_servers/mcqa/data/example.jsonl" \ - --output "$OUT/rollouts.jsonl" \ - --limit 5 --num-repeats 1 --concurrency 1 \ - --temperature 0 --max-output-tokens 4096 \ - ++route_failures_to_sidecar=true \ - ++observability_enabled=true \ - "++model_call_capture_dir=$OUT/model-calls" \ - "hydra.run.dir=$OUT/hydra-eval" +export NVIDIA_API_KEY="your-api-key" +bash benchmark/nemo_gym/run.sh ``` -These questions are included with Gym. The MCQA resources server checks the answer letter -against the expected answer; it does not call an LLM judge. A wrong answer earns zero reward. -An infrastructure failure is a different outcome, recorded separately. +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//`. Keep this unauthenticated development proxy local; do not share or publicly expose it. -The two-terminal flow is intentional: Gym's one-command evaluation mode does not accept -`--split example`. `--no-serve` collects against the servers you already started. +In `run.sh`, one loop runs both conditions: only the model group and output paths change. Leave the configuration and prepared data unchanged until both runs finish. -**After collection finishes, press Ctrl-C in Terminal 1 and wait for shutdown to finish.** -Gym writes `switchyard-stats.json` during shutdown, before stopping its hosted proxy. -Do not compare the runs before that file has been written. If Gym reports that a worker -exceeded its shutdown timeout, still wait for shutdown and check the statistics file. -Missing statistics are a failed run, not zero usage. +Gym components can outlive the command briefly; let them finish shutting down before an immediate rerun. -## 4. Repeat with routing +## 3. Read the comparison -Use the same two terminals, environment, API key, and `RUN_DIR` as the fixed run. -Do not change the TOML or generation settings. Run each **complete block** below: both -terminals must use the `routed` directory, not the earlier `fixed` directory. +Start with coverage and selected models, then compare rewards, tokens, and latency. Here is an excerpt from the **two-task local stub test**, not a real-model benchmark: -**Terminal 1 — start the routed servers, after stopping the fixed servers:** +```text +Pairing: matched=2, fixed-only=0, routed-only=0 -```bash -OUT="$RUN_DIR/routed" -mkdir "$OUT" && -git -C "$WORK/Gym" rev-parse HEAD > "$OUT/gym-commit.txt" && -gym env start --resources-server mcqa --model-type switchyard_model --model routed \ - "++policy_model.responses_api_models.switchyard_model.deployment=$EXAMPLE/routes.toml" \ - ++policy_model.responses_api_models.switchyard_model.switchyard_base_url=null \ - "++policy_model.responses_api_models.switchyard_model.condition_dir=$OUT" \ - ++mcqa_simple_agent.responses_api_agents.simple_agent.max_steps=1 \ - ++observability_enabled=true \ - "++model_call_capture_dir=$OUT/model-calls" \ - "++nemo_gym_log_dir=$OUT/server-logs" \ - "hydra.run.dir=$OUT/hydra-start" +Metric fixed routed +Paired rollouts 2 2 +Mean reward 0.500 0.500 +Terminal-answer tokens 30 30 +Gateway-reported tokens 30 30 + +fixed selected models: {"nvidia_nim/nvidia/nemotron-3-super-120b-a12b": 2} +routed selected models: {"nvidia_nim/nvidia/nemotron-3-super-120b-a12b": 1, "nvidia_nim/openai/gpt-oss-20b": 1} ``` -Wait for **`All 3 / 3 servers ready!`**. +- **Pairing:** both conditions completed the same two tasks. Incomplete or invalid evidence is rejected instead of producing partial averages. +- **Selection:** fixed stayed on Super; routed used both models. A small Random run need not split evenly. +- **Reward and usage:** MCQA scores the boxed answer letter (correct = 1, wrong = 0). Token columns sum input and output tokens across rollouts, not just generated answers. The stub supplies answers and token counts, so these values demonstrate the report, not model quality or savings. -**Terminal 2 — collect the routed results:** +For real runs, weigh reward against usage and latency rather than treating fewer tokens as a win by itself. 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. -```bash -OUT="$RUN_DIR/routed" -gym eval run --no-serve --agent mcqa_simple_agent \ - --input "$WORK/Gym/resources_servers/mcqa/data/example.jsonl" \ - --output "$OUT/rollouts.jsonl" \ - --limit 5 --num-repeats 1 --concurrency 1 \ - --temperature 0 --max-output-tokens 4096 \ - ++route_failures_to_sidecar=true \ - ++observability_enabled=true \ - "++model_call_capture_dir=$OUT/model-calls" \ - "hydra.run.dir=$OUT/hydra-eval" -``` +## 4. 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. This runner and comparator are MMLU-Redux-specific, not a general benchmark launcher. -**After collection finishes, press Ctrl-C in Terminal 1 and wait for shutdown.** -The comparison needs `routed/switchyard-stats.json`, which does not exist while those -servers are still running. Only then continue to Step 5. +See `bash benchmark/nemo_gym/run.sh --help` for profile, port, repeat, and concurrency options. -The route is selected when **starting the servers**. Changing a model flag only on -`gym eval run --no-serve` does not change a running server's route. +
+Saved files and token counts -On the clean path, these five tasks use 15 upstream calls in total: five fixed answers, -five classifier calls, and five routed answers. `routes.toml` disables Switchyard HTTP -retries, but this is not a hard spending cap: Gym retries and routing fallbacks can add calls. -The answer limit is 4,096 tokens; the separate classifier limit is 512 tokens. These are limits, -not fixed usage. Keep the same answer limit and per-model reasoning settings in both conditions. +Each `fixed/` and `routed/` folder contains: -## 5. Compare the runs +- `rollouts.jsonl`: model responses and rewards. +- `rollouts_materialized_inputs.jsonl`: the tasks and settings used. +- `rollouts_failures.jsonl`: failed tasks, if any. +- `model-calls/` and `litellm-calls.jsonl`: model request logs. +- `run-provenance.json`: version and configuration details. -After **both server runs have shut down**, run this in Terminal 2 with the same variables -and environment. This prints the comparison table; routed collection alone does not. -The command only reads saved files and makes no inference calls. +If something fails, start with that folder's `gym.log` or the result folder's `litellm.log`. + +The "Gateway-reported tokens" column includes the "Terminal-answer tokens", so don't add them together. Missing token counts are unknown, not zero, and some provider retries may not appear in the totals. Random makes no classifier calls; routing time is already included in rollout latency. + +To view the comparison again without calling the models, replace `my-run` with your results folder: ```bash -python "$EXAMPLE/compare.py" "$RUN_DIR/fixed" "$RUN_DIR/routed" +"$GYM_DIR/.venv/bin/python" benchmark/nemo_gym/compare.py \ + benchmark/nemo_gym/results/my-run/fixed benchmark/nemo_gym/results/my-run/routed ``` -The script first reports expected, completed, missing, unmatched, and failed rollouts. It -requires identical materialized tasks and generation settings, matching deployment hashes, -the pinned versions, one successful captured model call per task, and usable final answers. -It refuses partial comparisons: classifier statistics cover the whole run and cannot be -fairly combined with only a successful subset of tasks. - -Then it prints paired mean reward, selected-model tokens, classifier tokens, their combined -reported total, mean rollout latency, routing overhead, request/error counts, and selected -models. Model names come from `ng_model_call_capture.calls[].model`, not the agent's top-level -`response.model`, which may contain only the route name. Capture token totals must agree with -proxy totals. The capture summaries retain model/status/usage fields; the script does not -expect raw response payloads inside them. - -- **Rollout latency** includes the agent, model request, and verification. -- **Routing overhead** is Switchyard's reported routing time, including classifier work. - It is already part of end-to-end latency; do not add it again. -- **Tokens are not dollars.** Model prices differ, and provider-reported usage may be incomplete. -- **Zero errors does not prove every classifier decision was valid.** Switchyard 0.2.0 can - fall back to the strong model after an unusable verdict without incrementing classifier - errors. Inspect `server-logs/policy_model.log`; the report labels this limitation. - -### Interpreting your result - -The report describes your own runs: there are no fixed expected scores or routing proportions. -Fewer answer-model tokens do not necessarily mean fewer combined tokens once classifier -usage is included. Routing decisions and service latency can vary. Five questions demonstrate -the workflow, not a statistically meaningful routing advantage. - -## What was saved? - -Each condition has its own directory under `RUN_DIR`: - -| Artifact | What it tells you | -|---|---| -| `gym-commit.txt` | Exact Gym source revision. | -| `switchyard-condition.json` | Route, hosted Switchyard version, deployment hash, and redacted configuration. | -| `switchyard-stats.json` | Shutdown snapshot: selected-model and classifier usage, latency, errors, and routing statistics. | -| `rollouts_materialized_inputs.jsonl` | The exact tasks and generation settings, with task/repeat indexes for pairing. | -| `rollouts.jsonl` / `rollouts_failures.jsonl` | Completed results and separately recorded infrastructure failures. | -| `rollouts_aggregate_metrics.json` | Gym's aggregate evaluation metrics. | -| `model-calls/` | Raw per-rollout requests and responses. Normalized summaries appear in the rollout's `ng_model_call_capture`. | -| `server-logs/` | Gym component logs, including Switchyard routing decisions and classifier warnings. | -| `hydra-start/` / `hydra-eval/` | Resolved configuration snapshots from the two commands, kept out of the repository root. | - -Keep keys in environment variables, not in TOML or committed files. Captures contain task -prompts and model responses; review them before sharing. - -## Next experiment - -To change models or providers, edit `llm_clients` and `targets` in `routes.toml` **before** -running both conditions into a fresh result directory. The classifier provider must support -the pinned version's strict JSON Schema output. To make the smoke test smaller, use -`--limit 1` in both collections; the clean path then has three upstream calls. - -For other benchmarks, agent harnesses, and externally managed proxies, see the -[Gym Switchyard integration reference](https://docs.nvidia.com/nemo/gym/main/model-server/switchyard/). -This small comparison script deliberately supports only the hosted, single-turn setup above. +Keep the saved inputs because the source dataset can change. Request logs contain prompts and responses, so review them before sharing. + +
+ +**Tested:** with Gym `v0.6.0`. diff --git a/benchmark/nemo_gym/architecture.svg b/benchmark/nemo_gym/architecture.svg index d0d7259b9..be53aff88 100644 --- a/benchmark/nemo_gym/architecture.svg +++ b/benchmark/nemo_gym/architecture.svg @@ -1,6 +1,6 @@ - NeMo Gym and Switchyard for Evaluation - The evaluation flow follows four steps: dataset, initialize environment, agent executes, and verify. The simple agent runs in the agent server; the MCQA verifier runs in the resources server. Below the agent, Gym's model server hosts Switchyard in-process and exchanges requests and responses with upstream models. Run the same tasks twice: the fixed route always uses the stronger model, while the routed condition uses a classifier to choose the weaker or stronger model. Compare paired rewards, selected models, model and classifier tokens, latency, and routing statistics. + Evaluating Switchyard routing with NeMo Gym + The evaluation uses MMLU-Redux tasks, Gym's simple agent and MCQA verifier. Gym's litellm_model calls a separate runner-managed LiteLLM proxy through the OpenAI Responses API. Switchyard libsy Random routing runs directly inside LiteLLM, not as a Switchyard HTTP runtime. Compare fixed and routed runs on identical inputs using rewards, selected models, model tokens, latency and routing statistics. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 @@ -23,7 +23,7 @@ SPDX-License-Identifier: Apache-2.0 - NeMo Gym + Switchyard for Evaluation + Evaluating Switchyard routing with NeMo Gym Run the same tasks with a fixed model or a router, then compare rewards, tokens and latency @@ -32,7 +32,7 @@ SPDX-License-Identifier: Apache-2.0 Dataset (tasks) - MCQA examples + MMLU-Redux Initialize @@ -54,23 +54,25 @@ SPDX-License-Identifier: Apache-2.0 - - Model request / response - - - Model server - - Switchyard (hosted by Gym) - Fixed: always use the stronger model - Routed: classifier chooses weak or strong + + Model request / response + + Gym model server + litellm_model + + OpenAI Responses API + + LiteLLM proxy + + Switchyard library Upstream models - Weaker model + GPT-OSS 20B - Stronger model + Nemotron 3 Super Usage / @@ -81,6 +83,6 @@ SPDX-License-Identifier: Apache-2.0 Compare fixed vs. routed runs Per-task rewards Selected models - Model + classifier tokens + Model tokens Latency and routing statistics diff --git a/benchmark/nemo_gym/compare.py b/benchmark/nemo_gym/compare.py index 3854aaa5c..34917812d 100644 --- a/benchmark/nemo_gym/compare.py +++ b/benchmark/nemo_gym/compare.py @@ -1,23 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Compare complete, paired, single-turn MCQA runs from the hosted Gym tutorial.""" +"""Compare complete, paired MCQA runs through LiteLLM and Switchyard Random routing.""" from __future__ import annotations import argparse import json import math +import re import sys from collections import Counter from pathlib import Path from statistics import mean from typing import Any, cast -GYM_COMMIT = "3a26c35fa90c243427378569511f7b06f503e0fd" -SWITCHYARD_VERSION = "0.2.0" -HOSTED_SCOPE = "this run (proxy hosted for exactly this run)" - def require(condition: bool, message: str) -> None: """Reject incomplete or incompatible evidence before calculating metrics.""" @@ -34,6 +31,12 @@ def number(value: Any, name: str) -> int | float: return cast(int | float, value) +def count(value: Any, name: str) -> int: + """Read a nonnegative integer counter.""" + require(type(value) is int and value >= 0, f"{name} must be a nonnegative integer") + return cast(int, value) + + def read_object(path: Path) -> dict[str, Any]: """Read a JSON artifact with an object at its root.""" value = json.loads(path.read_text(encoding="utf-8")) @@ -87,112 +90,141 @@ def has_answer(response: Any) -> bool: def load_run(path: Path) -> dict[str, Any]: - """Load expected tasks, completed rollouts, failures, and hosted-proxy provenance.""" - failure_path = path / "rollouts_failures.jsonl" - rollout_path = path / "rollouts.jsonl" - condition_path = path / "switchyard-condition.json" - stats_path = path / "switchyard-stats.json" - require( - condition_path.is_file(), - f"Missing {condition_path}. This file is written when the model server starts. " - "Set Terminal 1's OUT and condition_dir to the same run directory as Terminal 2. " - "If another condition's metadata was overwritten, rerun that condition in a fresh directory.", - ) + """Require complete gateway request evidence alongside Gym's rollout artifacts.""" + failure_path, rollout_path = path / "rollouts_failures.jsonl", path / "rollouts.jsonl" + for filename in ("run-provenance.json", "litellm-calls.jsonl"): + require( + (path / filename).is_file(), + f"Missing {path / filename}. Inspect {path / 'gym.log'} and {path.parent / 'litellm.log'}.", + ) + events: dict[str, dict[str, dict[str, Any]]] = {"start": {}, "finish": {}} + for event in read_jsonl(path / "litellm-calls.jsonl"): + kind, request_id = event.get("event"), event.get("request_id") + require(kind in events, f"{path}: unknown gateway event") + kind = cast(str, kind) + require(isinstance(request_id, str) and bool(request_id), f"{path}: missing request ID") + request_id = cast(str, request_id) + require(request_id not in events[kind], f"{path}: duplicate gateway {kind} event") + events[kind][request_id] = event require( - stats_path.is_file(), - f"Missing {stats_path}. After evaluation finishes, press Ctrl-C in the terminal " - "running gym env start and wait for shutdown to save the statistics. " - "If the servers have already stopped, inspect server-logs/policy_model.log for snapshot errors.", + bool(events["start"]) and events["start"].keys() == events["finish"].keys(), + f"{path}: incomplete gateway request evidence", ) return { "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 [], - "condition": read_object(condition_path), - "snapshot": read_object(stats_path), - "gym_commit": (path / "gym-commit.txt").read_text(encoding="utf-8").strip(), + "provenance": read_object(path / "run-provenance.json"), + "events": events, } def summarize(run: dict[str, Any], route: str) -> tuple[dict[str, int | float], Counter[str]]: - """Calculate metrics only for this tutorial's complete, one-call-per-task runs.""" - condition, snapshot = run["condition"], run["snapshot"] - require( - condition["route"] == route, - f"{route}: manifest records route {condition['route']!r}. " - "The server's route and output directory must agree; rerun this condition in a fresh directory.", - ) - require(condition["mode"] == snapshot["mode"] == "hosted", f"{route}: expected hosted mode") - require(snapshot["scope"] == HOSTED_SCOPE, f"{route}: statistics are not run-scoped") - require( - condition["nemo_switchyard_version"] == SWITCHYARD_VERSION, - f"{route}: unexpected Switchyard version", - ) - require(run["gym_commit"] == GYM_COMMIT, f"{route}: unexpected Gym commit") - stats = snapshot["stats"] - classifier = stats["classifier"] - model_errors = number(stats["total_errors"], "model errors") - classifier_errors = number(classifier["total_errors"], "classifier errors") - require( - model_errors == classifier_errors == 0, - f"{route}: model errors={model_errors}, classifier errors={classifier_errors}", - ) + """Join final answers by response ID while retaining all recorded gateway work.""" + runtime = run["provenance"]["runtime"] + allowed_models = runtime["models"][route] + finishes = list(run["events"]["finish"].values()) + responses: dict[str, dict[str, Any]] = {} + tokens = [] + routing_times = [] + errors = unknown_usage = 0 + for phase in run["events"].values(): + for event in phase.values(): + require(event["route"] == route, f"{route}: gateway event belongs to another route") + require( + event["instance_id"] == runtime["instance_id"], + f"{route}: gateway events belong to another proxy instance", + ) + for event in finishes: + failed = event.get("status_code") != 200 or bool(event.get("error_type")) + errors += int(failed) + if not failed: + model, response_id = event.get("selected_model"), event.get("response_id") + require(model in allowed_models, f"{route}: missing or invalid Switchyard selection") + require( + event.get("deployment_model") == model, + f"{route}: selected model differs from the LiteLLM deployment", + ) + require( + isinstance(response_id, str) and bool(response_id), + f"{route}: missing gateway response ID", + ) + require(response_id not in responses, f"{route}: ambiguous gateway response ID") + responses[response_id] = event + number(event.get("tokens_total"), "gateway response tokens") + number(event.get("routing_ms"), "routing decision time") + if event.get("tokens_total") is None: + unknown_usage += 1 + else: + tokens.append(number(event["tokens_total"], "gateway tokens")) + if event.get("routing_ms") is not None: + routing_times.append(number(event["routing_ms"], "routing decision time")) - calls, rewards, latencies = [], [], [] + calls, rewards, latencies, models = [], [], [], [] + used_response_ids: set[str] = set() + captured_attempts = captured_errors = 0 for key, row in sorted(run["rows"].items()): label = f"{route} {key}" require(not row.get("_ng_failure_class"), f"{label}: failed rollout") require(has_answer(row["response"]), f"{label}: missing or incomplete final answer") + response_id = row["response"].get("id") + require(isinstance(response_id, str) and bool(response_id), f"{label}: missing response id") + require(response_id not in used_response_ids, f"{label}: response reused across rollouts") + used_response_ids.add(response_id) reward = number(row["reward"], f"{label}: reward") require(reward <= 1, f"{label}: MCQA reward must be between zero and one") capture = row["ng_model_call_capture"] require(isinstance(capture, dict), f"{label}: missing model-call capture") require(not capture.get("gaps"), f"{label}: incomplete model-call capture") - require(len(capture["calls"]) == 1, f"{label}: expected exactly one captured model call") - call = capture["calls"][0] + records = capture["calls"] + require( + isinstance(records, list) and all(isinstance(call, dict) for call in records), + f"{label}: invalid captures", + ) + 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}: model call failed", + f"{label}: terminal call failed", ) require(call["response_status"] == "completed", f"{label}: captured call is incomplete") - model = call.get("model") + require(call.get("model") == route, f"{label}: captured response has the wrong model group") + require(response_id in responses, f"{label}: final answer has no gateway evidence") + event = responses[response_id] require( - isinstance(model, str) and bool(model.strip()) and model not in {"fixed", "routed"}, - f"{label}: missing selected model", + event.get("response_status") == "completed", f"{label}: gateway response is incomplete" + ) + require( + number(call["tokens_total"], "terminal-answer tokens") == event["tokens_total"], + f"{label}: gateway/capture token totals differ", + ) + captured_attempts += len(records) + captured_errors += sum( + record.get("status_code") != 200 or bool(record.get("error_category")) + for record in records ) calls.append(call) rewards.append(reward) latencies.append(number(row["ng_perf"]["total_latency_ms"], f"{label}: rollout latency")) + models.append(event["selected_model"]) - selected_tokens = sum(number(call["tokens_total"], "captured model tokens") for call in calls) - require( - number(stats["total_requests"], "model requests") == len(calls), - f"{route}: proxy/capture request counts differ", - ) - require( - number(stats["total_tokens"]["total"], "proxy model tokens") == selected_tokens, - f"{route}: proxy/capture token totals differ", - ) - classifier_calls = number(classifier["total_requests"], "classifier requests") - require( - classifier_calls == (0 if route == "fixed" else len(calls)), - f"{route}: unexpected classifier request count", - ) - classifier_tokens = number(classifier["total_tokens"]["total"], "classifier tokens") return { "Paired rollouts": len(rewards), "Mean reward": mean(rewards), - "Selected-model tokens": selected_tokens, - "Classifier tokens": classifier_tokens, - "Combined reported tokens": selected_tokens + classifier_tokens, - "Mean rollout latency (ms)": mean(latencies), - "Mean routing overhead (ms)": number( - stats["routing_overhead"]["avg_ms"], "routing overhead" + "Terminal-answer tokens": sum( + number(call["tokens_total"], "terminal-answer tokens") for call in calls ), - "Classifier requests": classifier_calls, - "Model errors": model_errors, - "Classifier errors": classifier_errors, - }, Counter(call["model"] for call in calls) + "Gateway-reported tokens": sum(tokens), + "Gateway requests w/o usage": unknown_usage, + "Mean rollout latency (ms)": mean(latencies), + "Mean recorded routing (ms)": mean(routing_times), + "Recorded routing decisions": len(routing_times), + "Gateway requests": len(finishes), + "Gateway errors": errors, + "Captured model attempts": captured_attempts, + "Captured failed attempts": captured_errors, + }, Counter(models) def compare(fixed: Path, routed: Path) -> None: @@ -203,48 +235,94 @@ def compare(fixed: Path, routed: Path) -> None: expected, actual = set(run["inputs"]), set(run["rows"]) missing, unexpected = expected - actual, actual - expected print( - f"{name}: expected={len(expected)}, completed={len(actual)}, " - f"missing={len(missing)}, unexpected={len(unexpected)}, failures={len(run['failures'])}" + 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)}, " - f"fixed-only={len(fixed_keys - routed_keys)}, routed-only={len(routed_keys - fixed_keys)}" + 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", ) - hashes = [run["condition"].get("deployment_sha256") for run in runs.values()] - require( - all( - isinstance(value, str) - and len(value) == 64 - and all(char in "0123456789abcdef" for char in value) - for value in hashes + provenances = [run["provenance"] for run in runs.values()] + for provenance in provenances: + require( + all( + isinstance(provenance.get(key), str) and provenance[key].strip() + for key in ("gym_revision", "switchyard_revision") + ), + "Missing Gym or Switchyard revision", ) - and hashes[0] == hashes[1], - "Expected the same recorded deployment hash in both runs", - ) + runtime = provenance["runtime"] + require(runtime["mode"] == "litellm_libsy", "Expected the LiteLLM libsy integration") + require( + runtime["routing_plugin"] == "switchyard_litellm.RandomRoutingPlugin", + "Expected Switchyard Random routing", + ) + require( + all( + isinstance(runtime.get(key), str) and runtime[key].strip() + for key in ("instance_id", "litellm_version", "switchyard_version") + ), + "Missing LiteLLM runtime identity", + ) + require( + all( + re.fullmatch(r"[0-9a-f]{64}", str(runtime.get(key))) is not None + for key in ( + "profile_sha256", + "routing_sha256", + "callback_sha256", + "provider_base_sha256", + ) + ), + "Missing deployment fingerprint", + ) + models = runtime["models"] + require( + all( + isinstance(models.get(route), list) + and all(isinstance(model, str) and model for model in models[route]) + for route in ("fixed", "routed") + ), + "Missing configured models", + ) + require( + len(models["fixed"]) == 1 + and len(set(models["routed"])) == 2 + and models["fixed"][0] in models["routed"], + "Expected a fixed target and a routed pair containing it", + ) + require(provenances[0] == provenances[1], "Gym, Switchyard, or deployment provenance differs") summaries = {name: summarize(run, name) for name, run in runs.items()} + print(f"\nProvenance: {json.dumps(provenances[0], sort_keys=True)}") 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 = [str(value) if type(value) is int else f"{value:.3f}" for value in values] print(f"{metric:<29} {formatted[0]:>14} {formatted[1]:>14}") - for name, (_, models) in summaries.items(): - fallbacks = runs[name]["snapshot"]["stats"].get("routing_fallbacks", "unavailable") - print(f"\n{name} selected models: {json.dumps(models, sort_keys=True)}") - print(f"{name} reported fallbacks: {json.dumps(fallbacks, sort_keys=True)}") + for name, (summary, selected) in summaries.items(): + print(f"\n{name} selected models: {json.dumps(selected, sort_keys=True)}") + if ( + summary["Gateway errors"] + or summary["Gateway requests w/o usage"] + or summary["Captured failed attempts"] + or summary["Gateway requests"] != summary["Paired rollouts"] + or summary["Captured model attempts"] != summary["Paired rollouts"] + ): + print(f"WARNING: {name} completed with recovered errors or additional work.") + print("\nClassifier tokens: N/A (Random makes no classifier calls).") print( - "\nClassifier fail-open decisions are not counted in these v0.2.0 statistics; " - "inspect Gym's model-server log." + "Gateway totals include all recorded requests, including extra attempts; do not add terminal-answer tokens again." ) print( - "Reported tokens are not dollar costs. A small run demonstrates the workflow, not a routing advantage." + "Unreported failed-request usage is unknown, not free. Gateway counts are not exhaustive provider-attempt or fallback telemetry." ) + print("Routing time is already included in rollout latency. Tokens are not dollar costs.") + print("A small Random-routing run demonstrates the integration, not a routing advantage.") def main(argv: list[str] | None = None) -> int: diff --git a/benchmark/nemo_gym/gym_routing_plugin.py b/benchmark/nemo_gym/gym_routing_plugin.py new file mode 100644 index 000000000..866f6ded8 --- /dev/null +++ b/benchmark/nemo_gym/gym_routing_plugin.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Record LiteLLM request evidence and preserve Gym-compatible usage detail shapes.""" + +from __future__ import annotations + +import hashlib +import json +import os +from copy import deepcopy +from importlib.metadata import version +from pathlib import Path +from time import perf_counter +from typing import Any + +import yaml +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.router import RoutingContext +from switchyard_litellm import RandomRoutingPlugin +from switchyard_litellm.configuration import load_routing_plugin + + +def response_payload(response: Any) -> dict[str, Any]: + """Keep missing detail counts unknown while satisfying Gym's object-shaped fields.""" + payload = response.model_dump() if hasattr(response, "model_dump") else deepcopy(dict(response)) + 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): + """Join routing and callback instances using one runner-supplied gateway identity.""" + + def __init__( + self, routing_path: Path, profile_path: Path, results: Path, instance_id: str + ) -> None: + super().__init__() + if not instance_id.strip(): + raise ValueError("The runner must supply a gateway instance ID") + self.routing = load_routing_plugin(routing_path) + if not isinstance(self.routing, RandomRoutingPlugin): + raise ValueError("This example requires Switchyard Random routing") + self.results = results + profile = yaml.safe_load(profile_path.read_text(encoding="utf-8")) + models = { + route: [ + entry["litellm_params"]["model"] + for entry in profile["model_list"] + if entry["model_name"] == route + ] + for route in ("fixed", "routed") + } + if len(models["fixed"]) != 1 or len(set(models["routed"])) != 2: + raise ValueError("Define one fixed target and two distinct routed targets") + if models["fixed"][0] not in models["routed"]: + raise ValueError("The fixed target must be one of the routed targets") + self.runtime = { + "mode": "litellm_libsy", + "instance_id": instance_id, + "litellm_version": version("litellm"), + "switchyard_version": version("nemo-switchyard"), + "fastapi_version": version("fastapi"), + "starlette_version": version("starlette"), + "routing_plugin": "switchyard_litellm.RandomRoutingPlugin", + "models": models, + "profile_sha256": hashlib.sha256(profile_path.read_bytes()).hexdigest(), + "routing_sha256": hashlib.sha256(routing_path.read_bytes()).hexdigest(), + "callback_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "provider_base_sha256": hashlib.sha256( + os.environ.get("NVIDIA_BASE_URL", "").encode() + ).hexdigest(), + } + results.mkdir(parents=True, exist_ok=True) + (results / "litellm-runtime.json").write_text( + json.dumps(self.runtime, indent=2) + "\n", encoding="utf-8" + ) + + async def run(self, context: RoutingContext) -> RoutingContext: + """Measure the actual libsy decision without including provider inference.""" + started = perf_counter() + result = await self.routing.run(context) + result.signals["switchyard"]["routing_ms"] = (perf_counter() - started) * 1000 + return result + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, Any], call_type: Any + ) -> dict[str, Any] | None: + """Keep the configured plugin's deployment callback behavior intact.""" + return await self.routing.async_pre_call_deployment_hook(kwargs, call_type) + + def _record(self, data: dict[str, Any], event: dict[str, Any]) -> None: + """Write only allowlisted evidence, never request content or credentials.""" + route, request_id = data.get("model"), data.get("litellm_call_id") + if route not in ("fixed", "routed"): + raise ValueError("This example accepts only the fixed and routed model groups") + if not isinstance(request_id, str) or not request_id: + raise ValueError("LiteLLM request ID is missing") + directory = self.results / route + directory.mkdir(parents=True, exist_ok=True) + record = { + "instance_id": self.runtime["instance_id"], + "route": route, + "request_id": request_id, + **event, + } + with (directory / "litellm-calls.jsonl").open("a", encoding="utf-8") as stream: + stream.write(json.dumps(record, allow_nan=False) + "\n") + + async def async_pre_call_hook( + self, user_api_key_dict: Any, cache: Any, data: dict[str, Any], call_type: Any + ) -> dict[str, Any]: + """Record a request before routing so interrupted evidence can be detected.""" + self._record(data, {"event": "start"}) + if data.get("stream"): + raise ValueError("This example requires non-streaming requests") + return data + + async def async_post_call_success_hook( + self, data: dict[str, Any], user_api_key_dict: Any, response: Any + ) -> dict[str, Any]: + """Record the final response before LiteLLM replaces its model with the public alias.""" + payload = response_payload(response) + metadata = data.get("litellm_metadata") or {} + signal = (metadata.get("routing_plugin_signals") or {}).get("switchyard") or {} + self._record( + data, + { + "event": "finish", + "status_code": 200, + "error_type": None, + "response_id": payload.get("id"), + "response_status": payload.get("status"), + "selected_model": signal.get("selected_model_id"), + "deployment_model": metadata.get("deployment"), + "tokens_total": (payload.get("usage") or {}).get("total_tokens"), + "routing_ms": signal.get("routing_ms"), + }, + ) + return payload + + async def async_post_call_failure_hook( + self, + request_data: dict[str, Any], + original_exception: Exception, + user_api_key_dict: Any, + traceback_str: str | None = None, + ) -> None: + """Record failed gateway requests without inventing unreported token usage.""" + metadata = request_data.get("litellm_metadata") or {} + signal = (metadata.get("routing_plugin_signals") or {}).get("switchyard") or {} + self._record( + request_data, + { + "event": "finish", + "status_code": getattr(original_exception, "status_code", None), + "error_type": type(original_exception).__name__, + "response_id": None, + "response_status": None, + "selected_model": signal.get("selected_model_id"), + "deployment_model": metadata.get("deployment"), + "tokens_total": None, + "routing_ms": signal.get("routing_ms"), + }, + ) + + +PLUGIN = ( + GymRoutingPlugin( + Path(os.environ["SWITCHYARD_LITELLM_CONFIG"]), + Path(os.environ["NEMO_GYM_LITELLM_PROFILE"]), + Path(os.environ["NEMO_GYM_LITELLM_RESULTS"]), + os.environ["NEMO_GYM_LITELLM_INSTANCE_ID"], + ) + 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..91078488a --- /dev/null +++ b/benchmark/nemo_gym/litellm.yaml @@ -0,0 +1,28 @@ +model_list: + - model_name: fixed + litellm_params: &strong + model: nvidia_nim/nvidia/nemotron-3-super-120b-a12b + 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: *strong + - model_name: routed + litellm_params: + model: nvidia_nim/openai/gpt-oss-20b + api_base: os.environ/NVIDIA_BASE_URL + api_key: os.environ/NVIDIA_API_KEY + max_retries: 0 + extra_body: + reasoning_effort: low +router_settings: + num_retries: 0 + plugins: + - gym_routing_plugin.PLUGIN +litellm_settings: + num_retries: 0 + callbacks: + - gym_routing_plugin.PLUGIN diff --git a/benchmark/nemo_gym/routes.toml b/benchmark/nemo_gym/routes.toml index 1d1ae8914..00dcf73a6 100644 --- a/benchmark/nemo_gym/routes.toml +++ b/benchmark/nemo_gym/routes.toml @@ -1,35 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -schema_version = 1 - -[llm_clients.nvidia] -format = "openai_chat" -base_url = "https://integrate.api.nvidia.com/v1" -api_key_env = "NVIDIA_API_KEY" -max_retries = 0 - -[targets.strong] -id = "nvidia/nemotron-3-super-120b-a12b" -llm_client = "nvidia" -extra_body = { chat_template_kwargs = { enable_thinking = false } } - -[targets.weak] -id = "openai/gpt-oss-20b" -llm_client = "nvidia" -extra_body = { reasoning_effort = "low" } - -[routes.fixed] -id = "fixed" -type = "passthrough" -target = "strong" - -[routes.routed] -id = "routed" -type = "llm_classifier" -mode = "capability" -classifier_target = "weak" -strong_target = "strong" -weak_target = "weak" -base_threshold = 0.5 -max_output_tokens = 512 +algorithm = "random" +seed = 6 diff --git a/benchmark/nemo_gym/run.sh b/benchmark/nemo_gym/run.sh new file mode 100644 index 000000000..a463c1fac --- /dev/null +++ b/benchmark/nemo_gym/run.sh @@ -0,0 +1,200 @@ +#!/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="${LITELLM_CONFIG:-$SCRIPT_DIR/litellm.yaml}" +ROUTING="${SWITCHYARD_CONFIG:-$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/) + LITELLM_CONFIG Model inventory (litellm.yaml beside this script) + SWITCHYARD_CONFIG Routing policy (routes.toml beside this script) + NVIDIA_BASE_URL Provider endpoint (NVIDIA API Catalog /v1) +Provider credentials come from the environment variables named in the profile. +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)" +[[ "$PROFILE" = /* ]] || PROFILE="$PWD/$PROFILE" +[[ "$ROUTING" = /* ]] || ROUTING="$PWD/$ROUTING" +[[ "$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" rev-parse HEAD)" +[[ -z "$(git -C "$GYM_DIR" status --porcelain --untracked-files=no)" ]] || die "tracked Gym source must be clean" +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_PROFILE="$PROFILE" +export NEMO_GYM_LITELLM_RESULTS="$RESULTS_DIR" +NEMO_GYM_LITELLM_INSTANCE_ID="$("$PYTHON" -c 'from uuid import uuid4; print(uuid4().hex)')" +export NEMO_GYM_LITELLM_INSTANCE_ID +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" >"$RESULTS_DIR/litellm-health.json" 2>/dev/null && [[ -s "$RESULTS_DIR/litellm-runtime.json" ]]; then + ready=true + break + fi + sleep 1 +done +[[ "$ready" == true ]] || die "LiteLLM readiness timed out; see $RESULTS_DIR/litellm.log" + +# Step 4: Save the shared runtime details alongside each condition's results. +# The comparison uses these revisions and configuration fingerprints to check compatibility. +"$PYTHON" - "$RESULTS_DIR" "$GYM_REVISION" "$SWITCHYARD_REVISION" <<'PY' +import json +import pathlib +import sys + +results = pathlib.Path(sys.argv[1]) +runtime = json.loads((results / 'litellm-runtime.json').read_text()) +provenance = {'gym_revision': sys.argv[2], 'switchyard_revision': sys.argv[3], 'runtime': runtime} +for route in ('fixed', 'routed'): + run_dir = results / route + (run_dir / 'model-calls').mkdir(parents=True, exist_ok=True) + (run_dir / 'run-provenance.json').write_text(json.dumps(provenance, indent=2) + '\n', encoding='utf-8') +PY + +# 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 request ledgers. +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 4096 \ + "++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" + # Keep request evidence even when Gym collection fails. + [[ -s "$run_dir/litellm-calls.jsonl" ]] || die "missing LiteLLM request evidence for $route" + [[ ! -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/fixed" "$RESULTS_DIR/routed" 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 index 103d3334f..b5f2f497c 100644 --- a/tests/test_nemo_gym_compare.py +++ b/tests/test_nemo_gym_compare.py @@ -5,10 +5,6 @@ import importlib.util import json -import os -import re -import shutil -import subprocess from copy import deepcopy from pathlib import Path from types import ModuleType @@ -16,16 +12,14 @@ import pytest -MISSING = object() -BIG = "nvidia/nemotron-3-super-120b-a12b" -SMALL = "openai/gpt-oss-20b" +BIG = "nvidia_nim/nvidia/nemotron-3-super-120b-a12b" +SMALL = "nvidia_nim/openai/gpt-oss-20b" FILES = { "inputs": "rollouts_materialized_inputs.jsonl", "rows": "rollouts.jsonl", "failures": "rollouts_failures.jsonl", - "condition": "switchyard-condition.json", - "snapshot": "switchyard-stats.json", - "commit": "gym-commit.txt", + "events": "litellm-calls.jsonl", + "provenance": "run-provenance.json", } @@ -33,93 +27,120 @@ def comparator() -> ModuleType: 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 - assert spec.loader is not None + 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 _attempt(route: str, request_id: str, **fields: Any) -> list[dict[str, Any]]: + common = {"route": route, "request_id": request_id, "instance_id": "fixture-proxy"} + return [ + {**common, "event": "start"}, + { + **common, + "event": "finish", + "status_code": 200, + "error_type": None, + "response_id": request_id, + "response_status": "completed", + "selected_model": BIG, + "deployment_model": BIG, + "tokens_total": 30, + "routing_ms": 5, + **fields, + }, + ] + + @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 = [], [] + inputs, rows, events = [], [], [] for index in range(2): task = { "_ng_task_index": index, "_ng_rollout_index": 0, "expected_answer": "B", - "grading_mode": "strict_single_letter_boxed", - "agent_ref": {"name": "mcqa_simple_agent"}, + "agent_ref": {"name": "mmlu-redux_mcqa_simple_agent"}, "responses_create_params": { "input": [{"role": "user", "content": f"Question {index}"}], "temperature": 0, - "max_output_tokens": 512, + "max_output_tokens": 4096, }, } inputs.append(task) - response = { - "status": "completed", - "model": SMALL if route == "routed" and index == 0 else BIG, - "output": [ - { - "type": "message", - "role": "assistant", + response_id = f"{route}-{index}" + tokens = 30 + 10 * index + model = SMALL if route == "routed" and index == 0 else BIG + rows.append( + { + **deepcopy(task), + "response": { + "id": response_id, "status": "completed", - "content": [ + "model": route, + "output": [ { - "type": "output_text", - "text": "\\boxed{B}" if index == 0 else "\\boxed{A}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "\\boxed{B}" if index == 0 else "\\boxed{A}", + } + ], } ], - } - ], - } - rows.append( - { - **deepcopy(task), - "response": dict(deepcopy(response), model=route), + }, "reward": 1 - index, "ng_perf": {"total_latency_ms": 100 + 200 * index}, "ng_model_call_capture": { "gaps": [], "calls": [ { + "response_id": response_id, "status_code": 200, "error_category": None, "response_status": "completed", - "model": response["model"], - "tokens_total": 30 + 10 * index, + "model": route, + "tokens_total": tokens, } ], }, } ) + events.extend( + _attempt( + route, + response_id, + tokens_total=tokens, + selected_model=model, + deployment_model=model, + ) + ) runs[route] = { "inputs": inputs, "rows": rows, "failures": [], - "commit": "3a26c35fa90c243427378569511f7b06f503e0fd", - "condition": { - "route": route, - "mode": "hosted", - "nemo_switchyard_version": "0.2.0", - "deployment_sha256": "a" * 64, - }, - "snapshot": { - "mode": "hosted", - "scope": "this run (proxy hosted for exactly this run)", - "stats": { - "total_errors": 0, - "total_requests": 2, - "total_tokens": {"total": 70}, - "routing_overhead": {"avg_ms": 5}, - "classifier": { - "total_errors": 0, - "total_requests": 2 if route == "routed" else 0, - "total_tokens": {"total": 11 if route == "routed" else 0}, - }, + "events": events, + "provenance": { + "gym_revision": "b" * 40, + "switchyard_revision": "c" * 40, + "runtime": { + "mode": "litellm_libsy", + "instance_id": "fixture-proxy", + "litellm_version": "1.97.0", + "switchyard_version": "0.2.0", + "routing_plugin": "switchyard_litellm.RandomRoutingPlugin", + "models": {"fixed": [BIG], "routed": [BIG, SMALL]}, + "profile_sha256": "a" * 64, + "routing_sha256": "a" * 64, + "callback_sha256": "a" * 64, + "provider_base_sha256": "a" * 64, }, }, } @@ -132,262 +153,157 @@ def _write_runs(tmp_path: Path, artifacts: dict[str, dict[str, Any]]) -> list[st directory.mkdir() for name, filename in FILES.items(): value = run[name] - if filename.endswith(".jsonl"): - text = "".join(json.dumps(row) + "\n" for row in value) - else: - text = value + "\n" if name == "commit" else json.dumps(value) + text = ( + "".join(json.dumps(row) + "\n" for row in value) + if filename.endswith(".jsonl") + else json.dumps(value) + ) (directory / filename).write_text(text, encoding="utf-8") return [str(tmp_path / route) for route in ("fixed", "routed")] -def test_reordered_pairing_uses_captured_models_and_separate_classifier_tokens( +def test_complete_reordered_pair( tmp_path: Path, comparator: ModuleType, - artifacts: dict, + artifacts: dict[str, dict[str, Any]], capsys: pytest.CaptureFixture[str], ) -> None: artifacts["routed"]["rows"].reverse() artifacts["routed"]["inputs"].reverse() + artifacts["routed"]["events"].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 selected models: {{"{BIG}": 2}}' in output.out assert f"routed selected models: {json.dumps({SMALL: 1, BIG: 1}, sort_keys=True)}" in output.out - expected = { + for metric, values in { "Mean reward": ["0.500", "0.500"], - "Selected-model tokens": ["70", "70"], - "Classifier tokens": ["0", "11"], - "Combined reported tokens": ["70", "81"], + "Terminal-answer tokens": ["70", "70"], + "Gateway-reported tokens": ["70", "70"], "Mean rollout latency (ms)": ["200", "200"], - "Mean routing overhead (ms)": ["5", "5"], - "Classifier requests": ["0", "2"], - } - for metric, values in expected.items(): - line = next(line for line in output.out.splitlines() if line.startswith(metric)) - assert line[len(metric) :].split() == values + "Gateway requests": ["2", "2"], + "Gateway errors": ["0", "0"], + }.items(): + line = next(line for line in output.out.splitlines() if line[:29].rstrip() == metric) + assert line[29:].split() == values + assert "Classifier tokens: N/A" in output.out + assert "not exhaustive provider-attempt" in output.out @pytest.mark.parametrize( - ("path", "value"), + "problem", [ - ("routed.rows", []), - ( - "routed.failures", - [{"_ng_task_index": 0, "_ng_rollout_index": 0, "_ng_failure_class": "timeout"}], - ), - ("routed.rows.1._ng_task_index", 0), - ("routed.rows.0._ng_task_index", -1), - ("routed.rows.0._ng_task_index", True), - ("routed.rows.0._ng_rollout_index", "0"), - ("routed.rows.0._ng_rollout_index", MISSING), - ("routed.rows.0._ng_task_index", 99), - ("routed.inputs.1._ng_task_index", 0), - ("routed.inputs.0.expected_answer", "A"), - ("routed.inputs.0.responses_create_params.input.0.content", "Different question"), - ("routed.inputs.0.responses_create_params.temperature", 1), - ("routed.inputs.0.responses_create_params.max_output_tokens", 256), - ("routed.condition.deployment_sha256", "b" * 64), - ("routed.condition.deployment_sha256", MISSING), - ("routed.condition.nemo_switchyard_version", "0.3.0"), - ("routed.commit", "different-commit"), - ("routed.condition.route", "fixed"), - ("routed.condition.mode", "external"), - ("routed.snapshot.mode", "external"), - ("routed.snapshot.scope", "all runs"), - ("routed.condition", None), - ("routed.snapshot", []), - ("routed.rows.0._ng_failure_class", "timeout"), - ("routed.rows.0.reward", None), - ("routed.rows.0.reward", 1.1), - ("routed.rows.0.response", MISSING), - ("routed.rows.0.response", None), - ("routed.rows.0.response.status", "incomplete"), - ("routed.rows.0.response.output", []), - ("routed.rows.0.response.output", None), - ("routed.rows.0.response.output", [None]), - ("routed.rows.0.response.output.0.role", "user"), - ("routed.rows.0.response.output.0.content", None), - ("routed.rows.0.response.output.0.content", [None]), - ("routed.rows.0.response.output.0.content.0.text", " "), - ("routed.rows.0.ng_model_call_capture", MISSING), - ("routed.rows.0.ng_model_call_capture", None), - ("routed.rows.0.ng_model_call_capture.gaps", ["missing call"]), - ("routed.rows.0.ng_model_call_capture.calls", []), - ("routed.rows.0.ng_model_call_capture.calls", [{}, {}]), - ("routed.rows.0.ng_model_call_capture.calls", [None]), - ("routed.rows.0.ng_model_call_capture.calls.0.status_code", 500), - ("routed.rows.0.ng_model_call_capture.calls.0.error_category", "timeout"), - ("routed.rows.0.ng_model_call_capture.calls.0.response_status", None), - ("routed.rows.0.ng_model_call_capture.calls.0.response_status", "incomplete"), - ("routed.rows.0.ng_model_call_capture.calls.0.model", "routed"), - ("routed.rows.0.ng_model_call_capture.calls.0.model", MISSING), - ("routed.rows.0.ng_model_call_capture.calls.0.tokens_total", None), - ("routed.rows.0.ng_perf.total_latency_ms", None), - ("routed.snapshot.stats.total_requests", 3), - ("routed.snapshot.stats.total_tokens.total", 81), - ("routed.snapshot.stats.classifier.total_requests", 0), - ("fixed.snapshot.stats.classifier.total_requests", 2), - ("routed.snapshot.stats.classifier.total_tokens.total", None), + "missing", + "duplicate", + "input", + "provenance", + "capture", + "usage", + "failure", + "ledger_gap", + "instance", + "selection", + "gateway_response", ], ) -def test_rejects_incompatible_or_malformed_artifacts( +def test_invalid_evidence_never_prints_averages( tmp_path: Path, comparator: ModuleType, - artifacts: dict, + artifacts: dict[str, dict[str, Any]], capsys: pytest.CaptureFixture[str], - path: str, - value: Any, + problem: str, ) -> None: - keys = [int(key) if key.isdecimal() else key for key in path.split(".")] - target = artifacts - for key in keys[:-1]: - target = target[key] - if value is MISSING: - del target[keys[-1]] + routed = artifacts["routed"] + if problem == "missing": + for run in artifacts.values(): + run["rows"].pop() + elif problem == "duplicate": + routed["rows"].append(deepcopy(routed["rows"][0])) + elif problem == "input": + routed["inputs"][0]["expected_answer"] = "A" + elif problem == "provenance": + routed["provenance"]["gym_revision"] = "d" * 40 + elif problem == "capture": + routed["rows"][0]["ng_model_call_capture"]["gaps"] = ["missing exchange"] + elif problem == "usage": + routed["events"][1]["tokens_total"] = 29 + elif problem == "failure": + routed["failures"] = [{"_ng_task_index": 0, "_ng_rollout_index": 0}] + elif problem == "ledger_gap": + routed["events"].pop() + elif problem == "instance": + routed["events"][1]["instance_id"] = "another-proxy" + elif problem == "selection": + routed["events"][1]["deployment_model"] = BIG else: - target[keys[-1]] = value + routed["events"][1]["response_id"] = "not-the-final-response" assert comparator.main(_write_runs(tmp_path, artifacts)) == 1 output = capsys.readouterr() assert "Cannot compare:" in output.err assert "Mean reward" not in output.out -@pytest.mark.parametrize("side", ["fixed", "routed"]) -@pytest.mark.parametrize("classifier", [False, True]) -@pytest.mark.parametrize("value", [MISSING, None, 1]) -def test_error_counts_must_be_present_and_zero( +def test_recovery_keeps_extra_work_and_terminal_attribution( tmp_path: Path, comparator: ModuleType, - artifacts: dict, - side: str, - classifier: bool, - value: Any, -) -> None: - stats = artifacts[side]["snapshot"]["stats"] - target = stats["classifier"] if classifier else stats - if value is MISSING: - del target["total_errors"] - else: - target["total_errors"] = value - assert comparator.main(_write_runs(tmp_path, artifacts)) == 1 - - -@pytest.mark.parametrize("sides", [("fixed",), ("routed",), ("fixed", "routed")]) -def test_missing_expected_row_rejects_even_when_both_sides_match( - tmp_path: Path, - comparator: ModuleType, - artifacts: dict, + artifacts: dict[str, dict[str, Any]], capsys: pytest.CaptureFixture[str], - sides: tuple, ) -> None: - for side in sides: - artifacts[side]["rows"].pop() - assert comparator.main(_write_runs(tmp_path, artifacts)) == 1 - output = capsys.readouterr() - assert "missing=1" in output.out - assert "Incomplete runs" in output.err - assert "Mean reward" not in output.out + run = artifacts["routed"] + calls = run["rows"][0]["ng_model_call_capture"]["calls"] + calls.insert(0, {"status_code": 503, "error_category": "upstream", "tokens_total": None}) + calls.append({**calls[1], "response_id": "superseded", "tokens_total": 20}) + run["events"].extend( + _attempt( + "routed", + "failed", + status_code=503, + error_type="ServiceUnavailable", + response_id=None, + response_status=None, + tokens_total=None, + routing_ms=None, + ) + ) + run["events"].extend(_attempt("routed", "superseded", tokens_total=20)) + assert comparator.main(_write_runs(tmp_path, artifacts)) == 0 + output = capsys.readouterr().out + for metric, values in { + "Terminal-answer tokens": ["70", "70"], + "Gateway-reported tokens": ["70", "90"], + "Gateway requests": ["2", "4"], + "Gateway errors": ["0", "1"], + "Gateway requests w/o usage": ["0", "1"], + "Captured failed attempts": ["0", "1"], + }.items(): + line = next(line for line in output.splitlines() if line[:29].rstrip() == metric) + assert line[29:].split() == values + assert "WARNING: routed" in output @pytest.mark.parametrize( - ("filename", "hint"), - [ - ("switchyard-condition.json", "when the model server starts"), - ("switchyard-stats.json", "press Ctrl-C"), - ], + "problem", ["truncated", "gateway_truncated", "ambiguous", "missing_ledger"] ) -def test_missing_server_artifacts_explain_the_required_step( +def test_missing_terminal_or_snapshot_has_actionable_error( tmp_path: Path, comparator: ModuleType, - artifacts: dict, + artifacts: dict[str, dict[str, Any]], capsys: pytest.CaptureFixture[str], - filename: str, - hint: str, + problem: str, ) -> None: + row = artifacts["routed"]["rows"][0] + if problem == "truncated": + row["response"]["status"] = "incomplete" + elif problem == "gateway_truncated": + artifacts["routed"]["events"][1]["response_status"] = "incomplete" + elif problem == "ambiguous": + row["ng_model_call_capture"]["calls"] *= 2 args = _write_runs(tmp_path, artifacts) - missing = Path(args[1]) / filename - missing.unlink() + if problem == "missing_ledger": + (Path(args[1]) / FILES["events"]).unlink() assert comparator.main(args) == 1 output = capsys.readouterr() - assert str(missing) in output.err - assert hint in output.err assert "Mean reward" not in output.out - assert not missing.exists() - - -def test_misfiled_manifest_explains_route_directory_mismatch( - tmp_path: Path, - comparator: ModuleType, - artifacts: dict, - capsys: pytest.CaptureFixture[str], -) -> None: - artifacts["fixed"]["condition"]["route"] = "routed" - args = _write_runs(tmp_path, artifacts) - assert comparator.main(args) == 1 - output = capsys.readouterr() - assert "fixed: manifest records route 'routed'" in output.err - assert "fresh directory" in output.err - assert "Mean reward" not in output.out - manifest = json.loads((Path(args[0]) / "switchyard-condition.json").read_text()) - assert manifest["route"] == "routed" - - -@pytest.mark.skipif(shutil.which("bash") is None, reason="Bash is not installed") -@pytest.mark.parametrize("block_index", [0, 1]) -def test_routed_commands_reset_a_stale_fixed_output_directory( - tmp_path: Path, - block_index: int, -) -> None: - readme = Path(__file__).resolve().parents[1] / "benchmark/nemo_gym/README.md" - section = readme.read_text(encoding="utf-8").split("## 4. Repeat with routing", 1)[1] - section = section.split("## 5.", 1)[0] - blocks = re.findall(r"```bash\n(.*?)\n```", section, re.DOTALL) - assert len(blocks) == 2 - run_dir = tmp_path / "run" - run_dir.mkdir() - result = subprocess.run( - [ - "bash", - "-c", - "gym() { printf '%s\\n' \"$@\"; }; git() { printf 'test-revision\\n'; };\n" - + blocks[block_index], - ], - env={ - "PATH": os.defpath, - "EXAMPLE": str(tmp_path / "example"), - "WORK": str(tmp_path / "work"), - "RUN_DIR": str(run_dir), - "OUT": str(run_dir / "fixed"), - "ROUTE": "fixed", - }, - text=True, - capture_output=True, - check=False, - timeout=10, - ) - assert result.returncode == 0, result.stderr - args = result.stdout.splitlines() - assert f"++model_call_capture_dir={run_dir}/routed/model-calls" in args - if block_index == 0: - assert args[args.index("--model") + 1] == "routed" - assert ( - f"++policy_model.responses_api_models.switchyard_model.condition_dir={run_dir}/routed" - in args - ) - assert (run_dir / "routed/gym-commit.txt").read_text() == "test-revision\n" - else: - assert args[args.index("--output") + 1] == str(run_dir / "routed/rollouts.jsonl") - assert not (run_dir / "fixed").exists() - - -@pytest.mark.skipif(shutil.which("bash") is None, reason="Bash is not installed") -def test_readme_shell_blocks_have_valid_bash_syntax() -> None: - readme = Path(__file__).resolve().parents[1] / "benchmark/nemo_gym/README.md" - blocks = re.findall(r"```bash\n(.*?)\n```", readme.read_text(encoding="utf-8"), re.DOTALL) - assert blocks, "The tutorial must contain executable Bash examples" - for index, block in enumerate(blocks, start=1): - result = subprocess.run( - ["bash", "-n"], input=block, text=True, capture_output=True, check=False, timeout=10 - ) - assert result.returncode == 0, f"Bash example {index}: {result.stderr}" + if problem == "missing_ledger": + assert "litellm.log" in output.err diff --git a/tests/test_nemo_gym_litellm.py b/tests/test_nemo_gym_litellm.py new file mode 100644 index 000000000..bd4f1062f --- /dev/null +++ b/tests/test_nemo_gym_litellm.py @@ -0,0 +1,356 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import importlib.metadata +import importlib.util +import json +import math +from copy import deepcopy +from pathlib import Path +from types import ModuleType +from typing import Any + +import pytest + +litellm = pytest.importorskip("litellm") +pytest.importorskip("switchyard_litellm") + +from litellm import ResponsesAPIResponse # noqa: E402 +from litellm.types.router import RoutingContext # noqa: E402 + +BIG = "nvidia_nim/nvidia/nemotron-3-super-120b-a12b" +SMALL = "nvidia_nim/openai/gpt-oss-20b" +ROOT = Path(__file__).resolve().parents[1] +CALLBACK = ROOT / "benchmark/nemo_gym/gym_routing_plugin.py" +PROFILE = ROOT / "benchmark/nemo_gym/litellm.yaml" + + +@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 _plugin(callback: ModuleType, tmp_path: Path) -> Any: + routing = tmp_path / "routes.toml" + routing.write_text('algorithm = "random"\nseed = 6\n', encoding="utf-8") + return callback.GymRoutingPlugin(routing, PROFILE, tmp_path / "results", "fixture-proxy") + + +def _context(candidates: list[str], text: str) -> RoutingContext: + messages = [{"role": "user", "content": text}] + return RoutingContext( + raw_messages=messages, + structured_messages=messages, + candidate_models=candidates, + metadata={"model_group": "routed"}, + ) + + +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"]) +@pytest.mark.parametrize("as_dict", [False, True]) +def test_response_payload_normalizes_litellm_reasoning( + callback: ModuleType, finish_reason: str, as_dict: bool +) -> 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 = deepcopy(response.model_dump()) + source = deepcopy(before) if as_dict else response + payload = callback.response_payload(source) + 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 callback.response_payload(payload) == payload + assert (source if as_dict else source.model_dump()) == before + + +@pytest.mark.parametrize("summary", [None, [{"type": "summary_text", "text": "Existing summary"}]]) +def test_response_payload_preserves_native_reasoning_and_other_items( + callback: ModuleType, summary: Any +) -> None: + source = { + "id": "response-native", + "status": "incomplete", + "usage": None, + "output": [ + { + "type": "reasoning", + "id": "reasoning-native", + "summary": summary, + "encrypted_content": "opaque-test-content", + "content": [{"type": "reasoning_text", "text": "Existing reasoning"}], + }, + {"type": "message", "content": [{"type": "output_text", "text": "Answer"}]}, + {"type": "function_call", "name": "test_tool", "arguments": "{}"}, + ], + } + before = deepcopy(source) + expected = deepcopy(source) + if summary is None: + expected["output"][0]["summary"] = [] + assert callback.response_payload(source) == expected + assert source == before + + +async def test_real_libsy_random_routing_reaches_both_candidates( + callback: ModuleType, tmp_path: Path +) -> None: + plugin = _plugin(callback, tmp_path) + + first = await plugin.run(_context([BIG, SMALL], "first")) + second = await plugin.run(_context([BIG, SMALL], "second")) + fixed = await plugin.run(_context([BIG], "fixed")) + + selections = { + first.signals["switchyard"]["selected_model_id"], + second.signals["switchyard"]["selected_model_id"], + } + assert selections == {BIG, SMALL} + assert fixed.signals["switchyard"]["selected_model_id"] == BIG + for result in (first, second, fixed): + routing_ms = result.signals["switchyard"]["routing_ms"] + assert isinstance(routing_ms, float) and math.isfinite(routing_ms) and routing_ms >= 0 + + +def test_litellm_local_file_loader_shares_runner_instance_id( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + from litellm.proxy.types_utils.utils import get_instance_fn + + routing = tmp_path / "routes.toml" + routing.write_text('algorithm = "random"\nseed = 6\n', encoding="utf-8") + results = tmp_path / "results" + monkeypatch.setenv("SWITCHYARD_LITELLM_CONFIG", str(routing)) + monkeypatch.setenv("NEMO_GYM_LITELLM_PROFILE", str(PROFILE)) + monkeypatch.setenv("NEMO_GYM_LITELLM_RESULTS", str(results)) + monkeypatch.setenv("NEMO_GYM_LITELLM_INSTANCE_ID", "fixture-shared-run") + + first = get_instance_fn(value="gym_routing_plugin.PLUGIN", config_file_path=str(PROFILE)) + second = get_instance_fn(value="gym_routing_plugin.PLUGIN", config_file_path=str(PROFILE)) + + assert first is not second + assert first.runtime["instance_id"] == "fixture-shared-run" + assert second.runtime["instance_id"] == "fixture-shared-run" + runtime = json.loads((results / "litellm-runtime.json").read_text()) + assert runtime["instance_id"] == "fixture-shared-run" + data = {"model": "fixed", "litellm_call_id": "fixture-request"} + first._record(data, {"event": "start"}) + second._record(data, {"event": "finish"}) + ledger = [ + json.loads(line) + for line in (results / "fixed/litellm-calls.jsonl").read_text().splitlines() + ] + assert [record["instance_id"] for record in ledger] == [ + "fixture-shared-run", + "fixture-shared-run", + ] + + +async def test_callbacks_record_allowlisted_success_and_runtime( + callback: ModuleType, tmp_path: Path +) -> None: + plugin = _plugin(callback, tmp_path) + selected = SMALL + data = { + "model": "routed", + "litellm_call_id": "request-1", + "messages": [{"role": "user", "content": "secret-prompt-marker"}], + "headers": {"Authorization": "secret-header-marker"}, + "litellm_metadata": { + "deployment": selected, + "routing_plugin_signals": { + "switchyard": {"selected_model_id": selected, "routing_ms": 1.5} + }, + }, + } + response = _response( + status="completed", + usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + ) + + await plugin.async_pre_call_hook(None, None, data, "aresponses") + payload = await plugin.async_post_call_success_hook(data, None, response) + + records_path = tmp_path / "results/routed/litellm-calls.jsonl" + records = [json.loads(line) for line in records_path.read_text().splitlines()] + assert records == [ + { + "instance_id": plugin.runtime["instance_id"], + "route": "routed", + "request_id": "request-1", + "event": "start", + }, + { + "instance_id": plugin.runtime["instance_id"], + "route": "routed", + "request_id": "request-1", + "event": "finish", + "status_code": 200, + "error_type": None, + "response_id": "response-1", + "response_status": "completed", + "selected_model": selected, + "deployment_model": selected, + "tokens_total": 15, + "routing_ms": 1.5, + }, + ] + assert payload["id"] == "response-1" + assert payload["status"] == "completed" + assert payload["usage"]["total_tokens"] == 15 + assert payload["_hidden_params"] == response._hidden_params + serialized = records_path.read_text() + assert "secret-prompt-marker" not in serialized + assert "secret-header-marker" not in serialized + + runtime = json.loads((tmp_path / "results/litellm-runtime.json").read_text()) + assert runtime["profile_sha256"] == hashlib.sha256(PROFILE.read_bytes()).hexdigest() + assert runtime["callback_sha256"] == hashlib.sha256(CALLBACK.read_bytes()).hexdigest() + assert runtime["litellm_version"] == importlib.metadata.version("litellm") + assert runtime["switchyard_version"] == importlib.metadata.version("nemo-switchyard") + assert runtime["routing_plugin"] == "switchyard_litellm.RandomRoutingPlugin" + + +async def test_failure_hook_records_unknown_usage_without_secrets( + callback: ModuleType, tmp_path: Path +) -> None: + plugin = _plugin(callback, tmp_path) + selected = BIG + data = { + "model": "fixed", + "litellm_call_id": "request-failed", + "messages": [{"role": "user", "content": "secret-failure-message"}], + "headers": {"Authorization": "secret-failure-header"}, + "litellm_metadata": { + "deployment": selected, + "routing_plugin_signals": { + "switchyard": {"selected_model_id": selected, "routing_ms": 0.5} + }, + }, + } + + class ServiceUnavailable(Exception): + status_code = 503 + + await plugin.async_pre_call_hook(None, None, data, "aresponses") + await plugin.async_post_call_failure_hook(data, ServiceUnavailable("secret-error"), None) + + records_path = tmp_path / "results/fixed/litellm-calls.jsonl" + records = [json.loads(line) for line in records_path.read_text().splitlines()] + assert len(records) == 2 + finish = records[1] + assert finish["event"] == "finish" + assert finish["status_code"] == 503 + assert finish["error_type"] == "ServiceUnavailable" + assert finish["tokens_total"] is None + assert finish["response_id"] is None + serialized = records_path.read_text() + assert "secret-failure-message" not in serialized + assert "secret-failure-header" not in serialized + assert "secret-error" not in serialized diff --git a/tests/test_nemo_gym_run.py b/tests/test_nemo_gym_run.py new file mode 100644 index 000000000..9fe0c641e --- /dev/null +++ b/tests/test_nemo_gym_run.py @@ -0,0 +1,454 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +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" +if [[ "${FAKE_GYM_BEHAVIOR:-}" != "missing_ledger" ]]; then + printf '{"event":"start"}\\n{"event":"finish"}\\n' > "$run_dir/litellm-calls.jsonl" +fi +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 +set -u +printf 'uv:%s\\n' "$*" >> "$EVENT_LOG" +"$REAL_PYTHON" - "$NEMO_GYM_LITELLM_RESULTS" "$NEMO_GYM_LITELLM_PROFILE" \ + "$SWITCHYARD_LITELLM_CONFIG" "$FAKE_CALLBACK_PATH" "$NVIDIA_BASE_URL" <<'PY' +import hashlib +import json +import os +import pathlib +import sys + +results, profile, routing, callback = map(pathlib.Path, sys.argv[1:5]) +runtime = { + "mode": "litellm_libsy", + "instance_id": os.environ["NEMO_GYM_LITELLM_INSTANCE_ID"], + "litellm_version": "1.97.0", + "switchyard_version": "0.2.0", + "fastapi_version": "0.136.3", + "starlette_version": "1.3.1", + "routing_plugin": "switchyard_litellm.RandomRoutingPlugin", + "models": { + "fixed": ["nvidia_nim/nvidia/nemotron-3-super-120b-a12b"], + "routed": [ + "nvidia_nim/nvidia/nemotron-3-super-120b-a12b", + "nvidia_nim/openai/gpt-oss-20b", + ], + }, + "profile_sha256": hashlib.sha256(profile.read_bytes()).hexdigest(), + "routing_sha256": hashlib.sha256(routing.read_bytes()).hexdigest(), + "callback_sha256": hashlib.sha256(callback.read_bytes()).hexdigest(), + "provider_base_sha256": hashlib.sha256(sys.argv[5].encode()).hexdigest(), +} +results.mkdir(parents=True, exist_ok=True) +(results / "litellm-runtime.json").write_text(json.dumps(runtime), encoding="utf-8") +PY +exec "$REAL_PYTHON" -c ' +import os +import signal +import sys + +arguments = " ".join(sys.argv[1:]) + +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()}:{arguments}\\n") +signal.pause() +' "$@" +""", + ) + _write_executable( + bin_dir / "curl", + """#!/bin/bash +url="" +for argument in "$@"; do case "$argument" in http://*) url="$argument" ;; esac; done +printf 'curl:%s\\n' "$url" >> "$EVENT_LOG" +[[ -s "$NEMO_GYM_LITELLM_RESULTS/litellm-runtime.json" ]] || exit 22 +printf '{}\\n' +""", + ) + _write_executable( + bin_dir / "git", + """#!/bin/bash +printf 'git:%s\\n' "$*" >> "$EVENT_LOG" +case "$*" in + *" status "*) exit 0 ;; + *" rev-parse "*) printf '%039d1\\n' 0 ;; + *" describe "*) printf '%040d-dirty\\n' 2 ;; +esac +""", + ) + for tool in ("cargo", "switchyard-server"): + _write_executable( + bin_dir / tool, + f'#!/bin/bash\nprintf \'forbidden:{tool}:%s\\n\' "$*" >> "$EVENT_LOG"\nexit 99\n', + ) + + env = os.environ.copy() + for key in ( + "NVIDIA_API_KEY", + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "OPENROUTER_API_KEY", + "NVIDIA_BASE_URL", + "LITELLM_CONFIG", + "SWITCHYARD_CONFIG", + "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"), + "FAKE_CALLBACK_PATH": str(ROOT / "benchmark/nemo_gym/gym_routing_plugin.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", + "LITELLM_CONFIG", + "SWITCHYARD_CONFIG", + "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 + + +@pytest.mark.parametrize("invalid", ["existing-results", "zero-limit", "bad-port"]) +def test_preflight_rejects_before_operations( + fake_runner_env: tuple[dict[str, str], Path, Path], invalid: str +) -> None: + env, results_dir, event_log = fake_runner_env + if invalid == "existing-results": + results_dir.mkdir() + sentinel = results_dir / "sentinel" + sentinel.write_text("keep", encoding="utf-8") + else: + sentinel = None + env["LIMIT" if invalid == "zero-limit" else "LITELLM_PORT"] = "0" + result = _run(env) + assert result.returncode != 0 + assert _events(event_log) == [] + if sentinel is not None: + 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", + "NEMO_GYM_LITELLM_INSTANCE_ID": "must-not-be-reused", + } + ) + result = _run(env) + assert result.returncode == 0, result.stderr + events = _events(event_log) + + preparations = [line for line in events if line.startswith("prepare:")] + assert len(preparations) == 1 + assert "eval prepare --benchmark mmlu-redux" in preparations[0] + starts = [line for line in events if line.startswith("proxy_start:")] + assert len(starts) == 1 + proxy_pid = int(starts[0].split(":", 2)[1]) + for pin in ( + "litellm[proxy]==1.97.0", + "fastapi==0.136.3", + "starlette==1.3.1", + "--num_workers 1", + ): + assert pin in starts[0] + 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 route, invocation in zip(("fixed", "routed"), gym_runs, strict=True): + run_dir = results_dir / route + for expected in ( + "--benchmark mmlu-redux", + "--model-type litellm_model", + f"--model {route}", + "--split benchmark", + "--limit 2", + "--num-repeats 3", + "--concurrency 1", + "--temperature 0", + "--max-output-tokens 4096", + "++policy_base_url=http://127.0.0.1:4000/v1", + "++policy_api_key=unused", + f"++model_call_capture_dir={run_dir}/model-calls", + f"++nemo_gym_log_dir={run_dir}/server-logs", + "++mcqa_simple_agent.responses_api_agents.simple_agent.max_steps=1", + ): + assert expected in invocation + assert "switchyard_model" not in invocation + assert ".deployment=" not in invocation + assert "condition_dir" not in invocation + assert ( + "++mmlu-redux_mcqa_simple_agent.responses_api_agents.simple_agent.max_steps=" + not in invocation + ) + assert not any(line.startswith("forbidden:") for line in events) + + provenances = [ + json.loads((results_dir / route / "run-provenance.json").read_text()) + for route in ("fixed", "routed") + ] + assert provenances[0] == provenances[1] + runtime = provenances[0]["runtime"] + assert len(runtime["instance_id"]) == 32 + assert all(character in "0123456789abcdef" for character in runtime["instance_id"]) + assert runtime["instance_id"] != "must-not-be-reused" + assert runtime["models"] == { + "fixed": ["nvidia_nim/nvidia/nemotron-3-super-120b-a12b"], + "routed": [ + "nvidia_nim/nvidia/nemotron-3-super-120b-a12b", + "nvidia_nim/openai/gpt-oss-20b", + ], + } + 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 / "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 + + +def test_missing_ledger_stops_before_comparison( + fake_runner_env: tuple[dict[str, str], Path, Path], +) -> None: + env, _, event_log = fake_runner_env + env["FAKE_GYM_BEHAVIOR"] = "missing_ledger" + result = _run(env) + events = _events(event_log) + assert result.returncode != 0 + assert "missing LiteLLM request evidence for fixed" in result.stderr + 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 + + +@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) From 2b9a742c68bada9be50f8d8e7049d4ea91f7004b Mon Sep 17 00:00:00 2001 From: Shashank Verma Date: Mon, 14 Sep 2026 15:23:19 -0700 Subject: [PATCH 3/6] docs(tutorial): clarify evaluation diagram labels Distinguish the MMLU-Redux dataset from its multiple-choice verifier. Clarify task-session initialization and input/output token accounting. Signed-off-by: Shashank Verma --- benchmark/nemo_gym/architecture.svg | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/benchmark/nemo_gym/architecture.svg b/benchmark/nemo_gym/architecture.svg index be53aff88..f1a7351bc 100644 --- a/benchmark/nemo_gym/architecture.svg +++ b/benchmark/nemo_gym/architecture.svg @@ -1,6 +1,6 @@ Evaluating Switchyard routing with NeMo Gym - The evaluation uses MMLU-Redux tasks, Gym's simple agent and MCQA verifier. Gym's litellm_model calls a separate runner-managed LiteLLM proxy through the OpenAI Responses API. Switchyard libsy Random routing runs directly inside LiteLLM, not as a Switchyard HTTP runtime. Compare fixed and routed runs on identical inputs using rewards, selected models, model tokens, latency and routing statistics. + 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, selected models, input and output tokens, latency and routing statistics. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 @@ -32,11 +32,11 @@ SPDX-License-Identifier: Apache-2.0 Dataset (tasks) - MMLU-Redux + MMLU-Redux 2.0 Initialize - Environment + Task session Agent @@ -47,7 +47,7 @@ SPDX-License-Identifier: Apache-2.0 Verify Resources server - MCQA verifier + Multiple-choice verifier @@ -83,6 +83,6 @@ SPDX-License-Identifier: Apache-2.0 Compare fixed vs. routed runs Per-task rewards Selected models - Model tokens + Input + output tokens Latency and routing statistics From 74b8ca779ceeec8677200b4a8b593344fed6b1df Mon Sep 17 00:00:00 2001 From: Shashank Verma Date: Mon, 14 Sep 2026 16:26:48 -0700 Subject: [PATCH 4/6] test(tutorial): focus comparator coverage on core behaviors Signed-off-by: Shashank Verma --- tests/test_nemo_gym_compare.py | 66 ++++++---------------------------- 1 file changed, 10 insertions(+), 56 deletions(-) diff --git a/tests/test_nemo_gym_compare.py b/tests/test_nemo_gym_compare.py index b5f2f497c..efa0a1bb4 100644 --- a/tests/test_nemo_gym_compare.py +++ b/tests/test_nemo_gym_compare.py @@ -192,19 +192,12 @@ def test_complete_reordered_pair( @pytest.mark.parametrize( - "problem", + ("problem", "expected_error"), [ - "missing", - "duplicate", - "input", - "provenance", - "capture", - "usage", - "failure", - "ledger_gap", - "instance", - "selection", - "gateway_response", + ("incomplete", "Incomplete runs"), + ("mismatched", "Task inputs, verifier metadata, or generation settings differ"), + ("capture", "incomplete model-call capture"), + ("truncated", "missing or incomplete final answer"), ], ) def test_invalid_evidence_never_prints_averages( @@ -213,34 +206,22 @@ def test_invalid_evidence_never_prints_averages( artifacts: dict[str, dict[str, Any]], capsys: pytest.CaptureFixture[str], problem: str, + expected_error: str, ) -> None: routed = artifacts["routed"] - if problem == "missing": + if problem == "incomplete": for run in artifacts.values(): run["rows"].pop() - elif problem == "duplicate": - routed["rows"].append(deepcopy(routed["rows"][0])) - elif problem == "input": + elif problem == "mismatched": routed["inputs"][0]["expected_answer"] = "A" - elif problem == "provenance": - routed["provenance"]["gym_revision"] = "d" * 40 elif problem == "capture": routed["rows"][0]["ng_model_call_capture"]["gaps"] = ["missing exchange"] - elif problem == "usage": - routed["events"][1]["tokens_total"] = 29 - elif problem == "failure": - routed["failures"] = [{"_ng_task_index": 0, "_ng_rollout_index": 0}] - elif problem == "ledger_gap": - routed["events"].pop() - elif problem == "instance": - routed["events"][1]["instance_id"] = "another-proxy" - elif problem == "selection": - routed["events"][1]["deployment_model"] = BIG else: - routed["events"][1]["response_id"] = "not-the-final-response" + routed["rows"][0]["response"]["status"] = "incomplete" 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 @@ -280,30 +261,3 @@ def test_recovery_keeps_extra_work_and_terminal_attribution( line = next(line for line in output.splitlines() if line[:29].rstrip() == metric) assert line[29:].split() == values assert "WARNING: routed" in output - - -@pytest.mark.parametrize( - "problem", ["truncated", "gateway_truncated", "ambiguous", "missing_ledger"] -) -def test_missing_terminal_or_snapshot_has_actionable_error( - tmp_path: Path, - comparator: ModuleType, - artifacts: dict[str, dict[str, Any]], - capsys: pytest.CaptureFixture[str], - problem: str, -) -> None: - row = artifacts["routed"]["rows"][0] - if problem == "truncated": - row["response"]["status"] = "incomplete" - elif problem == "gateway_truncated": - artifacts["routed"]["events"][1]["response_status"] = "incomplete" - elif problem == "ambiguous": - row["ng_model_call_capture"]["calls"] *= 2 - args = _write_runs(tmp_path, artifacts) - if problem == "missing_ledger": - (Path(args[1]) / FILES["events"]).unlink() - assert comparator.main(args) == 1 - output = capsys.readouterr() - assert "Mean reward" not in output.out - if problem == "missing_ledger": - assert "litellm.log" in output.err From 6deb79debe7bd40bf1217ed9b7908ba57cb00b1e Mon Sep 17 00:00:00 2001 From: Shashank Verma Date: Tue, 15 Sep 2026 07:57:08 -0700 Subject: [PATCH 5/6] fix(tutorial): Update tutorial to simplify and user newer models - Simplifies tutorial language, tests, and overall LOC - Updates the models to use Nemotron 3 Ultra and Nemotron 3.5 Lightning - Updates the results captured. Signed-off-by: Shashank Verma --- benchmark/nemo_gym/README.md | 144 +++++------ benchmark/nemo_gym/architecture.svg | 14 +- benchmark/nemo_gym/compare.py | 312 +++++++---------------- benchmark/nemo_gym/gym_routing_plugin.py | 157 ++---------- benchmark/nemo_gym/litellm.yaml | 14 +- benchmark/nemo_gym/run.sh | 46 +--- tests/test_nemo_gym_compare.py | 219 +++++++--------- tests/test_nemo_gym_litellm.py | 232 ++--------------- tests/test_nemo_gym_run.py | 167 ++---------- 9 files changed, 352 insertions(+), 953 deletions(-) diff --git a/benchmark/nemo_gym/README.md b/benchmark/nemo_gym/README.md index 4466bbb1a..dce33571f 100644 --- a/benchmark/nemo_gym/README.md +++ b/benchmark/nemo_gym/README.md @@ -2,55 +2,14 @@ [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. -Gym provides the evaluation substrate: it supplies tasks, runs the agent, and verifies answers to report rewards. -Switchyard sits in the model-request path, selecting which upstream model serves each request. +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) -## Understand the wiring - -Gym calls a LiteLLM endpoint through its `litellm_model` adapter. The Switchyard library is integrated in LiteLLM, which is the path we'll use in this example. As such, Switchyard does not run as a separate server here. - -**LiteLLM defines the model groups.** This excerpt from [litellm.yaml](litellm.yaml) shows the candidates; provider settings are omitted here: - -```yaml -model_list: - - model_name: fixed - litellm_params: {model: nvidia_nim/nvidia/nemotron-3-super-120b-a12b} - - model_name: routed - litellm_params: {model: nvidia_nim/nvidia/nemotron-3-super-120b-a12b} - - model_name: routed - litellm_params: {model: nvidia_nim/openai/gpt-oss-20b} -``` - -`fixed` and `routed` are names we chose, not special Gym modes. The fixed group has one candidate, so it always uses Super. The routed group has two candidates for Switchyard to choose from. - -**Switchyard defines the selection policy.** [routes.toml](routes.toml) contains: - -```toml -algorithm = "random" -seed = 6 -``` - -LiteLLM uses the [Switchyard integration](../../examples/litellm/README.md) to choose a model using `routes.toml`. A small adapter records each choice and makes the response compatible with Gym. - -**Gym requests a group, not a concrete model.** These excerpts show the model wiring inside the runner; **these are wrapped in a [run.sh](./run.sh) script we'll run later, and are not additional steps to execute in this tutorial**: - -```text -gym eval run --benchmark mmlu-redux --model-type litellm_model --model fixed \ - ++policy_base_url=http://127.0.0.1:4000/v1 ++policy_api_key=unused - -gym eval run --benchmark mmlu-redux --model-type litellm_model --model routed \ - ++policy_base_url=http://127.0.0.1:4000/v1 ++policy_api_key=unused -``` -Note the following -- `--model-type` selects the adapter -- `policy_base_url` points it at LiteLLM -- `--model` selects the group. The local proxy uses provider credentials from the environment, not Gym's placeholder key. -- The runner adds identical task limits and separate output/capture paths. The dataset, agent, verifier, temperature, and 4,096-token answer limit stay unchanged. - ## 1. Set up You need: @@ -59,7 +18,7 @@ You need: - 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 `openai/gpt-oss-20b` and `nvidia/nemotron-3-super-120b-a12b` +- 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. @@ -79,43 +38,84 @@ Gym and LiteLLM use separate Python environments. The proxy builds Switchyard bi ## 2. Run both conditions -The default is five tasks per condition, normally ten upstream calls. Inference can consume credits, and retries can add calls. Replace the placeholder below with your NVIDIA API key, then run. Pasting a key into this command may save it in shell history. +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//`. Keep this unauthenticated development proxy local; do not share or publicly expose it. +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//`. -In `run.sh`, one loop runs both conditions: only the model group and output paths change. Leave the configuration and prepared data unchanged until both runs finish. +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. Read the comparison +## 3. Understand the result -Start with coverage and selected models, then compare rewards, tokens, and latency. Here is an excerpt from the **two-task local stub test**, not a real-model benchmark: +Start with task coverage and serving models, then compare rewards, tokens, and latency. This is the output from a live NVIDIA run of the default five tasks per condition: ```text -Pairing: matched=2, fixed-only=0, routed-only=0 +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 2 2 -Mean reward 0.500 0.500 -Terminal-answer tokens 30 30 -Gateway-reported tokens 30 30 - -fixed selected models: {"nvidia_nim/nvidia/nemotron-3-super-120b-a12b": 2} -routed selected models: {"nvidia_nim/nvidia/nemotron-3-super-120b-a12b": 1, "nvidia_nim/openai/gpt-oss-20b": 1} +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:** both conditions completed the same two tasks. Incomplete or invalid evidence is rejected instead of producing partial averages. -- **Selection:** fixed stayed on Super; routed used both models. A small Random run need not split evenly. -- **Reward and usage:** MCQA scores the boxed answer letter (correct = 1, wrong = 0). Token columns sum input and output tokens across rollouts, not just generated answers. The stub supplies answers and token counts, so these values demonstrate the report, not model quality or savings. +- **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. -For real runs, weigh reward against usage and latency rather than treating fewer tokens as a win by itself. 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. +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. -## 4. Try a small change +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: @@ -126,9 +126,9 @@ For real runs, weigh reward against usage and latency rather than treating fewer - **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. This runner and comparator are MMLU-Redux-specific, not a general benchmark launcher. +- **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 profile, port, repeat, and concurrency options. +See `bash benchmark/nemo_gym/run.sh --help` for output, port, repeat, and concurrency options.
Saved files and token counts @@ -138,22 +138,24 @@ 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/` and `litellm-calls.jsonl`: model request logs. -- `run-provenance.json`: version and configuration details. +- `model-calls/`: Gym's captured model exchanges. +- `models.jsonl`: response IDs and serving models. -If something fails, start with that folder's `gym.log` or the result folder's `litellm.log`. +The result root also contains `comparison.txt` and `versions.txt`, recorded for reference. -The "Gateway-reported tokens" column includes the "Terminal-answer tokens", so don't add them together. Missing token counts are unknown, not zero, and some provider retries may not appear in the totals. Random makes no classifier calls; routing time is already included in rollout latency. +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/fixed benchmark/nemo_gym/results/my-run/routed + 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.
-**Tested:** with Gym `v0.6.0`. +**Validation:** tested with Gym `v0.6.0`, and a fresh live NVIDIA 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 index f1a7351bc..4bc1f9fb1 100644 --- a/benchmark/nemo_gym/architecture.svg +++ b/benchmark/nemo_gym/architecture.svg @@ -1,6 +1,6 @@ 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, selected models, input and output tokens, latency and routing statistics. + 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 @@ -70,19 +70,19 @@ SPDX-License-Identifier: Apache-2.0 Upstream models - GPT-OSS 20B + Nemotron 3.5 Lightning - Nemotron 3 Super + Nemotron 3 Ultra - Usage / - routing + Serving + model Output: Compare fixed vs. routed runs Per-task rewards - Selected models + Serving models Input + output tokens - Latency and routing statistics + Latency and captured errors diff --git a/benchmark/nemo_gym/compare.py b/benchmark/nemo_gym/compare.py index 34917812d..584634779 100644 --- a/benchmark/nemo_gym/compare.py +++ b/benchmark/nemo_gym/compare.py @@ -1,14 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Compare complete, paired MCQA runs through LiteLLM and Switchyard Random routing.""" +"""Compare one runner-produced experiment using Gym-captured usage and model attribution.""" from __future__ import annotations import argparse import json import math -import re import sys from collections import Counter from pathlib import Path @@ -31,156 +30,81 @@ def number(value: Any, name: str) -> int | float: return cast(int | float, value) -def count(value: Any, name: str) -> int: - """Read a nonnegative integer counter.""" - require(type(value) is int and value >= 0, f"{name} must be a nonnegative integer") - return cast(int, value) - - -def read_object(path: Path) -> dict[str, Any]: - """Read a JSON artifact with an object at its root.""" - value = json.loads(path.read_text(encoding="utf-8")) - require(isinstance(value, dict), f"{path}: expected a JSON object") - return cast(dict[str, Any], value) - - def read_jsonl(path: Path) -> list[dict[str, Any]]: - """Read JSONL objects, ignoring empty lines.""" - rows = [ - json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip() - ] - require(all(isinstance(row, dict) for row in rows), f"{path}: expected JSON objects") - return rows + """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 by Gym task and repeat, rejecting duplicate or malformed identities.""" - indexed: 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.get("_ng_task_index"), row.get("_ng_rollout_index")) - require( - all(type(part) is int and part >= 0 for part in key), - f"{path}: invalid task/repeat index", - ) - key = cast(tuple[int, int], key) + 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(response: Any) -> bool: - """Require a completed Responses API message, not just reasoning or a tool call.""" - return ( - isinstance(response, dict) - and response.get("status") == "completed" - and any( - item.get("type") == "message" - and item.get("role") == "assistant" - and any( - part.get("type") == "output_text" - and isinstance(part.get("text"), str) - and part["text"].strip() - for part in (item.get("content") or []) - if isinstance(part, dict) - ) - for item in (response.get("output") or []) - if isinstance(item, dict) - ) +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]: - """Require complete gateway request evidence alongside Gym's rollout artifacts.""" + """Load Gym's expected inputs, completed rollouts, and terminal failure sidecar.""" failure_path, rollout_path = path / "rollouts_failures.jsonl", path / "rollouts.jsonl" - for filename in ("run-provenance.json", "litellm-calls.jsonl"): - require( - (path / filename).is_file(), - f"Missing {path / filename}. Inspect {path / 'gym.log'} and {path.parent / 'litellm.log'}.", - ) - events: dict[str, dict[str, dict[str, Any]]] = {"start": {}, "finish": {}} - for event in read_jsonl(path / "litellm-calls.jsonl"): - kind, request_id = event.get("event"), event.get("request_id") - require(kind in events, f"{path}: unknown gateway event") - kind = cast(str, kind) - require(isinstance(request_id, str) and bool(request_id), f"{path}: missing request ID") - request_id = cast(str, request_id) - require(request_id not in events[kind], f"{path}: duplicate gateway {kind} event") - events[kind][request_id] = event - require( - bool(events["start"]) and events["start"].keys() == events["finish"].keys(), - f"{path}: incomplete gateway request evidence", - ) 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 [], - "provenance": read_object(path / "run-provenance.json"), - "events": events, } -def summarize(run: dict[str, Any], route: str) -> tuple[dict[str, int | float], Counter[str]]: - """Join final answers by response ID while retaining all recorded gateway work.""" - runtime = run["provenance"]["runtime"] - allowed_models = runtime["models"][route] - finishes = list(run["events"]["finish"].values()) - responses: dict[str, dict[str, Any]] = {} - tokens = [] - routing_times = [] - errors = unknown_usage = 0 - for phase in run["events"].values(): - for event in phase.values(): - require(event["route"] == route, f"{route}: gateway event belongs to another route") - require( - event["instance_id"] == runtime["instance_id"], - f"{route}: gateway events belong to another proxy instance", - ) - for event in finishes: - failed = event.get("status_code") != 200 or bool(event.get("error_type")) - errors += int(failed) - if not failed: - model, response_id = event.get("selected_model"), event.get("response_id") - require(model in allowed_models, f"{route}: missing or invalid Switchyard selection") - require( - event.get("deployment_model") == model, - f"{route}: selected model differs from the LiteLLM deployment", - ) - require( - isinstance(response_id, str) and bool(response_id), - f"{route}: missing gateway response ID", - ) - require(response_id not in responses, f"{route}: ambiguous gateway response ID") - responses[response_id] = event - number(event.get("tokens_total"), "gateway response tokens") - number(event.get("routing_ms"), "routing decision time") - if event.get("tokens_total") is None: - unknown_usage += 1 - else: - tokens.append(number(event["tokens_total"], "gateway tokens")) - if event.get("routing_ms") is not None: - routing_times.append(number(event["routing_ms"], "routing decision time")) - - calls, rewards, latencies, models = [], [], [], [] - used_response_ids: set[str] = set() - captured_attempts = captured_errors = 0 +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") - require(has_answer(row["response"]), f"{label}: missing or incomplete final answer") - response_id = row["response"].get("id") - require(isinstance(response_id, str) and bool(response_id), f"{label}: missing response id") - require(response_id not in used_response_ids, f"{label}: response reused across rollouts") - used_response_ids.add(response_id) - reward = number(row["reward"], f"{label}: reward") - require(reward <= 1, f"{label}: MCQA reward must be between zero and one") + 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(isinstance(capture, dict), f"{label}: missing model-call capture") require(not capture.get("gaps"), f"{label}: incomplete model-call capture") records = capture["calls"] - require( - isinstance(records, list) and all(isinstance(call, dict) for call in records), - f"{label}: invalid captures", - ) 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] @@ -188,48 +112,38 @@ def summarize(run: dict[str, Any], route: str) -> tuple[dict[str, int | float], call["status_code"] == 200 and not call.get("error_category"), f"{label}: terminal call failed", ) - require(call["response_status"] == "completed", f"{label}: captured call is incomplete") - require(call.get("model") == route, f"{label}: captured response has the wrong model group") - require(response_id in responses, f"{label}: final answer has no gateway evidence") - event = responses[response_id] - require( - event.get("response_status") == "completed", f"{label}: gateway response is incomplete" - ) require( - number(call["tokens_total"], "terminal-answer tokens") == event["tokens_total"], - f"{label}: gateway/capture token totals differ", + call["response_status"] == response["status"], + f"{label}: response and capture status differ", ) - captured_attempts += len(records) - captured_errors += sum( - record.get("status_code") != 200 or bool(record.get("error_category")) - for record in records - ) - calls.append(call) + 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(event["selected_model"]) - + 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), - "Terminal-answer tokens": sum( - number(call["tokens_total"], "terminal-answer tokens") for call in calls - ), - "Gateway-reported tokens": sum(tokens), - "Gateway requests w/o usage": unknown_usage, + "Captured input tokens": token_total(calls, "tokens_in"), + "Captured output tokens": token_total(calls, "tokens_out"), "Mean rollout latency (ms)": mean(latencies), - "Mean recorded routing (ms)": mean(routing_times), - "Recorded routing decisions": len(routing_times), - "Gateway requests": len(finishes), - "Gateway errors": errors, - "Captured model attempts": captured_attempts, - "Captured failed attempts": captured_errors, + "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(fixed: Path, routed: Path) -> None: - """Report coverage first, then compare like-for-like complete runs.""" - runs = {"fixed": load_run(fixed), "routed": load_run(routed)} +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"]) @@ -247,92 +161,44 @@ def compare(fixed: Path, routed: Path) -> None: runs["fixed"]["inputs"] == runs["routed"]["inputs"], "Task inputs, verifier metadata, or generation settings differ", ) - provenances = [run["provenance"] for run in runs.values()] - for provenance in provenances: - require( - all( - isinstance(provenance.get(key), str) and provenance[key].strip() - for key in ("gym_revision", "switchyard_revision") - ), - "Missing Gym or Switchyard revision", - ) - runtime = provenance["runtime"] - require(runtime["mode"] == "litellm_libsy", "Expected the LiteLLM libsy integration") - require( - runtime["routing_plugin"] == "switchyard_litellm.RandomRoutingPlugin", - "Expected Switchyard Random routing", - ) - require( - all( - isinstance(runtime.get(key), str) and runtime[key].strip() - for key in ("instance_id", "litellm_version", "switchyard_version") - ), - "Missing LiteLLM runtime identity", - ) - require( - all( - re.fullmatch(r"[0-9a-f]{64}", str(runtime.get(key))) is not None - for key in ( - "profile_sha256", - "routing_sha256", - "callback_sha256", - "provider_base_sha256", - ) - ), - "Missing deployment fingerprint", - ) - models = runtime["models"] - require( - all( - isinstance(models.get(route), list) - and all(isinstance(model, str) and model for model in models[route]) - for route in ("fixed", "routed") - ), - "Missing configured models", - ) - require( - len(models["fixed"]) == 1 - and len(set(models["routed"])) == 2 - and models["fixed"][0] in models["routed"], - "Expected a fixed target and a routed pair containing it", - ) - require(provenances[0] == provenances[1], "Gym, Switchyard, or deployment provenance differs") summaries = {name: summarize(run, name) for name, run in runs.items()} - print(f"\nProvenance: {json.dumps(provenances[0], sort_keys=True)}") 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 = [str(value) if type(value) is int else f"{value:.3f}" for value in values] + 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, selected) in summaries.items(): - print(f"\n{name} selected models: {json.dumps(selected, sort_keys=True)}") + for name, (summary, models) in summaries.items(): + print(f"\n{name} serving models: {json.dumps(models, sort_keys=True)}") if ( - summary["Gateway errors"] - or summary["Gateway requests w/o usage"] - or summary["Captured failed attempts"] - or summary["Gateway requests"] != summary["Paired rollouts"] - or summary["Captured model attempts"] != summary["Paired rollouts"] + 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} completed with recovered errors or additional work.") - print("\nClassifier tokens: N/A (Random makes no classifier calls).") + print( + f"WARNING: {name} has incomplete responses, recovered errors, additional calls, or unknown usage." + ) print( - "Gateway totals include all recorded requests, including extra attempts; do not add terminal-answer tokens again." + "\nGym-scored incomplete responses remain in the comparison. " + "Gym-captured tokens include extra attempts; +unknown means some usage was not reported." ) print( - "Unreported failed-request usage is unknown, not free. Gateway counts are not exhaustive provider-attempt or fallback telemetry." + "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." ) - print("Routing time is already included in rollout latency. Tokens are not dollar costs.") - print("A small Random-routing 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("fixed", type=Path, help="Fixed run directory") - parser.add_argument("routed", type=Path, help="Routed run directory") + parser.add_argument( + "results", type=Path, help="Experiment directory containing fixed/ and routed/" + ) args = parser.parse_args(argv) try: - compare(args.fixed, args.routed) + compare(args.results) except (OSError, ValueError, KeyError, TypeError) as error: print(f"Cannot compare: {error}", file=sys.stderr) return 1 diff --git a/benchmark/nemo_gym/gym_routing_plugin.py b/benchmark/nemo_gym/gym_routing_plugin.py index 866f6ded8..333853d88 100644 --- a/benchmark/nemo_gym/gym_routing_plugin.py +++ b/benchmark/nemo_gym/gym_routing_plugin.py @@ -1,29 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Record LiteLLM request evidence and preserve Gym-compatible usage detail shapes.""" +"""Normalize Gym responses and record the serving model before LiteLLM applies its alias.""" from __future__ import annotations -import hashlib import json import os -from copy import deepcopy -from importlib.metadata import version from pathlib import Path -from time import perf_counter from typing import Any -import yaml from litellm.integrations.custom_logger import CustomLogger -from litellm.types.router import RoutingContext -from switchyard_litellm import RandomRoutingPlugin -from switchyard_litellm.configuration import load_routing_plugin def response_payload(response: Any) -> dict[str, Any]: """Keep missing detail counts unknown while satisfying Gym's object-shaped fields.""" - payload = response.model_dump() if hasattr(response, "model_dump") else deepcopy(dict(response)) + payload = response.model_dump() for item in payload.get("output") or []: if not isinstance(item, dict) or item.get("type") != "reasoning": continue @@ -49,148 +41,35 @@ def response_payload(response: Any) -> dict[str, Any]: class GymRoutingPlugin(CustomLogger): - """Join routing and callback instances using one runner-supplied gateway identity.""" + """Provide response compatibility and final-model attribution, not routing or accounting.""" - def __init__( - self, routing_path: Path, profile_path: Path, results: Path, instance_id: str - ) -> None: + def __init__(self, results: Path) -> None: super().__init__() - if not instance_id.strip(): - raise ValueError("The runner must supply a gateway instance ID") - self.routing = load_routing_plugin(routing_path) - if not isinstance(self.routing, RandomRoutingPlugin): - raise ValueError("This example requires Switchyard Random routing") self.results = results - profile = yaml.safe_load(profile_path.read_text(encoding="utf-8")) - models = { - route: [ - entry["litellm_params"]["model"] - for entry in profile["model_list"] - if entry["model_name"] == route - ] - for route in ("fixed", "routed") - } - if len(models["fixed"]) != 1 or len(set(models["routed"])) != 2: - raise ValueError("Define one fixed target and two distinct routed targets") - if models["fixed"][0] not in models["routed"]: - raise ValueError("The fixed target must be one of the routed targets") - self.runtime = { - "mode": "litellm_libsy", - "instance_id": instance_id, - "litellm_version": version("litellm"), - "switchyard_version": version("nemo-switchyard"), - "fastapi_version": version("fastapi"), - "starlette_version": version("starlette"), - "routing_plugin": "switchyard_litellm.RandomRoutingPlugin", - "models": models, - "profile_sha256": hashlib.sha256(profile_path.read_bytes()).hexdigest(), - "routing_sha256": hashlib.sha256(routing_path.read_bytes()).hexdigest(), - "callback_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), - "provider_base_sha256": hashlib.sha256( - os.environ.get("NVIDIA_BASE_URL", "").encode() - ).hexdigest(), - } - results.mkdir(parents=True, exist_ok=True) - (results / "litellm-runtime.json").write_text( - json.dumps(self.runtime, indent=2) + "\n", encoding="utf-8" - ) - - async def run(self, context: RoutingContext) -> RoutingContext: - """Measure the actual libsy decision without including provider inference.""" - started = perf_counter() - result = await self.routing.run(context) - result.signals["switchyard"]["routing_ms"] = (perf_counter() - started) * 1000 - return result - - async def async_pre_call_deployment_hook( - self, kwargs: dict[str, Any], call_type: Any - ) -> dict[str, Any] | None: - """Keep the configured plugin's deployment callback behavior intact.""" - return await self.routing.async_pre_call_deployment_hook(kwargs, call_type) - - def _record(self, data: dict[str, Any], event: dict[str, Any]) -> None: - """Write only allowlisted evidence, never request content or credentials.""" - route, request_id = data.get("model"), data.get("litellm_call_id") - if route not in ("fixed", "routed"): - raise ValueError("This example accepts only the fixed and routed model groups") - if not isinstance(request_id, str) or not request_id: - raise ValueError("LiteLLM request ID is missing") - directory = self.results / route - directory.mkdir(parents=True, exist_ok=True) - record = { - "instance_id": self.runtime["instance_id"], - "route": route, - "request_id": request_id, - **event, - } - with (directory / "litellm-calls.jsonl").open("a", encoding="utf-8") as stream: - stream.write(json.dumps(record, allow_nan=False) + "\n") - - async def async_pre_call_hook( - self, user_api_key_dict: Any, cache: Any, data: dict[str, Any], call_type: Any - ) -> dict[str, Any]: - """Record a request before routing so interrupted evidence can be detected.""" - self._record(data, {"event": "start"}) - if data.get("stream"): - raise ValueError("This example requires non-streaming requests") - return data async def async_post_call_success_hook( self, data: dict[str, Any], user_api_key_dict: Any, response: Any ) -> dict[str, Any]: - """Record the final response before LiteLLM replaces its model with the public alias.""" + """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 {} - signal = (metadata.get("routing_plugin_signals") or {}).get("switchyard") or {} - self._record( - data, - { - "event": "finish", - "status_code": 200, - "error_type": None, - "response_id": payload.get("id"), - "response_status": payload.get("status"), - "selected_model": signal.get("selected_model_id"), - "deployment_model": metadata.get("deployment"), - "tokens_total": (payload.get("usage") or {}).get("total_tokens"), - "routing_ms": signal.get("routing_ms"), - }, - ) + 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 - async def async_post_call_failure_hook( - self, - request_data: dict[str, Any], - original_exception: Exception, - user_api_key_dict: Any, - traceback_str: str | None = None, - ) -> None: - """Record failed gateway requests without inventing unreported token usage.""" - metadata = request_data.get("litellm_metadata") or {} - signal = (metadata.get("routing_plugin_signals") or {}).get("switchyard") or {} - self._record( - request_data, - { - "event": "finish", - "status_code": getattr(original_exception, "status_code", None), - "error_type": type(original_exception).__name__, - "response_id": None, - "response_status": None, - "selected_model": signal.get("selected_model_id"), - "deployment_model": metadata.get("deployment"), - "tokens_total": None, - "routing_ms": signal.get("routing_ms"), - }, - ) - PLUGIN = ( - GymRoutingPlugin( - Path(os.environ["SWITCHYARD_LITELLM_CONFIG"]), - Path(os.environ["NEMO_GYM_LITELLM_PROFILE"]), - Path(os.environ["NEMO_GYM_LITELLM_RESULTS"]), - os.environ["NEMO_GYM_LITELLM_INSTANCE_ID"], - ) + 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 index 91078488a..6ff745c7d 100644 --- a/benchmark/nemo_gym/litellm.yaml +++ b/benchmark/nemo_gym/litellm.yaml @@ -1,7 +1,7 @@ model_list: - model_name: fixed - litellm_params: &strong - model: nvidia_nim/nvidia/nemotron-3-super-120b-a12b + 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 @@ -9,20 +9,22 @@ model_list: chat_template_kwargs: enable_thinking: false - model_name: routed - litellm_params: *strong + litellm_params: *ultra - model_name: routed litellm_params: - model: nvidia_nim/openai/gpt-oss-20b + 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: - reasoning_effort: low + chat_template_kwargs: + enable_thinking: false router_settings: num_retries: 0 plugins: - - gym_routing_plugin.PLUGIN + - 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/run.sh b/benchmark/nemo_gym/run.sh index a463c1fac..1beff744b 100644 --- a/benchmark/nemo_gym/run.sh +++ b/benchmark/nemo_gym/run.sh @@ -11,8 +11,8 @@ set -euo pipefail # 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="${LITELLM_CONFIG:-$SCRIPT_DIR/litellm.yaml}" -ROUTING="${SWITCHYARD_CONFIG:-$SCRIPT_DIR/routes.toml}" +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}" @@ -27,10 +27,8 @@ 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/) - LITELLM_CONFIG Model inventory (litellm.yaml beside this script) - SWITCHYARD_CONFIG Routing policy (routes.toml beside this script) NVIDIA_BASE_URL Provider endpoint (NVIDIA API Catalog /v1) -Provider credentials come from the environment variables named in the profile. +Provider credentials come from the environment variables named in litellm.yaml. EOF [[ $# -eq 1 && ( "$1" == "-h" || "$1" == "--help" ) ]] && exit 0 exit 2 @@ -39,8 +37,6 @@ fi die() { echo "error: $*" >&2; exit 1; } [[ -n "${GYM_DIR:-}" ]] || die "set GYM_DIR; see the README setup" GYM_DIR="$(cd "$GYM_DIR" && pwd)" -[[ "$PROFILE" = /* ]] || PROFILE="$PWD/$PROFILE" -[[ "$ROUTING" = /* ]] || ROUTING="$PWD/$ROUTING" [[ "$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" @@ -53,8 +49,7 @@ 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" rev-parse HEAD)" -[[ -z "$(git -C "$GYM_DIR" status --porcelain --untracked-files=no)" ]] || die "tracked Gym source must be clean" +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 @@ -129,10 +124,7 @@ fi # 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_PROFILE="$PROFILE" export NEMO_GYM_LITELLM_RESULTS="$RESULTS_DIR" -NEMO_GYM_LITELLM_INSTANCE_ID="$("$PYTHON" -c 'from uuid import uuid4; print(uuid4().hex)')" -export NEMO_GYM_LITELLM_INSTANCE_ID 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 \ @@ -143,7 +135,7 @@ 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" >"$RESULTS_DIR/litellm-health.json" 2>/dev/null && [[ -s "$RESULTS_DIR/litellm-runtime.json" ]]; then + if curl -fsS --max-time 2 "$ROOT_URL/health/readiness/details" >/dev/null 2>&1; then ready=true break fi @@ -151,26 +143,15 @@ for _ in {1..240}; do done [[ "$ready" == true ]] || die "LiteLLM readiness timed out; see $RESULTS_DIR/litellm.log" -# Step 4: Save the shared runtime details alongside each condition's results. -# The comparison uses these revisions and configuration fingerprints to check compatibility. -"$PYTHON" - "$RESULTS_DIR" "$GYM_REVISION" "$SWITCHYARD_REVISION" <<'PY' -import json -import pathlib -import sys - -results = pathlib.Path(sys.argv[1]) -runtime = json.loads((results / 'litellm-runtime.json').read_text()) -provenance = {'gym_revision': sys.argv[2], 'switchyard_revision': sys.argv[3], 'runtime': runtime} -for route in ('fixed', 'routed'): - run_dir = results / route - (run_dir / 'model-calls').mkdir(parents=True, exist_ok=True) - (run_dir / 'run-provenance.json').write_text(json.dumps(provenance, indent=2) + '\n', encoding='utf-8') -PY +# 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 request ledgers. +# 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" @@ -179,15 +160,14 @@ for route in fixed routed; do --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 4096 \ + --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" - # Keep request evidence even when Gym collection fails. - [[ -s "$run_dir/litellm-calls.jsonl" ]] || die "missing LiteLLM request evidence for $route" + # 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 @@ -196,5 +176,5 @@ done stop_process "$PROXY_PID" PROXY_PID="" echo "Comparing fixed and routed results" -"$PYTHON" "$SCRIPT_DIR/compare.py" "$RESULTS_DIR/fixed" "$RESULTS_DIR/routed" 2>&1 | tee "$RESULTS_DIR/comparison.txt" +"$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 index efa0a1bb4..26b770c6a 100644 --- a/tests/test_nemo_gym_compare.py +++ b/tests/test_nemo_gym_compare.py @@ -12,19 +12,13 @@ import pytest -BIG = "nvidia_nim/nvidia/nemotron-3-super-120b-a12b" -SMALL = "nvidia_nim/openai/gpt-oss-20b" -FILES = { - "inputs": "rollouts_materialized_inputs.jsonl", - "rows": "rollouts.jsonl", - "failures": "rollouts_failures.jsonl", - "events": "litellm-calls.jsonl", - "provenance": "run-provenance.json", -} +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 @@ -33,60 +27,36 @@ def comparator() -> ModuleType: return module -def _attempt(route: str, request_id: str, **fields: Any) -> list[dict[str, Any]]: - common = {"route": route, "request_id": request_id, "instance_id": "fixture-proxy"} - return [ - {**common, "event": "start"}, - { - **common, - "event": "finish", - "status_code": 200, - "error_type": None, - "response_id": request_id, - "response_status": "completed", - "selected_model": BIG, - "deployment_model": BIG, - "tokens_total": 30, - "routing_ms": 5, - **fields, - }, - ] - - @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, events = [], [], [] + inputs, rows, models = [], [], [] for index in range(2): task = { "_ng_task_index": index, "_ng_rollout_index": 0, "expected_answer": "B", - "agent_ref": {"name": "mmlu-redux_mcqa_simple_agent"}, "responses_create_params": { "input": [{"role": "user", "content": f"Question {index}"}], "temperature": 0, - "max_output_tokens": 4096, + "max_output_tokens": 8192, }, } - inputs.append(task) response_id = f"{route}-{index}" - tokens = 30 + 10 * index - model = SMALL if route == "routed" and index == 0 else BIG + inputs.append(task) rows.append( { **deepcopy(task), + "reward": 1 - index, "response": { "id": response_id, "status": "completed", - "model": route, "output": [ { "type": "message", "role": "assistant", - "status": "completed", "content": [ { "type": "output_text", @@ -96,70 +66,52 @@ def artifacts() -> dict[str, dict[str, Any]]: } ], }, - "reward": 1 - index, "ng_perf": {"total_latency_ms": 100 + 200 * index}, "ng_model_call_capture": { - "gaps": [], "calls": [ { + "model_call_id": response_id, "response_id": response_id, + "model": route, "status_code": 200, "error_category": None, "response_status": "completed", - "model": route, - "tokens_total": tokens, + "tokens_in": 10 + 10 * index, + "tokens_out": 20, } ], }, } ) - events.extend( - _attempt( - route, - response_id, - tokens_total=tokens, - selected_model=model, - deployment_model=model, - ) + models.append( + { + "response_id": response_id, + "model": SMALL if route == "routed" and index == 0 else BIG, + } ) runs[route] = { - "inputs": inputs, - "rows": rows, - "failures": [], - "events": events, - "provenance": { - "gym_revision": "b" * 40, - "switchyard_revision": "c" * 40, - "runtime": { - "mode": "litellm_libsy", - "instance_id": "fixture-proxy", - "litellm_version": "1.97.0", - "switchyard_version": "0.2.0", - "routing_plugin": "switchyard_litellm.RandomRoutingPlugin", - "models": {"fixed": [BIG], "routed": [BIG, SMALL]}, - "profile_sha256": "a" * 64, - "routing_sha256": "a" * 64, - "callback_sha256": "a" * 64, - "provider_base_sha256": "a" * 64, - }, - }, + "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, run in artifacts.items(): + for route, files in artifacts.items(): directory = tmp_path / route directory.mkdir() - for name, filename in FILES.items(): - value = run[name] - text = ( - "".join(json.dumps(row) + "\n" for row in value) - if filename.endswith(".jsonl") - else json.dumps(value) + for name, rows in files.items(): + (directory / name).write_text( + "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" ) - (directory / filename).write_text(text, encoding="utf-8") - return [str(tmp_path / route) for route in ("fixed", "routed")] + 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( @@ -168,27 +120,30 @@ def test_complete_reordered_pair( artifacts: dict[str, dict[str, Any]], capsys: pytest.CaptureFixture[str], ) -> None: - artifacts["routed"]["rows"].reverse() - artifacts["routed"]["inputs"].reverse() - artifacts["routed"]["events"].reverse() + """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 selected models: {{"{BIG}": 2}}' in output.out - assert f"routed selected models: {json.dumps({SMALL: 1, BIG: 1}, sort_keys=True)}" 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"], - "Terminal-answer tokens": ["70", "70"], - "Gateway-reported tokens": ["70", "70"], + "Captured input tokens": ["30", "30"], + "Captured output tokens": ["40", "40"], "Mean rollout latency (ms)": ["200", "200"], - "Gateway requests": ["2", "2"], - "Gateway errors": ["0", "0"], + "Captured calls": ["2", "2"], + "Captured failed calls": ["0", "0"], + "Calls with unknown usage": ["0", "0"], }.items(): - line = next(line for line in output.out.splitlines() if line[:29].rstrip() == metric) - assert line[29:].split() == values - assert "Classifier tokens: N/A" in output.out - assert "not exhaustive provider-attempt" in output.out + _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( @@ -197,7 +152,6 @@ def test_complete_reordered_pair( ("incomplete", "Incomplete runs"), ("mismatched", "Task inputs, verifier metadata, or generation settings differ"), ("capture", "incomplete model-call capture"), - ("truncated", "missing or incomplete final answer"), ], ) def test_invalid_evidence_never_prints_averages( @@ -208,16 +162,15 @@ def test_invalid_evidence_never_prints_averages( 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["rows"].pop() + run["rollouts.jsonl"].pop() elif problem == "mismatched": - routed["inputs"][0]["expected_answer"] = "A" - elif problem == "capture": - routed["rows"][0]["ng_model_call_capture"]["gaps"] = ["missing exchange"] + routed["rollouts_materialized_inputs.jsonl"][0]["expected_answer"] = "A" else: - routed["rows"][0]["response"]["status"] = "incomplete" + 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 @@ -225,39 +178,63 @@ def test_invalid_evidence_never_prints_averages( 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: - run = artifacts["routed"] - calls = run["rows"][0]["ng_model_call_capture"]["calls"] - calls.insert(0, {"status_code": 503, "error_category": "upstream", "tokens_total": None}) - calls.append({**calls[1], "response_id": "superseded", "tokens_total": 20}) - run["events"].extend( - _attempt( - "routed", - "failed", - status_code=503, - error_type="ServiceUnavailable", - response_id=None, - response_status=None, - tokens_total=None, - routing_ms=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, + } ) - run["events"].extend(_attempt("routed", "superseded", tokens_total=20)) + 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 { - "Terminal-answer tokens": ["70", "70"], - "Gateway-reported tokens": ["70", "90"], - "Gateway requests": ["2", "4"], - "Gateway errors": ["0", "1"], - "Gateway requests w/o usage": ["0", "1"], - "Captured failed attempts": ["0", "1"], + "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(): - line = next(line for line in output.splitlines() if line[:29].rstrip() == metric) - assert line[29:].split() == values + _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 index bd4f1062f..ffd718fb8 100644 --- a/tests/test_nemo_gym_litellm.py +++ b/tests/test_nemo_gym_litellm.py @@ -3,12 +3,8 @@ from __future__ import annotations -import hashlib -import importlib.metadata import importlib.util import json -import math -from copy import deepcopy from pathlib import Path from types import ModuleType from typing import Any @@ -16,16 +12,13 @@ import pytest litellm = pytest.importorskip("litellm") -pytest.importorskip("switchyard_litellm") from litellm import ResponsesAPIResponse # noqa: E402 -from litellm.types.router import RoutingContext # noqa: E402 -BIG = "nvidia_nim/nvidia/nemotron-3-super-120b-a12b" -SMALL = "nvidia_nim/openai/gpt-oss-20b" +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" -PROFILE = ROOT / "benchmark/nemo_gym/litellm.yaml" @pytest.fixture @@ -58,22 +51,6 @@ def _response( return response -def _plugin(callback: ModuleType, tmp_path: Path) -> Any: - routing = tmp_path / "routes.toml" - routing.write_text('algorithm = "random"\nseed = 6\n', encoding="utf-8") - return callback.GymRoutingPlugin(routing, PROFILE, tmp_path / "results", "fixture-proxy") - - -def _context(candidates: list[str], text: str) -> RoutingContext: - messages = [{"role": "user", "content": text}] - return RoutingContext( - raw_messages=messages, - structured_messages=messages, - candidate_models=candidates, - metadata={"model_group": "routed"}, - ) - - def test_response_payload_preserves_wire_identity_and_unknown_usage_details( callback: ModuleType, ) -> None: @@ -115,9 +92,8 @@ def test_response_payload_preserves_wire_identity_and_unknown_usage_details( @pytest.mark.parametrize("finish_reason", ["stop", "length"]) -@pytest.mark.parametrize("as_dict", [False, True]) def test_response_payload_normalizes_litellm_reasoning( - callback: ModuleType, finish_reason: str, as_dict: bool + callback: ModuleType, finish_reason: str ) -> None: from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, @@ -146,9 +122,8 @@ def test_response_payload_normalizes_litellm_reasoning( "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, }, ) - before = deepcopy(response.model_dump()) - source = deepcopy(before) if as_dict else response - payload = callback.response_payload(source) + before = response.model_dump() + payload = callback.response_payload(response) reasoning, answer = payload["output"] validated = ResponseReasoningItem.model_validate(reasoning) @@ -166,191 +141,38 @@ def test_response_payload_normalizes_litellm_reasoning( 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 callback.response_payload(payload) == payload - assert (source if as_dict else source.model_dump()) == before - - -@pytest.mark.parametrize("summary", [None, [{"type": "summary_text", "text": "Existing summary"}]]) -def test_response_payload_preserves_native_reasoning_and_other_items( - callback: ModuleType, summary: Any -) -> None: - source = { - "id": "response-native", - "status": "incomplete", - "usage": None, - "output": [ - { - "type": "reasoning", - "id": "reasoning-native", - "summary": summary, - "encrypted_content": "opaque-test-content", - "content": [{"type": "reasoning_text", "text": "Existing reasoning"}], - }, - {"type": "message", "content": [{"type": "output_text", "text": "Answer"}]}, - {"type": "function_call", "name": "test_tool", "arguments": "{}"}, - ], - } - before = deepcopy(source) - expected = deepcopy(source) - if summary is None: - expected["output"][0]["summary"] = [] - assert callback.response_payload(source) == expected - assert source == before - - -async def test_real_libsy_random_routing_reaches_both_candidates( - callback: ModuleType, tmp_path: Path -) -> None: - plugin = _plugin(callback, tmp_path) - - first = await plugin.run(_context([BIG, SMALL], "first")) - second = await plugin.run(_context([BIG, SMALL], "second")) - fixed = await plugin.run(_context([BIG], "fixed")) - - selections = { - first.signals["switchyard"]["selected_model_id"], - second.signals["switchyard"]["selected_model_id"], - } - assert selections == {BIG, SMALL} - assert fixed.signals["switchyard"]["selected_model_id"] == BIG - for result in (first, second, fixed): - routing_ms = result.signals["switchyard"]["routing_ms"] - assert isinstance(routing_ms, float) and math.isfinite(routing_ms) and routing_ms >= 0 - - -def test_litellm_local_file_loader_shares_runner_instance_id( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - from litellm.proxy.types_utils.utils import get_instance_fn - - routing = tmp_path / "routes.toml" - routing.write_text('algorithm = "random"\nseed = 6\n', encoding="utf-8") - results = tmp_path / "results" - monkeypatch.setenv("SWITCHYARD_LITELLM_CONFIG", str(routing)) - monkeypatch.setenv("NEMO_GYM_LITELLM_PROFILE", str(PROFILE)) - monkeypatch.setenv("NEMO_GYM_LITELLM_RESULTS", str(results)) - monkeypatch.setenv("NEMO_GYM_LITELLM_INSTANCE_ID", "fixture-shared-run") - - first = get_instance_fn(value="gym_routing_plugin.PLUGIN", config_file_path=str(PROFILE)) - second = get_instance_fn(value="gym_routing_plugin.PLUGIN", config_file_path=str(PROFILE)) - - assert first is not second - assert first.runtime["instance_id"] == "fixture-shared-run" - assert second.runtime["instance_id"] == "fixture-shared-run" - runtime = json.loads((results / "litellm-runtime.json").read_text()) - assert runtime["instance_id"] == "fixture-shared-run" - data = {"model": "fixed", "litellm_call_id": "fixture-request"} - first._record(data, {"event": "start"}) - second._record(data, {"event": "finish"}) - ledger = [ - json.loads(line) - for line in (results / "fixed/litellm-calls.jsonl").read_text().splitlines() - ] - assert [record["instance_id"] for record in ledger] == [ - "fixture-shared-run", - "fixture-shared-run", - ] + assert response.model_dump() == before -async def test_callbacks_record_allowlisted_success_and_runtime( +async def test_callback_records_serving_model_without_request_content( callback: ModuleType, tmp_path: Path ) -> None: - plugin = _plugin(callback, tmp_path) - selected = SMALL + plugin = callback.GymRoutingPlugin(tmp_path) + response = _response( + status="completed", usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15} + ) data = { "model": "routed", - "litellm_call_id": "request-1", "messages": [{"role": "user", "content": "secret-prompt-marker"}], "headers": {"Authorization": "secret-header-marker"}, "litellm_metadata": { - "deployment": selected, - "routing_plugin_signals": { - "switchyard": {"selected_model_id": selected, "routing_ms": 1.5} - }, + "deployment": SMALL, + "routing_plugin_signals": {"switchyard": {"selected_model_id": BIG}}, }, } - response = _response( - status="completed", - usage={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, - ) - - await plugin.async_pre_call_hook(None, None, data, "aresponses") payload = await plugin.async_post_call_success_hook(data, None, response) - - records_path = tmp_path / "results/routed/litellm-calls.jsonl" - records = [json.loads(line) for line in records_path.read_text().splitlines()] - assert records == [ - { - "instance_id": plugin.runtime["instance_id"], - "route": "routed", - "request_id": "request-1", - "event": "start", - }, - { - "instance_id": plugin.runtime["instance_id"], - "route": "routed", - "request_id": "request-1", - "event": "finish", - "status_code": 200, - "error_type": None, - "response_id": "response-1", - "response_status": "completed", - "selected_model": selected, - "deployment_model": selected, - "tokens_total": 15, - "routing_ms": 1.5, - }, + records = [ + json.loads(line) for line in (tmp_path / "routed/models.jsonl").read_text().splitlines() ] - assert payload["id"] == "response-1" - assert payload["status"] == "completed" - assert payload["usage"]["total_tokens"] == 15 - assert payload["_hidden_params"] == response._hidden_params - serialized = records_path.read_text() - assert "secret-prompt-marker" not in serialized - assert "secret-header-marker" not in serialized - - runtime = json.loads((tmp_path / "results/litellm-runtime.json").read_text()) - assert runtime["profile_sha256"] == hashlib.sha256(PROFILE.read_bytes()).hexdigest() - assert runtime["callback_sha256"] == hashlib.sha256(CALLBACK.read_bytes()).hexdigest() - assert runtime["litellm_version"] == importlib.metadata.version("litellm") - assert runtime["switchyard_version"] == importlib.metadata.version("nemo-switchyard") - assert runtime["routing_plugin"] == "switchyard_litellm.RandomRoutingPlugin" - - -async def test_failure_hook_records_unknown_usage_without_secrets( - callback: ModuleType, tmp_path: Path -) -> None: - plugin = _plugin(callback, tmp_path) - selected = BIG - data = { - "model": "fixed", - "litellm_call_id": "request-failed", - "messages": [{"role": "user", "content": "secret-failure-message"}], - "headers": {"Authorization": "secret-failure-header"}, - "litellm_metadata": { - "deployment": selected, - "routing_plugin_signals": { - "switchyard": {"selected_model_id": selected, "routing_ms": 0.5} - }, - }, - } - - class ServiceUnavailable(Exception): - status_code = 503 - - await plugin.async_pre_call_hook(None, None, data, "aresponses") - await plugin.async_post_call_failure_hook(data, ServiceUnavailable("secret-error"), None) - - records_path = tmp_path / "results/fixed/litellm-calls.jsonl" - records = [json.loads(line) for line in records_path.read_text().splitlines()] - assert len(records) == 2 - finish = records[1] - assert finish["event"] == "finish" - assert finish["status_code"] == 503 - assert finish["error_type"] == "ServiceUnavailable" - assert finish["tokens_total"] is None - assert finish["response_id"] is None - serialized = records_path.read_text() - assert "secret-failure-message" not in serialized - assert "secret-failure-header" not in serialized - assert "secret-error" not in serialized + 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 index 9fe0c641e..49c240887 100644 --- a/tests/test_nemo_gym_run.py +++ b/tests/test_nemo_gym_run.py @@ -3,7 +3,6 @@ from __future__ import annotations -import json import os import signal import subprocess @@ -87,9 +86,6 @@ def stop(number, _frame): run_dir="$(dirname "$output")" mkdir -p "$run_dir" printf '{"route": "%s"}\\n' "$route" > "$output" -if [[ "${FAKE_GYM_BEHAVIOR:-}" != "missing_ledger" ]]; then - printf '{"event":"start"}\\n{"event":"finish"}\\n' > "$run_dir/litellm-calls.jsonl" -fi 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 @@ -100,46 +96,10 @@ def stop(number, _frame): _write_executable( bin_dir / "uv", """#!/bin/bash -set -u -printf 'uv:%s\\n' "$*" >> "$EVENT_LOG" -"$REAL_PYTHON" - "$NEMO_GYM_LITELLM_RESULTS" "$NEMO_GYM_LITELLM_PROFILE" \ - "$SWITCHYARD_LITELLM_CONFIG" "$FAKE_CALLBACK_PATH" "$NVIDIA_BASE_URL" <<'PY' -import hashlib -import json -import os -import pathlib -import sys - -results, profile, routing, callback = map(pathlib.Path, sys.argv[1:5]) -runtime = { - "mode": "litellm_libsy", - "instance_id": os.environ["NEMO_GYM_LITELLM_INSTANCE_ID"], - "litellm_version": "1.97.0", - "switchyard_version": "0.2.0", - "fastapi_version": "0.136.3", - "starlette_version": "1.3.1", - "routing_plugin": "switchyard_litellm.RandomRoutingPlugin", - "models": { - "fixed": ["nvidia_nim/nvidia/nemotron-3-super-120b-a12b"], - "routed": [ - "nvidia_nim/nvidia/nemotron-3-super-120b-a12b", - "nvidia_nim/openai/gpt-oss-20b", - ], - }, - "profile_sha256": hashlib.sha256(profile.read_bytes()).hexdigest(), - "routing_sha256": hashlib.sha256(routing.read_bytes()).hexdigest(), - "callback_sha256": hashlib.sha256(callback.read_bytes()).hexdigest(), - "provider_base_sha256": hashlib.sha256(sys.argv[5].encode()).hexdigest(), -} -results.mkdir(parents=True, exist_ok=True) -(results / "litellm-runtime.json").write_text(json.dumps(runtime), encoding="utf-8") -PY exec "$REAL_PYTHON" -c ' import os import signal -import sys - -arguments = " ".join(sys.argv[1:]) +from pathlib import Path def stop(_number, _frame): with open(os.environ["EVENT_LOG"], "a", encoding="utf-8") as stream: @@ -149,38 +109,24 @@ def stop(_number, _frame): 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()}:{arguments}\\n") + 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 -url="" -for argument in "$@"; do case "$argument" in http://*) url="$argument" ;; esac; done -printf 'curl:%s\\n' "$url" >> "$EVENT_LOG" -[[ -s "$NEMO_GYM_LITELLM_RESULTS/litellm-runtime.json" ]] || exit 22 -printf '{}\\n' +[[ -f "$NEMO_GYM_LITELLM_RESULTS/proxy-ready" ]] """, ) _write_executable( bin_dir / "git", """#!/bin/bash -printf 'git:%s\\n' "$*" >> "$EVENT_LOG" -case "$*" in - *" status "*) exit 0 ;; - *" rev-parse "*) printf '%039d1\\n' 0 ;; - *" describe "*) printf '%040d-dirty\\n' 2 ;; -esac +printf '%040d-dirty\\n' 2 """, ) - for tool in ("cargo", "switchyard-server"): - _write_executable( - bin_dir / tool, - f'#!/bin/bash\nprintf \'forbidden:{tool}:%s\\n\' "$*" >> "$EVENT_LOG"\nexit 99\n', - ) - env = os.environ.copy() for key in ( "NVIDIA_API_KEY", @@ -188,8 +134,6 @@ def stop(_number, _frame): "ANTHROPIC_API_KEY", "OPENROUTER_API_KEY", "NVIDIA_BASE_URL", - "LITELLM_CONFIG", - "SWITCHYARD_CONFIG", "LITELLM_PORT", ): env.pop(key, None) @@ -203,7 +147,6 @@ def stop(_number, _frame): gym_dir / "benchmarks/mmlu-redux/data/mmlu-redux_benchmark.jsonl" ), "FAKE_COMPARE_PATH": str(ROOT / "benchmark/nemo_gym/compare.py"), - "FAKE_CALLBACK_PATH": str(ROOT / "benchmark/nemo_gym/gym_routing_plugin.py"), "REAL_PYTHON": sys.executable, } ) @@ -248,8 +191,6 @@ def test_help_and_shell_syntax_need_no_setup(tmp_path: Path) -> None: assert result.stderr == "" for name in ( "GYM_DIR", - "LITELLM_CONFIG", - "SWITCHYARD_CONFIG", "NVIDIA_BASE_URL", "LITELLM_PORT", "LIMIT", @@ -261,105 +202,49 @@ def test_help_and_shell_syntax_need_no_setup(tmp_path: Path) -> None: assert subprocess.run(["/bin/bash", "-n", str(RUNNER)], check=False).returncode == 0 -@pytest.mark.parametrize("invalid", ["existing-results", "zero-limit", "bad-port"]) -def test_preflight_rejects_before_operations( - fake_runner_env: tuple[dict[str, str], Path, Path], invalid: str +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 - if invalid == "existing-results": - results_dir.mkdir() - sentinel = results_dir / "sentinel" - sentinel.write_text("keep", encoding="utf-8") - else: - sentinel = None - env["LIMIT" if invalid == "zero-limit" else "LITELLM_PORT"] = "0" + 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) == [] - if sentinel is not None: - assert sentinel.read_text(encoding="utf-8") == "keep" + 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", - "NEMO_GYM_LITELLM_INSTANCE_ID": "must-not-be-reused", - } - ) + env.update({"LIMIT": "2", "REPEATS": "3", "CONCURRENCY": "1"}) result = _run(env) assert result.returncode == 0, result.stderr events = _events(event_log) - - preparations = [line for line in events if line.startswith("prepare:")] - assert len(preparations) == 1 - assert "eval prepare --benchmark mmlu-redux" in preparations[0] + 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]) - for pin in ( - "litellm[proxy]==1.97.0", - "fastapi==0.136.3", - "starlette==1.3.1", - "--num_workers 1", - ): - assert pin in starts[0] 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 route, invocation in zip(("fixed", "routed"), gym_runs, strict=True): - run_dir = results_dir / route - for expected in ( + for invocation in gym_runs: + for flag in ( "--benchmark mmlu-redux", "--model-type litellm_model", - f"--model {route}", - "--split benchmark", + "--max-output-tokens 8192", "--limit 2", "--num-repeats 3", "--concurrency 1", - "--temperature 0", - "--max-output-tokens 4096", - "++policy_base_url=http://127.0.0.1:4000/v1", - "++policy_api_key=unused", - f"++model_call_capture_dir={run_dir}/model-calls", - f"++nemo_gym_log_dir={run_dir}/server-logs", - "++mcqa_simple_agent.responses_api_agents.simple_agent.max_steps=1", ): - assert expected in invocation - assert "switchyard_model" not in invocation - assert ".deployment=" not in invocation - assert "condition_dir" not in invocation - assert ( - "++mmlu-redux_mcqa_simple_agent.responses_api_agents.simple_agent.max_steps=" - not in invocation - ) - assert not any(line.startswith("forbidden:") for line in events) - - provenances = [ - json.loads((results_dir / route / "run-provenance.json").read_text()) - for route in ("fixed", "routed") - ] - assert provenances[0] == provenances[1] - runtime = provenances[0]["runtime"] - assert len(runtime["instance_id"]) == 32 - assert all(character in "0123456789abcdef" for character in runtime["instance_id"]) - assert runtime["instance_id"] != "must-not-be-reused" - assert runtime["models"] == { - "fixed": ["nvidia_nim/nvidia/nemotron-3-super-120b-a12b"], - "routed": [ - "nvidia_nim/nvidia/nemotron-3-super-120b-a12b", - "nvidia_nim/openai/gpt-oss-20b", - ], - } + 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" @@ -382,20 +267,6 @@ def test_fixed_failure_stops_before_routed_and_keeps_logs( assert (results_dir / "fixed/rollouts_failures.jsonl").stat().st_size > 0 -def test_missing_ledger_stops_before_comparison( - fake_runner_env: tuple[dict[str, str], Path, Path], -) -> None: - env, _, event_log = fake_runner_env - env["FAKE_GYM_BEHAVIOR"] = "missing_ledger" - result = _run(env) - events = _events(event_log) - assert result.returncode != 0 - assert "missing LiteLLM request evidence for fixed" in result.stderr - 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 - - @pytest.mark.parametrize( ("signal_number", "expected_status"), [(signal.SIGINT, 130), (signal.SIGTERM, 143)], From d3e63ff944ec337f925cb7dd34d04465e6042b0b Mon Sep 17 00:00:00 2001 From: Shashank Verma Date: Tue, 15 Sep 2026 08:14:31 -0700 Subject: [PATCH 6/6] docs(tutorial): Generalize live run wording Signed-off-by: Shashank Verma --- benchmark/nemo_gym/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmark/nemo_gym/README.md b/benchmark/nemo_gym/README.md index dce33571f..9e00ee48f 100644 --- a/benchmark/nemo_gym/README.md +++ b/benchmark/nemo_gym/README.md @@ -79,7 +79,7 @@ Gym components can outlive the command briefly; let them finish shutting down be ## 3. Understand the result -Start with task coverage and serving models, then compare rewards, tokens, and latency. This is the output from a live NVIDIA run of the default five tasks per condition: +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 @@ -158,4 +158,4 @@ Keep the saved inputs because the source dataset can change. Request logs contai -**Validation:** tested with Gym `v0.6.0`, and a fresh live NVIDIA run of five tasks per condition. Ultra and Lightning were both exercised, and all ten captured calls succeeded. +**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.