From 39d424033245a613dca2b402d14770865299752b Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Tue, 1 Sep 2026 17:43:05 -0700 Subject: [PATCH] feat(benchmark): add NeMo Gym routing comparison Signed-off-by: Alex Fournier --- benchmark/README.md | 7 +- benchmark/nemo_gym/README.md | 121 ++++++++++++++++++ benchmark/nemo_gym/compare.py | 225 +++++++++++++++++++++++++++++++++ benchmark/nemo_gym/routes.toml | 40 ++++++ benchmark/nemo_gym/run.sh | 168 ++++++++++++++++++++++++ tests/test_nemo_gym_compare.py | 157 +++++++++++++++++++++++ tests/test_nemo_gym_run.py | 42 ++++++ 7 files changed, 759 insertions(+), 1 deletion(-) create mode 100644 benchmark/nemo_gym/README.md create mode 100644 benchmark/nemo_gym/compare.py create mode 100644 benchmark/nemo_gym/routes.toml create mode 100755 benchmark/nemo_gym/run.sh create mode 100644 tests/test_nemo_gym_compare.py create mode 100644 tests/test_nemo_gym_run.py diff --git a/benchmark/README.md b/benchmark/README.md index dfaa36a52..64c1bff6c 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -1,7 +1,12 @@ -# Harbor Benchmarks +# Switchyard Benchmarks + +For a fixed-model versus routed comparison with NeMo Gym, see the +[NeMo Gym example](nemo_gym/README.md). + +## Harbor Use this guide to run Harbor Terminal-Bench Lite from a fresh Switchyard clone. It covers the two smoke paths most people need first: diff --git a/benchmark/nemo_gym/README.md b/benchmark/nemo_gym/README.md new file mode 100644 index 000000000..43f6aacb4 --- /dev/null +++ b/benchmark/nemo_gym/README.md @@ -0,0 +1,121 @@ + + + +# Compare Switchyard routes with NeMo Gym + +This example evaluates the same MMLU-Redux tasks through two routes in [`routes.toml`](routes.toml): + +- `strong-only` always uses the strong target. +- `policy-model` uses the efficient model as a classifier, then routes to the efficient or strong + target. + +Gym owns the tasks, verifier, rewards, and rollout capture. Switchyard owns routing, model calls, +and routing statistics. Gym's `switchyard_model` adapter joins them over HTTP, so neither project +imports the other's core library. The script starts a fresh server from this Switchyard checkout +for each condition to isolate its statistics. + +## Prerequisites + +Install Python 3.13.14 or newer, Rust 1.96.1, Cargo, `curl`, Git, and `uv`. You also +need credentials for the model endpoints configured in [`routes.toml`](routes.toml). The bundled +configuration uses two NVIDIA-hosted models and reads `NVIDIA_API_KEY`; create a key on the +[NVIDIA API key page](https://build.nvidia.com/settings/api-keys). Switchyard can use another +supported provider or compatible endpoint by changing the LLM client and targets in the TOML. + +Use this clean, pinned Gym checkout: + +```bash +git clone https://github.com/NVIDIA-NeMo/Gym.git /path/to/Gym +git -C /path/to/Gym checkout e044a8ca795ece2c69b053d30c0a8dea7fa3b9f3 +cd /path/to/Gym +uv sync --frozen --no-dev +``` + +## Run + +From the Switchyard repository: + +```bash +export NVIDIA_API_KEY="nvapi-..." +export GYM_DIR=/path/to/Gym +bash benchmark/nemo_gym/run.sh +``` + +To use another provider, copy `routes.toml`, retain the `strong-only` and `policy-model` route IDs, +and update its LLM client, targets, and `api_key_env`. Export each credential named by the +deployment, then point the runner at that TOML: + +```bash +export OPENAI_API_KEY="..." +export GYM_DIR=/path/to/Gym +SWITCHYARD_CONFIG=/path/to/routes.toml bash benchmark/nemo_gym/run.sh +``` + +Switchyard validates the `api_key_env` entries when it loads the deployment. An unauthenticated +endpoint does not need a credential variable. + +Run `bash benchmark/nemo_gym/run.sh --help` to see the optional environment overrides. + +The default run evaluates five tasks and writes a timestamped directory under +`benchmark/nemo_gym/results/`. It is a workflow smoke test, not a benchmark result. The first run +also downloads and prepares the dataset. Gym starts its serving environment for each condition. +For a larger workflow check: + +```bash +LIMIT=100 REPEATS=3 CONCURRENCY=4 RESULTS_DIR=/tmp/routing-eval \ + bash benchmark/nemo_gym/run.sh +``` + +`LIMIT` takes the first tasks in Gym's prepared file, so a small limit is not a representative +sample. The example above can make roughly 600 answer calls and 300 classifier calls before +retries. Use a recorded stratified subset or the full benchmark for representative results. The +script refuses to overwrite a result directory. + +## Read the result + +`comparison.json` pairs completed rollouts by task and repeat index and verifies identical inputs. +It reports: + +- mean reward and routed-versus-baseline wins, ties, and losses; +- missing and unpaired completions, so failed tasks are not silently scored or discarded; +- paired answer-model tokens and endpoint latency from Gym; +- classifier tokens, answer and classifier latency, routing overhead, model totals, and + classifier fail-open counts from Switchyard. + +Gym's endpoint latency covers the whole routed request. Switchyard routing overhead includes the +classifier call, so routing overhead and classifier latency overlap and must not be added together. +Answer and classifier tokens remain separate. Switchyard totals are condition-wide, while quality +and answer usage are paired only over tasks completed by both conditions. + +The result directory contains the following artifacts. Gym may also write its best-effort +`switchyard-stats.json` wrapper when its shutdown hook completes: + +| Artifact | Meaning | +|---|---| +| `comparison.json` | Paired quality and usage comparison for both conditions. | +| `/rollouts.jsonl` | Completed Gym rollouts and rewards. | +| `/rollouts_materialized_inputs.jsonl` | Exact task/repeat inputs used by Gym. | +| `/rollouts_failures.jsonl` | Rollouts that failed before producing a scored result. | +| `/rollouts_aggregate_metrics.json` | Gym's aggregate benchmark metrics. | +| `/switchyard-condition.json` | Route and attached-proxy provenance written by Gym. | +| `/switchyard-stats-raw.json` | Raw `/v1/stats` captured while the proxy is alive. | +| `/switchyard-stats.json` | Best-effort Gym wrapper around `/v1/stats`. | +| `/switchyard-metrics.prom` | Prometheus metrics, including classifier fail-open reasons. | +| `/model-calls/` | Per-rollout model-call captures, including the served model. | +| `/routes.toml` | Exact Switchyard deployment copied for the condition. | +| `/switchyard.log` | Switchyard server output for diagnosis. | + +Gym excludes failure-sidecar rows from its reward calculation. `comparison.json` reports those +rollouts as missing or unpaired instead of treating them as zero-reward answers. + +Keep the materialized inputs because Gym's MMLU-Redux loader does not pin a Hugging Face dataset +revision. The workflow and evaluated inputs are reproducible; hosted model outputs, token counts, +and latency can still change between runs. Run from a clean Switchyard checkout when you need to +reproduce the exact server build. Dirty builds are labeled `-dirty`, but the source diff is not +archived with the result. + +See the +[NeMo Gym Switchyard model-server documentation](https://docs.nvidia.com/nemo/gym/main/model-server/switchyard/) +for other benchmarks and hosted mode, and the +[LLM classifier guide](../../docs/routing_algorithms/llm_classifier_routing.md) for the routing +policy used here. diff --git a/benchmark/nemo_gym/compare.py b/benchmark/nemo_gym/compare.py new file mode 100644 index 000000000..4fba01de2 --- /dev/null +++ b/benchmark/nemo_gym/compare.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compare paired quality and usage for two attached NeMo Gym runs. + +Per-rollout measurements include only shared task/repeat pairs. Switchyard proxy +statistics remain condition-wide. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +RolloutKey = tuple[int, int] + + +def _read_object(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as error: + raise ValueError(f"invalid JSON in {path}: {error.msg}") from error + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object in {path}") + return value + + +def _read_jsonl(path: Path) -> dict[RolloutKey, dict[str, Any]]: + rows: dict[RolloutKey, dict[str, Any]] = {} + with path.open(encoding="utf-8") as lines: + for line_number, line in enumerate(lines, start=1): + try: + row = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError(f"{path}:{line_number}: invalid JSON: {error.msg}") from error + if not isinstance(row, dict): + raise ValueError(f"{path}:{line_number}: expected a JSON object") + try: + task_index = row["_ng_task_index"] + rollout_index = row["_ng_rollout_index"] + except KeyError as error: + raise ValueError(f"{path}:{line_number}: missing {error.args[0]}") from error + if type(task_index) is not int or type(rollout_index) is not int: + raise ValueError(f"{path}:{line_number}: rollout indices must be integers") + key = (task_index, rollout_index) + if key in rows: + raise ValueError(f"duplicate rollout key {key} in {path}") + rows[key] = row + return rows + + +def _fail_opens(run_dir: Path) -> dict[str, int]: + """Aggregate Switchyard's fixed classifier fail-open metric by reason.""" + path = run_dir / "switchyard-metrics.prom" + if not path.exists(): + raise ValueError(f"missing {path}; capture /metrics before stopping the proxy") + + counts: dict[str, int] = {} + prefix = "switchyard_classifier_fail_open_total{" + marker = 'reason="' + for line in path.read_text().splitlines(): + if not line.startswith(prefix) or marker not in line: + continue + reason = line.split(marker, 1)[1].split('"', 1)[0] + value = int(float(line.rsplit(maxsplit=1)[1])) + counts[reason] = counts.get(reason, 0) + value + return counts + + +def _models(stats: dict[str, Any]) -> dict[str, dict[str, int | float]]: + return { + name: { + "calls": model["calls"], + "errors": model["errors"], + "tokens": model["total_tokens"], + "model_call_latency_mean_ms": model["model_call_latency"]["avg_ms"], + } + for name, model in stats.items() + } + + +def _reward(row: dict[str, Any], key: RolloutKey) -> float: + try: + value = float(row["reward"]) + except (KeyError, TypeError, ValueError) as error: + raise ValueError(f"rollout {key} has no numeric reward") from error + if not math.isfinite(value): + raise ValueError(f"rollout {key} has a non-finite reward") + return value + + +def _condition_summary( + run_dir: Path, + rows: dict[RolloutKey, dict[str, Any]], + shared: set[RolloutKey], +) -> dict[str, Any]: + """Combine paired rollout measurements with condition-wide proxy statistics.""" + stats = _read_object(run_dir / "switchyard-stats-raw.json") + rewards = [_reward(rows[key], key) for key in sorted(shared)] + try: + captures = [rows[key]["ng_model_call_capture"]["metrics"] for key in sorted(shared)] + except (KeyError, TypeError) as error: + raise ValueError( + "rollout is missing ng_model_call_capture.metrics; enable Gym observability" + ) from error + + try: + answer_model_tokens = sum(int(item["tokens_total"]) for item in captures) + endpoint_latency_mean_ms = sum(float(item["latency_total_ms"]) for item in captures) / len( + captures + ) + except (KeyError, TypeError, ValueError) as error: + raise ValueError( + "rollout capture metrics must include numeric tokens and latency" + ) from error + + classifier = stats["classifier"] + return { + "paired": { + "mean_reward": sum(rewards) / len(rewards), + "answer_model_tokens": answer_model_tokens, + "endpoint_latency_mean_ms": endpoint_latency_mean_ms, + }, + "condition_totals": { + "classifier_tokens": classifier["total_tokens"]["total"], + "routing_overhead_mean_ms": stats["routing_overhead"]["avg_ms"], + "classifier_fail_opens": _fail_opens(run_dir), + "answer_models": _models(stats["models"]), + "classifier_models": _models(classifier["models"]), + }, + } + + +def compare(baseline_dir: Path, routed_dir: Path) -> dict[str, Any]: + """Return paired quality and usage summaries for two result directories.""" + conditions = { + "baseline": _read_object(baseline_dir / "switchyard-condition.json"), + "routed": _read_object(routed_dir / "switchyard-condition.json"), + } + if any(condition.get("mode") != "attached" for condition in conditions.values()): + raise ValueError("this example compares attached Switchyard runs") + provenances = { + name: condition.get("proxy_provenance") for name, condition in conditions.items() + } + required_revisions = ("gym_revision", "switchyard_revision") + if any( + not isinstance(provenance, dict) + or any( + not isinstance(provenance.get(key), str) or not provenance[key] + for key in required_revisions + ) + for provenance in provenances.values() + ): + raise ValueError("the runs have incomplete Switchyard provenance") + provenance = provenances["baseline"] + if provenance != provenances["routed"]: + raise ValueError("the runs used different or incomplete Switchyard provenance") + if (baseline_dir / "routes.toml").read_bytes() != (routed_dir / "routes.toml").read_bytes(): + raise ValueError("the runs used different Switchyard deployments") + + runs = { + "baseline": _read_jsonl(baseline_dir / "rollouts.jsonl"), + "routed": _read_jsonl(routed_dir / "rollouts.jsonl"), + } + inputs = { + "baseline": _read_jsonl(baseline_dir / "rollouts_materialized_inputs.jsonl"), + "routed": _read_jsonl(routed_dir / "rollouts_materialized_inputs.jsonl"), + } + expected = set(inputs["baseline"]) + if expected != set(inputs["routed"]): + raise ValueError("the runs materialized different rollout keys") + # Matching indices are insufficient because the source dataset revision is not pinned. + for key in sorted(expected): + if inputs["baseline"][key] != inputs["routed"][key]: + raise ValueError(f"paired rollout {key} used different materialized inputs") + for name, rows in runs.items(): + if extra := set(rows) - expected: + raise ValueError(f"{name} completed unknown rollout keys: {sorted(extra)}") + + shared = set(runs["baseline"]) & set(runs["routed"]) + if not shared: + raise ValueError("the runs have no completed rollouts in common") + rewards = { + name: {key: _reward(rows[key], key) for key in shared} for name, rows in runs.items() + } + return { + "provenance": provenance, + "routes": {name: condition["route"] for name, condition in conditions.items()}, + "coverage": { + "expected": len(expected), + "paired": len(shared), + "completed": {name: len(rows) for name, rows in runs.items()}, + "unpaired": {name: len(set(rows) - shared) for name, rows in runs.items()}, + "missing": {name: len(expected - set(rows)) for name, rows in runs.items()}, + }, + "routed_vs_baseline": { + "wins": sum(rewards["routed"][key] > rewards["baseline"][key] for key in shared), + "ties": sum(rewards["routed"][key] == rewards["baseline"][key] for key in shared), + "losses": sum(rewards["routed"][key] < rewards["baseline"][key] for key in shared), + }, + "baseline": _condition_summary(baseline_dir, runs["baseline"], shared), + "routed": _condition_summary(routed_dir, runs["routed"], shared), + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("baseline_dir", type=Path, help="fixed-model condition directory") + parser.add_argument("routed_dir", type=Path, help="routed condition directory") + args = parser.parse_args() + try: + result = compare(args.baseline_dir, args.routed_dir) + except KeyError as error: + parser.exit(1, f"{parser.prog}: error: missing expected field {error}\n") + except (OSError, TypeError, ValueError) as error: + parser.exit(1, f"{parser.prog}: error: {error}\n") + print(json.dumps(result, indent=2)) + 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..03b5a8727 --- /dev/null +++ b/benchmark/nemo_gym/routes.toml @@ -0,0 +1,40 @@ +# 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" + +# Reuse the efficient model for classification; Switchyard tracks its judge and answer calls +# separately. +[targets.efficient] +id = "nvidia/nemotron-3.5-lightning-30b-a3b" +llm_client = "nvidia" +extra_body = { temperature = 0, chat_template_kwargs = { enable_thinking = false } } + +# The strong model is both the fixed baseline and the routed condition's capable target. +[targets.strong] +id = "nvidia/nemotron-3-super-120b-a12b" +llm_client = "nvidia" +extra_body = { chat_template_kwargs = { enable_thinking = false } } + +# Fixed control: every request is served by the strong target. +[routes.strong-only] +id = "strong-only" +type = "passthrough" +target = "strong" + +# Routed condition: the classifier chooses the weak or strong answer target. +[routes.policy-model] +id = "policy-model" +type = "llm_classifier" +mode = "capability" +classifier_target = "efficient" +strong_target = "strong" +weak_target = "efficient" +base_threshold = 0.5 +threshold_step = 0.1 +max_output_tokens = 1024 diff --git a/benchmark/nemo_gym/run.sh b/benchmark/nemo_gym/run.sh new file mode 100755 index 000000000..7bb651aa1 --- /dev/null +++ b/benchmark/nemo_gym/run.sh @@ -0,0 +1,168 @@ +#!/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 classifier-routed conditions over the same NeMo Gym rollouts. + +set -euo pipefail + +GYM_REVISION="e044a8ca795ece2c69b053d30c0a8dea7fa3b9f3" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SWITCHYARD_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +DEFAULT_DEPLOYMENT="$SCRIPT_DIR/routes.toml" +DEPLOYMENT="${SWITCHYARD_CONFIG:-$DEFAULT_DEPLOYMENT}" +RESULTS_DIR="${RESULTS_DIR:-$SCRIPT_DIR/results/$(date -u +%Y%m%dT%H%M%SZ)}" +PORT="${SWITCHYARD_PORT:-4000}" +LIMIT="${LIMIT:-5}" +REPEATS="${REPEATS:-1}" +CONCURRENCY="${CONCURRENCY:-1}" +[[ "$DEPLOYMENT" = /* ]] || DEPLOYMENT="$PWD/$DEPLOYMENT" +[[ "$RESULTS_DIR" = /* ]] || RESULTS_DIR="$PWD/$RESULTS_DIR" + +usage() { + cat <) + SWITCHYARD_PORT Local proxy port (default: 4000) + LIMIT Number of prepared tasks (default: 5) + REPEATS Rollouts per task (default: 1) + CONCURRENCY Concurrent Gym samples (default: 1) +EOF +} + +die() { + echo "error: $*" >&2 + exit 1 +} + +if [[ $# -gt 0 ]]; then + if [[ $# -eq 1 && ( "$1" == "-h" || "$1" == "--help" ) ]]; then + usage + exit 0 + fi + echo "error: unexpected argument: $1" >&2 + usage >&2 + exit 2 +fi + +[[ -n "${GYM_DIR:-}" ]] || die "set GYM_DIR to a NeMo Gym checkout at $GYM_REVISION" +[[ -f "$DEPLOYMENT" ]] || die "Switchyard configuration does not exist: $DEPLOYMENT" + +GYM="$GYM_DIR/.venv/bin/gym" +BENCHMARK_DATA="$GYM_DIR/benchmarks/mmlu-redux/data/mmlu-redux_benchmark.jsonl" +TARGET_DIR="${CARGO_TARGET_DIR:-$SWITCHYARD_ROOT/target}" +[[ "$TARGET_DIR" = /* ]] || TARGET_DIR="$PWD/$TARGET_DIR" +SERVER="$TARGET_DIR/release/switchyard-server" +[[ -x "$GYM" ]] || die "run 'uv sync --frozen --no-dev' in $GYM_DIR" +[[ "$(git -C "$GYM_DIR" rev-parse HEAD)" == "$GYM_REVISION" ]] || { + die "Gym must be checked out at $GYM_REVISION" +} +[[ -z "$(git -C "$GYM_DIR" status --porcelain)" ]] || { + die "Gym checkout must be clean so its recorded revision is exact" +} +[[ ! -e "$RESULTS_DIR" ]] || die "results path already exists: $RESULTS_DIR" +if curl -sS --max-time 1 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then + die "port $PORT is already in use; set SWITCHYARD_PORT to a free port" +fi + +echo "Results: $RESULTS_DIR" +cargo build --manifest-path "$SWITCHYARD_ROOT/Cargo.toml" --release -p switchyard-server +if [[ ! -s "$BENCHMARK_DATA" ]]; then + (cd "$GYM_DIR" && "$GYM" eval prepare --benchmark mmlu-redux) +fi + +SWITCHYARD_REVISION="$(git -C "$SWITCHYARD_ROOT" describe --always --dirty)" +SERVER_PID="" + +stop_server() { + if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then + kill -INT "$SERVER_PID" + wait "$SERVER_PID" || true + fi + SERVER_PID="" +} +trap stop_server EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +wait_for_server() { + local log_file="$1" + for _ in {1..80}; do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + cat "$log_file" >&2 + return 1 + fi + if curl -fsS "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then + return + fi + sleep 0.25 + done + echo "Switchyard did not become healthy; see $log_file" >&2 + return 1 +} + +run_condition() { + local route="$1" + local run_dir="$RESULTS_DIR/$route" + local root_url="http://127.0.0.1:$PORT" + local run_status=0 + echo "Running condition: $route" + mkdir -p "$run_dir/model-calls" + cp "$DEPLOYMENT" "$run_dir/routes.toml" + + # Router state and /v1/stats are process-wide, so each condition gets a fresh proxy. + "$SERVER" --config "$DEPLOYMENT" --host 127.0.0.1 --port "$PORT" \ + >"$run_dir/switchyard.log" 2>&1 & + SERVER_PID=$! + wait_for_server "$run_dir/switchyard.log" + + ( + cd "$GYM_DIR" || exit 1 + "$GYM" eval run \ + --benchmark mmlu-redux \ + --model-type switchyard_model \ + --model "$route" \ + --output "$run_dir/rollouts.jsonl" \ + --split benchmark \ + --limit "$LIMIT" \ + --num-repeats "$REPEATS" \ + --concurrency "$CONCURRENCY" \ + --temperature 0 \ + --max-output-tokens 1024 \ + +route_failures_to_sidecar=true \ + ++observability_enabled=true \ + ++model_call_capture_dir="$run_dir/model-calls" \ + ++policy_model.responses_api_models.switchyard_model.switchyard_base_url="$root_url/v1" \ + ++policy_model.responses_api_models.switchyard_model.condition_dir="$run_dir" \ + ++policy_model.responses_api_models.switchyard_model.proxy_provenance.gym_revision="$GYM_REVISION" \ + ++policy_model.responses_api_models.switchyard_model.proxy_provenance.switchyard_revision="$SWITCHYARD_REVISION" + ) || run_status=$? + + # Capture proxy diagnostics even when Gym fails. + if ! curl -fsS "$root_url/v1/stats" -o "$run_dir/switchyard-stats-raw.json"; then + echo "warning: could not capture Switchyard stats for $route" >&2 + [[ "$run_status" -ne 0 ]] || run_status=1 + fi + if ! curl -fsS "$root_url/metrics" -o "$run_dir/switchyard-metrics.prom"; then + echo "warning: could not capture Switchyard metrics for $route" >&2 + [[ "$run_status" -ne 0 ]] || run_status=1 + fi + stop_server + return "$run_status" +} + +run_condition strong-only +run_condition policy-model +"$GYM_DIR/.venv/bin/python" "$SCRIPT_DIR/compare.py" \ + "$RESULTS_DIR/strong-only" "$RESULTS_DIR/policy-model" | tee "$RESULTS_DIR/comparison.json" +echo "Comparison written to $RESULTS_DIR/comparison.json" diff --git a/tests/test_nemo_gym_compare.py b/tests/test_nemo_gym_compare.py new file mode 100644 index 000000000..f1e7cf48e --- /dev/null +++ b/tests/test_nemo_gym_compare.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import runpy +from pathlib import Path +from typing import Any + +import pytest + +COMPARATOR = Path(__file__).parents[1] / "benchmark" / "nemo_gym" / "compare.py" +compare = runpy.run_path(str(COMPARATOR))["compare"] +RolloutKey = tuple[int, int] + + +def _model(calls: int, tokens: int, latency_ms: float) -> dict[str, Any]: + return { + "calls": calls, + "errors": 0, + "total_tokens": tokens, + "model_call_latency": {"avg_ms": latency_ms}, + } + + +def _write_run( + root: Path, + rewards: dict[RolloutKey, float | None], + *, + classifier_tokens: int, + input_suffix: str = "", + provenance: str = "same-build", + fail_open_reason: str | None = None, +) -> None: + root.mkdir() + rollouts: list[dict[str, Any]] = [] + inputs: list[dict[str, Any]] = [] + for (task, rollout), reward in rewards.items(): + key = {"_ng_task_index": task, "_ng_rollout_index": rollout} + inputs.append({**key, "prompt": f"task-{task}{input_suffix}"}) + if reward is None: + continue + rollouts.append( + { + **key, + "reward": reward, + "ng_model_call_capture": { + "metrics": { + "tokens_total": 10 + task + rollout, + "latency_total_ms": 20.0 + task + rollout, + } + }, + } + ) + answer_tokens = sum( + 10 + task + rollout for (task, rollout), reward in rewards.items() if reward is not None + ) + (root / "rollouts.jsonl").write_text( + "".join(json.dumps(row) + "\n" for row in rollouts), encoding="utf-8" + ) + (root / "rollouts_materialized_inputs.jsonl").write_text( + "".join(json.dumps(row) + "\n" for row in inputs), encoding="utf-8" + ) + (root / "switchyard-stats-raw.json").write_text( + json.dumps( + { + "classifier": { + "total_tokens": {"total": classifier_tokens}, + "models": {"judge": _model(1, classifier_tokens, 4.0)} + if classifier_tokens + else {}, + }, + "routing_overhead": {"avg_ms": 3.5}, + "models": {"model-a": _model(len(rollouts), answer_tokens, 8.0)}, + } + ), + encoding="utf-8", + ) + (root / "switchyard-condition.json").write_text( + json.dumps( + { + "route": root.name, + "mode": "attached", + "proxy_provenance": { + "gym_revision": "gym-revision", + "switchyard_revision": provenance, + }, + } + ), + encoding="utf-8", + ) + metric = "" + if fail_open_reason is not None: + metric = ( + "switchyard_classifier_fail_open_total" + f'{{judge_model="classifier",reason="{fail_open_reason}"}} 1\n' + ) + (root / "switchyard-metrics.prom").write_text(metric, encoding="utf-8") + (root / "routes.toml").write_text("same deployment", encoding="utf-8") + + +def test_compares_paired_quality_and_usage(tmp_path: Path) -> None: + baseline = tmp_path / "baseline" + routed = tmp_path / "routed" + _write_run(baseline, {(0, 0): 0.0, (0, 1): 1.0}, classifier_tokens=0) + _write_run( + routed, + {(0, 0): 1.0, (0, 1): None}, + classifier_tokens=17, + fail_open_reason="parse_error", + ) + + result = compare(baseline, routed) + + assert result["coverage"]["paired"] == 1 + assert result["coverage"]["unpaired"] == {"baseline": 1, "routed": 0} + assert result["coverage"]["missing"] == {"baseline": 0, "routed": 1} + assert result["routed_vs_baseline"] == {"wins": 1, "ties": 0, "losses": 0} + assert result["baseline"]["paired"]["mean_reward"] == 0.0 + assert result["routed"]["paired"]["mean_reward"] == 1.0 + assert result["baseline"]["paired"]["answer_model_tokens"] == 10 + assert result["baseline"]["paired"]["endpoint_latency_mean_ms"] == 20.0 + assert result["routed"]["condition_totals"]["classifier_tokens"] == 17 + assert result["routed"]["condition_totals"]["classifier_fail_opens"] == {"parse_error": 1} + answer = result["routed"]["condition_totals"]["answer_models"]["model-a"] + assert answer["model_call_latency_mean_ms"] == 8.0 + + +def test_rejects_different_materialized_inputs(tmp_path: Path) -> None: + baseline = tmp_path / "baseline" + routed = tmp_path / "routed" + _write_run(baseline, {(0, 0): 1.0}, classifier_tokens=0) + _write_run(routed, {(0, 0): 1.0}, classifier_tokens=5, input_suffix="-changed") + + with pytest.raises(ValueError, match="different materialized inputs"): + compare(baseline, routed) + + +def test_rejects_different_switchyard_builds(tmp_path: Path) -> None: + baseline = tmp_path / "baseline" + routed = tmp_path / "routed" + _write_run(baseline, {(0, 0): 1.0}, classifier_tokens=0, provenance="build-a") + _write_run(routed, {(0, 0): 1.0}, classifier_tokens=5, provenance="build-b") + + with pytest.raises(ValueError, match="different or incomplete Switchyard provenance"): + compare(baseline, routed) + + +def test_rejects_non_finite_rewards(tmp_path: Path) -> None: + baseline = tmp_path / "baseline" + routed = tmp_path / "routed" + _write_run(baseline, {(0, 0): float("nan")}, classifier_tokens=0) + _write_run(routed, {(0, 0): 1.0}, classifier_tokens=5) + + with pytest.raises(ValueError, match="non-finite reward"): + compare(baseline, routed) diff --git a/tests/test_nemo_gym_run.py b/tests/test_nemo_gym_run.py new file mode 100644 index 000000000..08aafbcad --- /dev/null +++ b/tests/test_nemo_gym_run.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import subprocess +from pathlib import Path + +RUNNER = Path(__file__).parents[1] / "benchmark" / "nemo_gym" / "run.sh" + + +def test_help_documents_custom_switchyard_config() -> None: + result = subprocess.run( + ["bash", str(RUNNER), "--help"], + check=True, + capture_output=True, + text=True, + ) + + assert "SWITCHYARD_CONFIG" in result.stdout + assert "NVIDIA_API_KEY" not in result.stdout + + +def test_accepts_relative_custom_config_without_nvidia_key(tmp_path: Path) -> None: + (tmp_path / "routes.toml").write_text("schema_version = 1\n", encoding="utf-8") + result = subprocess.run( + ["bash", str(RUNNER)], + cwd=tmp_path, + env={ + "PATH": os.environ["PATH"], + "GYM_DIR": str(tmp_path / "gym"), + "SWITCHYARD_CONFIG": "routes.toml", + }, + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 1 + assert "run 'uv sync --frozen --no-dev'" in result.stderr + assert "NVIDIA_API_KEY" not in result.stderr