diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 0973fc23e9..0ba0738b0b 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -381,6 +381,50 @@ jobs: ;; esac + benchmark-guardrails: + name: Guardrails plugin benchmark + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - name: Checkout nemo-platform + uses: actions/checkout@v6 + with: + path: nemo-platform + - name: Checkout NeMo-Guardrails + uses: actions/checkout@v6 + with: + repository: NVIDIA/NeMo-Guardrails + path: NeMo-Guardrails + - name: Install uv + uses: astral-sh/setup-uv@v7 + with: + # Point setup-uv at the nemo-platform checkout so it picks up the + # repo's `required-version` constraint from pyproject.toml. Other + # jobs check out into the workspace root and get this implicitly. + working-directory: nemo-platform + python-version: "3.11" + enable-cache: true + - name: Bootstrap Python environment + working-directory: nemo-platform + run: make bootstrap-python + env: + PYTORCH_DEPS: cpu + - name: Run benchmark sweep + working-directory: nemo-platform + run: make benchmark-guardrails + env: + NEMO_GUARDRAILS_REPO_ROOT: ${{ github.workspace }}/NeMo-Guardrails + _TYPER_FORCE_DISABLE_TERMINAL: "1" + - name: Upload benchmark artifacts + if: always() + uses: actions/upload-artifact@v6 + with: + name: benchmark-guardrails-results + retention-days: 30 + path: | + nemo-platform/plugins/nemo-guardrails/benchmarks/artifacts/runs/ + coverage-comment: name: Post coverage comment needs: [python-unit-test, python-integration-test] diff --git a/Makefile b/Makefile index 7709b62d2d..7e4a4930d6 100644 --- a/Makefile +++ b/Makefile @@ -440,3 +440,9 @@ test-e2e-kubernetes-gpu: ## Run GPU e2e tests against Kubernetes (requires GPU n test-e2e-kubernetes-gpu-customizer: ## Run GPU customizer e2e tests against Kubernetes (requires GPU nodes; set NMP_E2E_CLUSTER_URL) @echo "Running GPU customizer e2e tests with Kubernetes..." uv run --frozen pytest e2e/test_customizer.py --kubernetes --feature gpu --feature customizer --log-cli-level=INFO -v --junitxml=report-kubernetes-gpu-customizer.xml + +.PHONY: benchmark-guardrails +benchmark-guardrails: ## Run nemo-guardrails IGW benchmark sweep (set BENCHMARK_ARGS for extra flags) + @echo "Running nemo-guardrails IGW benchmark..." + uv run --frozen --package nemo-guardrails-plugin --extra bench \ + python -m nemo_guardrails_plugin.benchmarks.run $(BENCHMARK_ARGS) diff --git a/plugins/nemo-guardrails/benchmarks/README.md b/plugins/nemo-guardrails/benchmarks/README.md new file mode 100644 index 0000000000..4a923f9833 --- /dev/null +++ b/plugins/nemo-guardrails/benchmarks/README.md @@ -0,0 +1,234 @@ +# NeMo Guardrails Plugin Benchmarks + +Local harness for benchmarking the `nemo-guardrails` Inference Gateway +middleware against the upstream NeMo Guardrails benchmark suite (mock LLMs + +AIPerf sweep). + +The implementation lives in `nemo_guardrails_plugin.benchmarks` (under +`plugins/nemo-guardrails/src/`). The harness does **not** copy benchmark code +from the NeMo Guardrails repository; it expects a local checkout and runs its +benchmark modules with `PYTHONPATH` pointed at that checkout. + +## Layout + +```text +plugins/nemo-guardrails/benchmarks/ + configs/ + nmp_igw_guardrails_sweep_concurrency.yaml # AIPerf sweep template + artifacts/ # per-run outputs (gitignored) +plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/ + run.py # entrypoint: `python -m nemo_guardrails_plugin.benchmarks.run` + paths.py # filesystem layout + constants.py # workspace / VM / provider names + processes.py # subprocess supervision (process groups + ExitStack) + seeding.py # NMP SDK calls to create the required entities + aiperf_runner.py # rewrite AIPerf config + invoke upstream sweep + collect results + bootstrap.py # manage the isolated venv that hosts the `aiperf` CLI + shim.py # tiny HTTP shim that satisfies AIPerf's `/v1/models` pre-check +``` + +## Prerequisites + +- This repo bootstrapped via `make bootstrap-python` (the harness runs in + `.venv` and imports the NMP SDK from the workspace). +- A local NeMo Guardrails checkout. By default the harness looks at + `../NeMo-Guardrails` relative to the NMP repo root. +- `uv` available on `PATH`. +- Ports `8000`, `8001`, `8080`, and `8090` available — unless you opt into + reusing an already-running local NMP via `--reuse-services`. Port `8090` is + used by an internal shim that satisfies AIPerf's hard-coded `/v1/models` + health probe. + +The harness has its own dependencies that are declared as +the `bench` extra on `nemo-guardrails-plugin`. The `make benchmark-guardrails` +target installs them automatically via `uv run --extra bench`; they are not +part of the plugin's runtime install. + +The upstream `aiperf` CLI itself pins `aiofiles<24.2`, which conflicts with +NMP's evaluator-service. To avoid downgrading the shared workspace venv, the +harness creates an isolated venv at +`plugins/nemo-guardrails/benchmarks/artifacts/venvs/aiperf/` on first run and +reuses it on subsequent runs. CI gets a fresh one each invocation; locally +this caches across runs for fast iteration. + +## Run locally + +From the NMP repo root: + +```bash +make benchmark-guardrails +``` + +If your NeMo Guardrails checkout is somewhere else: + +```bash +NEMO_GUARDRAILS_REPO_ROOT=/path/to/NeMo-Guardrails make benchmark-guardrails +``` + +To pass through arbitrary harness flags: + +```bash +make benchmark-guardrails BENCHMARK_ARGS="--verbose --reuse-services" +``` + +The default sweep runs concurrency levels: + +```text +1, 2, 4, 8, 16, 32, 64 +``` + +With the default 60-second benchmark duration, expect the benchmark to run for ~10 minutes after service bootstrap. + +### Monitoring progress + +After bootstrap and seeding, the terminal prints `Running aiperf sweep: ...` while +AIPerf runs. AIPerf stdout/stderr is redirected to `logs/aiperf.log`, not the +harness terminal. When the run finishes successfully, the last harness line looks +like: + +```text +Sweep summary: 7 run(s), 0 failure(s); per-sweep outputs under ... +``` + +**Tail the sweep log** (replace `` with your `--run-id` or timestamp +directory name): + +```bash +tail -f plugins/nemo-guardrails/benchmarks/artifacts/runs//logs/aiperf.log +``` + +Look for lines like `Run 3/7` and `Run 3 completed successfully`. + +**Watch completed sweep directories**: + +```bash +ls plugins/nemo-guardrails/benchmarks/artifacts/runs//aiperf_results/*/*/ +``` + +Each finished level appears as `concurrency1/`, `concurrency2/`, etc., with +`process_result.json` and `profile_export_aiperf.csv` inside. A directory with +an empty `profile_export.jsonl` is usually the sweep currently in progress. + +**Confirm the process is still running**: + +```bash +pgrep -fl "benchmark.aiperf" +``` + +**Normal log noise during bootstrap**: + +- With the `--verbose` flag, `Connection refused` on `:8080` for up to ~1–3 minutes while `nemo services run` starts. +- `409 Conflict` during seeding when resources from a prior run already exist. +- One failed smoke-test attempt (`404` on the VirtualModel) before the IGW route propagates. + +## What the harness starts + +- The upstream benchmark **mock app LLM** on `http://localhost:8000`, +- The upstream benchmark **mock content-safety LLM** on `http://localhost:8001`, +- Local **NMP services** on `http://localhost:8080`, unless `--reuse-services`. + +It then seeds NMP via the SDK with: + +- workspace `benchmark`, +- app model provider `benchmark-app-llm`, +- content-safety model provider `benchmark-content-safety-llm`, +- guardrail config `content-safety-local`, +- VirtualModel `benchmark/guardrails-vm` with `nemo-guardrails` attached to + both request and response middleware. + +The benchmark target for inference requests is: + +```text +http://localhost:8080/apis/inference-gateway/v2/workspaces/benchmark/openai/-/v1/chat/completions +``` + +## Useful flags / environment + +The harness accepts both CLI flags and environment variables: + +| CLI flag | Environment variable | Default | +|-----------------------------------|----------------------------------|------------------------| +| `--nemo-guardrails-repo-root` | `NEMO_GUARDRAILS_REPO_ROOT` | `../NeMo-Guardrails` | +| `--reuse-services` | `NMP_BENCHMARK_REUSE_SERVICES=1` | start `nemo services run` | +| `--keep-running` | `NMP_BENCHMARK_KEEP_RUNNING=1` | tear down on exit | +| `--mock-workers` | `NMP_BENCHMARK_MOCK_WORKERS` | `4` | +| `--run-id` | _n/a_ | current timestamp | + +`--keep-running` leaves child processes alive for post-mortem inspection; the +harness logs each child's PID as it starts. + +## Outputs + +Each run writes artifacts under: + +```text +plugins/nemo-guardrails/benchmarks/artifacts/runs// + logs/ + mock-app-llm.log + mock-content-safety-llm.log + nmp-services.log + aiperf.log + generated/ + app_provider.json + content_safety_provider.json + virtual_model.json + content_safety_local_nmp_request.json + nmp_igw_guardrails_sweep_concurrency.yaml # runtime AIPerf config + aiperf_results//// + run_metadata.json + process_result.json + profile_export*.json # written by aiperf +``` + +## CI + +A `benchmark-guardrails` job in `.github/workflows/ci.yaml` checks out both +this repo and `NVIDIA/NeMo-Guardrails`, runs `make bootstrap-python` and +`make benchmark-guardrails`, and uploads the per-run artifacts directory +(`logs/`, `generated/`, `aiperf_results/`) on success or failure. + +Pass/fail is driven by the harness's exit code, which is non-zero if `aiperf` +itself exits non-zero or any sweep returns a non-zero exit code. No latency +thresholds are enforced — those can be layered on later by a separate +analyzer that reads the per-sweep CSVs. + +## Cleanup + +The harness only stops the processes it started. It will not kill unrelated +processes on ports `8000`, `8001`, `8080`, or `8090`. + +Across runs, NMP's data dir is reused so subsequent benchmarks start +faster. The harness redirects NMP's writes via the `NMP_DATA_DIR` env var to +a per-checkout directory: + +```text +plugins/nemo-guardrails/benchmarks/artifacts/nmp-data +``` + +This holds NMP's SQLite database, Secrets vault, and other persistent +service state. It is gitignored and isolated from `~/.nmp/`, so the +benchmark will not pollute your normal local NMP setup. Reuse is what +lets repeat runs skip the workspace / provider / VirtualModel creation +work (the harness treats `409 Conflict` as "already exists, carry on"). + +If a benchmark misbehaves and you suspect stale state (ex. after pulling +a schema change), delete that directory for a fully fresh run: + +```bash +rm -rf plugins/nemo-guardrails/benchmarks/artifacts/nmp-data +``` + +To remove outputs from a specific run (logs, generated configs, AIPerf results): + +```bash +rm -rf plugins/nemo-guardrails/benchmarks/artifacts/runs/ +``` + +To clear all run outputs: + +```bash +rm -rf plugins/nemo-guardrails/benchmarks/artifacts/runs/* +``` + +In CI this is automatic — every job gets a fresh runner, so `nmp-data` +does not exist. diff --git a/plugins/nemo-guardrails/benchmarks/artifacts/.gitignore b/plugins/nemo-guardrails/benchmarks/artifacts/.gitignore new file mode 100644 index 0000000000..d6b7ef32c8 --- /dev/null +++ b/plugins/nemo-guardrails/benchmarks/artifacts/.gitignore @@ -0,0 +1,2 @@ +* +!.gitignore diff --git a/plugins/nemo-guardrails/benchmarks/configs/nmp_igw_guardrails_sweep_concurrency.yaml b/plugins/nemo-guardrails/benchmarks/configs/nmp_igw_guardrails_sweep_concurrency.yaml new file mode 100644 index 0000000000..6fe53c5c93 --- /dev/null +++ b/plugins/nemo-guardrails/benchmarks/configs/nmp_igw_guardrails_sweep_concurrency.yaml @@ -0,0 +1,33 @@ +# Benchmark NMP IGW VirtualModel with nemo-guardrails request/response middleware. + +batch_name: nmp_igw_guardrails_sweep_concurrency +output_base_dir: plugins/nemo-guardrails/benchmarks/artifacts/aiperf_results + +base_config: + model: benchmark/guardrails-vm + # AIPerf only needs the tokenizer to count input/output tokens for its + # metrics; the actual model under test is the mock IGW VirtualModel above. + # We use Qwen3-8B (instead of Llama 3.3, which is gated) so CI can fetch + # the tokenizer without an HF Token. + tokenizer: Qwen/Qwen3-8B + # The benchmark runs a required health check at `/v1/models`. Until this + # endpoint is configurable, we run a thin shim on :8090 that satisfies that probe. + # See nemo_guardrails_plugin.benchmarks.shim for the implementation. + url: "http://localhost:8090" + endpoint: "/v1/chat/completions" + endpoint_type: chat + streaming: false + + warmup_request_count: 10 + benchmark_duration: 60 + concurrency: 0 + request_rate_mode: constant + + random_seed: 12345 + prompt_input_tokens_mean: 100 + prompt_input_tokens_stddev: 10 + prompt_output_tokens_mean: 50 + prompt_output_tokens_stddev: 5 + +sweeps: + concurrency: [1, 2, 4, 8, 16, 32, 64] diff --git a/plugins/nemo-guardrails/pyproject.toml b/plugins/nemo-guardrails/pyproject.toml index 72b07b447d..83a6db4797 100644 --- a/plugins/nemo-guardrails/pyproject.toml +++ b/plugins/nemo-guardrails/pyproject.toml @@ -12,6 +12,19 @@ dependencies = [ "dataclasses-json>=0.6.7", ] +[project.optional-dependencies] +# Dependencies for the local IGW benchmark harness in +# `nemo_guardrails_plugin.benchmarks` +bench = [ + "httpx>=0.27", + "pyyaml>=6.0", + # NOTE: aiperf itself is *not* listed here. Its 0.x line pins + # aiofiles<24.2 which conflicts with evaluator-service's + # aiofiles>=25.1. To avoid downgrading the shared workspace venv, the + # harness installs aiperf into a dedicated venv at run time; see + # `nemo_guardrails_plugin.benchmarks.bootstrap`. +] + [project.entry-points."nemo.inference_middleware"] nemo-guardrails = "nemo_guardrails_plugin.middleware:GuardrailsMiddleware" diff --git a/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/aiperf_runner.py b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/aiperf_runner.py new file mode 100644 index 0000000000..f7ff6729c1 --- /dev/null +++ b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/aiperf_runner.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Invoke the upstream NeMo Guardrails AIPerf sweep runner as a subprocess. + +The upstream code lives in ``${NEMO_GUARDRAILS_REPO_ROOT}/benchmark/`` and is not +installed as a package; we set ``PYTHONPATH`` so ``python -m benchmark.aiperf`` +resolves. +""" + +from __future__ import annotations + +import json +import logging +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import yaml +from nemo_guardrails_plugin.benchmarks.bootstrap import build_env + +log = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class SweepRunResult: + """Outcome of a single ``aiperf profile`` invocation (one concurrency level).""" + + # Name of the per-sweep subdirectory AIPerf created, e.g. ``"concurrency16"``. + # Surfaced in the harness's summary log so failures point at a directory name. + sweep_label: str + # Absolute path to the sweep's output directory under + # ``aiperf_results////``. Contains the AIPerf + # CSV, per-request JSONL, run metadata, and the wrapper's process_result.json. + output_dir: Path + # Exit code of the AIPerf subprocess for this sweep (0 = success). Sourced from + # ``process_result.json``; defaults to 1 if that file is missing or unparseable + # so a crashed subprocess does not silently pass. + return_code: int + # Wall-clock duration of the sweep in seconds, read from ``run_metadata.json``. + # 0.0 when AIPerf never wrote metadata (typically because it crashed early). + duration_seconds: float + # Path to ``run_metadata.json`` if AIPerf wrote it, else ``None``. Kept on the + # result for downstream analyzers that want to inspect AIPerf's view of the run. + metadata_path: Path | None + # Path to ``process_result.json`` if AIPerf wrote it, else ``None``. This is + # the upstream wrapper's record of the subprocess exit; useful when + # ``return_code`` itself looks suspicious. + process_result_path: Path | None + + @property + def passed(self) -> bool: + return self.return_code == 0 + + +def prepare_runtime_aiperf_config( + *, + template_path: Path, + runtime_config_path: Path, + aiperf_output_dir: Path, +) -> dict[str, Any]: + """Materialize the AIPerf config this run will use. + + Reads the checked-in ``template_path`` config, overrides its + ``output_base_dir`` to point inside the current run's directory, and writes + the result to ``runtime_config_path``. AIPerf is later invoked with + ``--config-file `` so every artifact lands under a + separate per-run directory. + + Returns the parsed config dict so callers can log fields (sweep params, + benchmark_duration) without re-reading the file. + """ + if not template_path.is_file(): + raise FileNotFoundError(f"AIPerf template not found: {template_path}") + + config = yaml.safe_load(template_path.read_text(encoding="utf-8")) + if not isinstance(config, dict): + raise ValueError(f"Expected a YAML mapping at {template_path}, got {type(config).__name__}") + + # Point AIPerf's output_base_dir at this run's directory so its results + # nest under our per-run artifacts tree. + config["output_base_dir"] = str(aiperf_output_dir) + runtime_config_path.parent.mkdir(parents=True, exist_ok=True) + runtime_config_path.write_text(yaml.safe_dump(config, sort_keys=False), encoding="utf-8") + + return config + + +def run_aiperf_sweep( + *, + nemoguardrails_repo_root: Path, + runtime_config: Path, + log_path: Path, + python_executable: str | None = None, + venv_bin_path: Path | str | None = None, + extra_env: dict[str, str] | None = None, +) -> int: + """Run ``python -m benchmark.aiperf --config-file ...`` and tee output. + + Returns the subprocess exit code. The caller decides whether to treat a + non-zero code as a sweep-level failure or as a fail-fast. + + ``python_executable`` should point at the python in the dedicated aiperf + venv (see :mod:`nemo_guardrails_plugin.benchmarks.bootstrap`). ``venv_bin_path`` + prepends that venv's ``bin/`` to ``PATH`` so the ``aiperf`` CLI is resolvable + when the wrapper shells out to it. + + Note: AIPerf's built-in pre-flight check does a GET on + ``urljoin(config.base_config.url, "/v1/models")`` with no override + available upstream yet. The harness runs a tiny shim that satisfies this + probe; see :mod:`nemo_guardrails_plugin.benchmarks.shim`. + """ + python_bin = python_executable or sys.executable + # Upstream `benchmark.aiperf` is a single-command Typer app; invoking + # `python -m benchmark.aiperf --config-file ...` runs the only command. + # Passing a literal `run` subcommand confuses Typer. + cmd = [ + python_bin, + "-m", + "benchmark.aiperf", + "--config-file", + str(runtime_config), + ] + + env = build_env( + venv_bin_path=venv_bin_path, + extra_env={ + "PYTHONPATH": str(nemoguardrails_repo_root), + **(extra_env or {}), + }, + ) + + log_path.parent.mkdir(parents=True, exist_ok=True) + + log.info("Running aiperf sweep: %s (cwd=%s)", " ".join(cmd), nemoguardrails_repo_root) + + with log_path.open("wb") as log_fh: + proc = subprocess.run( # noqa: S603 - command is constructed internally + cmd, + cwd=str(nemoguardrails_repo_root), + env=env, + stdout=log_fh, + stderr=subprocess.STDOUT, + check=False, + ) + + return proc.returncode + + +def collect_sweep_results(aiperf_output_dir: Path) -> list[SweepRunResult]: + """Walk the AIPerf output tree and surface per-sweep exit status. + + Layout: + ``////{run_metadata.json,process_result.json}`` + + Missing ``process_result.json`` is treated as a failure for that sweep so a + crashed AIPerf subprocess does not silently pass. + """ + if not aiperf_output_dir.is_dir(): + return [] + + results: list[SweepRunResult] = [] + for batch_dir in sorted(p for p in aiperf_output_dir.iterdir() if p.is_dir()): + for timestamp_dir in sorted(p for p in batch_dir.iterdir() if p.is_dir()): + for sweep_dir in sorted(p for p in timestamp_dir.iterdir() if p.is_dir()): + results.append(_load_sweep_result(sweep_dir)) + return results + + +def _load_sweep_result(sweep_dir: Path) -> SweepRunResult: + """Load the sweep result from the given directory. + + Layout: + ``////{run_metadata.json,process_result.json}`` + """ + metadata_path = sweep_dir / "run_metadata.json" + process_result_path = sweep_dir / "process_result.json" + + return_code = 1 + duration = 0.0 + + if process_result_path.is_file(): + try: + data = json.loads(process_result_path.read_text(encoding="utf-8")) + return_code = int(data.get("returncode", 1)) + except (json.JSONDecodeError, OSError, ValueError, TypeError) as exc: + log.warning("Could not parse %s: %s", process_result_path, exc) + + if metadata_path.is_file(): + try: + md = json.loads(metadata_path.read_text(encoding="utf-8")) + duration = float(md.get("duration_seconds", 0.0) or 0.0) + except (json.JSONDecodeError, OSError, ValueError, TypeError) as exc: + log.warning("Could not parse %s: %s", metadata_path, exc) + + return SweepRunResult( + sweep_label=sweep_dir.name, + output_dir=sweep_dir, + return_code=return_code, + duration_seconds=duration, + metadata_path=metadata_path if metadata_path.is_file() else None, + process_result_path=process_result_path if process_result_path.is_file() else None, + ) diff --git a/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/bootstrap.py b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/bootstrap.py new file mode 100644 index 0000000000..761a5f56c3 --- /dev/null +++ b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/bootstrap.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bootstrap an isolated venv for the upstream AIPerf load generator. + +``aiperf`` pins ``aiofiles<24.2`` which conflicts with NMP's evaluator-service +requirement of ``aiofiles>=25.1``, so we install it into a dedicated venv +instead of the shared workspace one. The venv is reused across local runs; +CI gets a fresh one each invocation. +""" + +from __future__ import annotations + +import logging +import os +import subprocess +from pathlib import Path + +log = logging.getLogger("nemo_guardrails_plugin.benchmarks.bootstrap") + +# Mirrors the upstream NeMo-Guardrails AIPerf README. We don't pin further; +# aiperf itself pins its transitives. +_AIPERF_PACKAGES = ("aiperf", "huggingface_hub", "typer>=0.9", "httpx>=0.27") + + +def ensure_aiperf_venv(venv_dir: Path) -> Path: + """Idempotently create the aiperf venv. Returns the venv's python path. + + Uses ``uv venv`` + ``uv pip install`` since the harness is only ever invoked + via ``make benchmark-guardrails``, which already requires ``uv`` to be on + PATH. Skips both steps if the venv and the ``aiperf`` binary already exist. + """ + python_bin = venv_dir / "bin" / "python" + aiperf_bin = venv_dir / "bin" / "aiperf" + + if aiperf_bin.exists() and python_bin.exists(): + log.info("Reusing existing aiperf venv at %s", venv_dir) + return python_bin + + log.info("Creating aiperf venv at %s", venv_dir) + venv_dir.parent.mkdir(parents=True, exist_ok=True) + subprocess.run( # noqa: S603 - command is constructed internally + ["uv", "venv", "--python", "3.11", str(venv_dir)], + check=True, + capture_output=True, + ) + + log.info("Installing %s into %s", ", ".join(_AIPERF_PACKAGES), venv_dir) + subprocess.run( # noqa: S603 - command is constructed internally + ["uv", "pip", "install", "--python", str(python_bin), *_AIPERF_PACKAGES], + check=True, + ) + + if not aiperf_bin.exists(): + raise RuntimeError(f"aiperf install completed but {aiperf_bin} is missing") + return python_bin + + +def build_env( + *, + venv_bin_path: Path | str | None = None, + extra_env: dict[str, str] | None = None, +) -> dict[str, str]: + """Return a child-process environment based on ``os.environ``. + + Optionally prepends ``venv_bin_path`` to ``PATH`` and overlays ``extra_env``. + """ + env = dict(os.environ) + if venv_bin_path: + bin_path = str(venv_bin_path) + env["PATH"] = f"{bin_path}{os.pathsep}{env.get('PATH', '')}" + if extra_env: + env.update(extra_env) + return env diff --git a/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/constants.py b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/constants.py new file mode 100644 index 0000000000..cbe4b19960 --- /dev/null +++ b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/constants.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared constants for the benchmark harness.""" + +from __future__ import annotations + +WORKSPACE = "benchmark" +GUARDRAIL_CONFIG = "content-safety-local" +VM_NAME = "guardrails-vm" + +# ModelProvider that proxies requests to the mock main model +APP_PROVIDER = "benchmark-app-llm" +APP_PROVIDER_URL = "http://localhost:8000" +APP_MODEL_NAME = "meta/llama-3.3-70b-instruct" + +# ModelProvider that proxies requests to the mock content-safety model +CS_PROVIDER = "benchmark-content-safety-llm" +CS_PROVIDER_URL = "http://localhost:8001" +CS_MODEL_NAME = "nvidia/llama-3.1-nemoguard-8b-content-safety" + +NMP_BASE_URL = "http://localhost:8080" +NMP_HEALTH_PATH = "/health/ready" +IGW_CHAT_PATH = f"/apis/inference-gateway/v2/workspaces/{WORKSPACE}/openai/-/v1/chat/completions" + +# Local shim that satisfies AIPerf's pre-check and reverse-proxies chat +# completion requests through to NMP's IGW. See +# `nemo_guardrails_plugin.benchmarks.shim` for the implementation. +AIPERF_SHIM_HOST = "127.0.0.1" +AIPERF_SHIM_PORT = 8090 +AIPERF_SHIM_BASE_URL = f"http://{AIPERF_SHIM_HOST}:{AIPERF_SHIM_PORT}" + +GUARDRAILS_MIDDLEWARE_NAME = "nemo-guardrails" +GUARDRAILS_MIDDLEWARE_CONFIG_TYPE = "guardrail_config" diff --git a/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/paths.py b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/paths.py new file mode 100644 index 0000000000..b1113fa2c2 --- /dev/null +++ b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/paths.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Filesystem layout for the nemo-guardrails IGW benchmark harness. + +A benchmark invocation produces a `runs//` directory that holds +everything generated by that run: AIPerf results, supervised-process logs, +and the materialized AIPerf config. Some artifacts are intentionally kept +*outside* the per-run dir so they can be reused across runs (the NMP data +dir, the aiperf venv). + +`RunPaths` is the single source of truth for those locations; `build_run_paths` +composes it from the two repo roots the harness needs (NMP and NeMo-Guardrails). +""" + +from __future__ import annotations + +import datetime as dt +from dataclasses import dataclass +from pathlib import Path + + +@dataclass(frozen=True) +class RunPaths: + """Filesystem layout for a single benchmark invocation.""" + + # Root of the nemo-platform repo we're running from. Used as the cwd for + # `nemo services run`. + nmp_repo_root: Path + # Root of a local NeMo-Guardrails checkout. The harness reads mock-LLM + # configs and the content_safety_local example from here. + nemoguardrails_repo_root: Path + # Convenience anchor for everything under + # `plugins/nemo-guardrails/benchmarks/` in the NMP repo. + benchmark_dir: Path + # Per-invocation directory under `benchmarks/artifacts/runs//`. + # Everything produced by this run lives here so CI artifact upload and + # local cleanup can treat it as a single unit. + run_dir: Path + # Stdout/stderr for every supervised child (mocks, nmp-services, shim, + # aiperf orchestrator). Lives under `run_dir`. + log_dir: Path + # Scratch dir for files the harness materializes at runtime: the + # rewritten AIPerf config, NMP API request bodies we logged for debugging, + # etc. + generated_dir: Path + # Where AIPerf writes its sweep outputs (CSVs, JSONL records, metadata). + # The harness rewrites the AIPerf config's `output_base_dir` to point here. + aiperf_output_dir: Path + # NMP's persistent data directory (SQLite, Secrets vault, etc.). Lives + # *outside* `run_dir` and is reused across runs so subsequent benchmarks + # skip workspace / provider / VM creation work; see the README's Cleanup + # section. Passed to `nemo services run` via `NMP_DATA_DIR`. + nmp_data_dir: Path + # Checked-in YAML template for the AIPerf sweep config. + config_template: Path + # Per-run materialized copy of `config_template` with `output_base_dir` + # overridden to `aiperf_output_dir`. AIPerf is invoked against this file. + runtime_config: Path + # Isolated venv for the upstream `aiperf` CLI. Lives outside `run_dir` + # (sibling to `nmp_data_dir`) so subsequent local runs reuse the install; + # see `bootstrap.py` for why we don't share the workspace venv. + aiperf_venv_dir: Path + + @property + def run_id(self) -> str: + return self.run_dir.name + + def ensure_directories(self) -> None: + for path in ( + self.log_dir, + self.generated_dir, + self.aiperf_output_dir, + self.nmp_data_dir, + ): + path.mkdir(parents=True, exist_ok=True) + + +def _now_run_id() -> str: + return dt.datetime.now().strftime("%Y%m%d_%H%M%S") + + +def discover_nmp_repo_root(start: Path | None = None) -> Path: + """Walk up from ``start`` until a directory containing pyproject.toml + plugins/ is found.""" + candidate = (start or Path(__file__)).resolve() + for parent in (candidate, *candidate.parents): + if (parent / "pyproject.toml").is_file() and (parent / "plugins").is_dir(): + return parent + raise RuntimeError(f"Could not locate NMP repo root from {start or Path(__file__)}") + + +def default_nemoguardrails_repo_root(nmp_repo_root: Path) -> Path: + """Default to a sibling ``NeMo-Guardrails`` checkout next to ``nmp_repo_root``.""" + return (nmp_repo_root.parent / "NeMo-Guardrails").resolve() + + +def build_run_paths( + *, + nmp_repo_root: Path, + nemoguardrails_repo_root: Path, + run_id: str | None = None, +) -> RunPaths: + """Compose the benchmark filesystem layout under the plugin's artifacts dir. + + ``run_id`` overrides the per-run directory name (default: current timestamp). + """ + benchmark_dir = nmp_repo_root / "plugins" / "nemo-guardrails" / "benchmarks" + artifacts_dir = benchmark_dir / "artifacts" + run_dir = artifacts_dir / "runs" / (run_id or _now_run_id()) + + return RunPaths( + nmp_repo_root=nmp_repo_root, + nemoguardrails_repo_root=nemoguardrails_repo_root, + benchmark_dir=benchmark_dir, + run_dir=run_dir, + log_dir=run_dir / "logs", + generated_dir=run_dir / "generated", + aiperf_output_dir=run_dir / "aiperf_results", + nmp_data_dir=artifacts_dir / "nmp-data", + config_template=benchmark_dir / "configs" / "nmp_igw_guardrails_sweep_concurrency.yaml", + runtime_config=run_dir / "generated" / "nmp_igw_guardrails_sweep_concurrency.yaml", + aiperf_venv_dir=artifacts_dir / "venvs" / "aiperf", + ) diff --git a/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/processes.py b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/processes.py new file mode 100644 index 0000000000..7702b4e081 --- /dev/null +++ b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/processes.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Process supervision for the benchmark harness. + +The harness fans out several long-lived child processes per run (two mock LLM +``uvicorn`` servers, ``nemo services run``, and the AIPerf shim) and must clean +them all up — including any workers they fork — even when the parent dies +mid-run. This module wraps ``subprocess.Popen`` in a context-managed +``SupervisedProcess`` that: + +* runs each child in a new session (``start_new_session=True``) so we can send + signals to the whole process group via ``os.killpg`` and reap workers that + ``uvicorn --workers N`` and ``nemo services`` fork off, +* merges stdout/stderr into a per-process log file under ``run_dir/logs/``, +* escalates SIGTERM → SIGKILL on shutdown if the child doesn't exit in time. + +``supervised_processes`` starts the children in order, polls each spec's +``health_url`` before moving on, and guarantees every already-started child is +stopped if a later one fails to come up. ``wait_http`` implements the readiness +probe used by ``supervised_processes`` and for externally managed dependencies +(ex. ``--reuse-services``). +""" + +from __future__ import annotations + +import logging +import os +import signal +import subprocess +import time +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import IO, Iterator + +import httpx +from nemo_guardrails_plugin.benchmarks.bootstrap import build_env + +log = logging.getLogger(__name__) + +# How long to wait for a child (and its process group) to exit after SIGTERM +# before we escalate to SIGKILL. +_TERMINATE_TIMEOUT_SECONDS = 20 + + +@dataclass +class SupervisedProcess: + """A long-lived child managed as a context manager. + + Stdout and stderr are merged into a single log file. The child is placed in a + new session via ``start_new_session=True`` so ``os.killpg`` reaps any workers + it forks. + """ + + # Human-readable identifier used in log messages (e.g. ``"mock-llama"``, + # ``"nmp-services"``). Also drives the log filename via ``log_path``. + name: str + # The argv list passed to ``subprocess.Popen``. Constructed by the caller; + # never assembled from user input here (hence the ``noqa: S603`` in ``start``). + cmd: list[str] + # File to merge stdout + stderr into. Parent dirs are created on ``start``. + # The harness uploads these as CI artifacts on failure. + log_path: Path + # Working directory for the child process. Set per-process so e.g. the mock + # LLM servers run from inside the NeMo-Guardrails checkout where their + # configs live. + cwd: Path + # Extra env vars to overlay on top of the parent's ``os.environ`` (e.g. + # ``PYTHONPATH``, ``NMP_DATA_DIR``). ``None`` means "inherit unchanged". + env: dict[str, str] | None = None + # Readiness probe polled after ``start()`` when this process is entered via + # ``supervised_processes``. ``None`` skips the probe (e.g. when reusing an + # externally managed dependency). + health_url: str | None = None + health_timeout_seconds: float = 60.0 + # The live ``Popen`` handle, populated by ``start()`` and consulted by + # ``stop()``. Excluded from ``__init__`` and ``repr`` since it's pure state. + _proc: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False) + # Open file handle for ``log_path``, held so ``stop()`` can close it in a + # ``finally`` block regardless of how termination unwinds. + _log_fh: IO[bytes] | None = field(default=None, init=False, repr=False) + + def start(self) -> None: + """Spawn the child, redirecting stdout+stderr to ``log_path``. + + Raises ``RuntimeError`` if called twice on the same instance. A + ``SupervisedProcess`` wraps exactly one child process over its lifetime; + to restart, build a new ``SupervisedProcess`` rather than calling + ``start()`` again on a stopped one. + """ + if self._proc is not None: + raise RuntimeError(f"Process {self.name!r} already started") + + # Ensure the log file's parent directory exists. + self.log_path.parent.mkdir(parents=True, exist_ok=True) + # Open the log file for writing. + self._log_fh = self.log_path.open("wb") + + log.info("Starting %s; log=%s", self.name, self.log_path) + try: + self._proc = subprocess.Popen( + self.cmd, + cwd=str(self.cwd), + env=build_env(extra_env=self.env), + stdout=self._log_fh, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + except BaseException: + # Popen failed; close the log handle so we don't leak it. + self._log_fh.close() + self._log_fh = None + raise + + def stop(self) -> None: + """Terminate the child's process group and close the log file. + + Sends SIGTERM to the whole group, waits up to ``_TERMINATE_TIMEOUT_SECONDS``, + then escalates to SIGKILL. Safe to call when the child was never started, + has already exited on its own, or races with us between signals — every + such case becomes a no-op that still closes the log handle. + """ + proc = self._proc + if proc is None: + return + + try: + # Child already exited on its own (crashed, finished early, etc.) — + # nothing to signal, just fall through to the log-close in `finally`. + if proc.poll() is not None: + return + + # Look up the process group id so we can signal the child *and* + # every worker it forked (uvicorn --workers, nemo services). + try: + pgid = os.getpgid(proc.pid) + log.info("Stopping %s pid=%d (pgid=%d)", self.name, proc.pid, pgid) + os.killpg(pgid, signal.SIGTERM) + except ProcessLookupError: + return + + # Give the group a chance to shut down gracefully. If SIGTERM + # doesn't take effect in time, escalate to SIGKILL on the whole + # group and then wait unconditionally. + try: + proc.wait(timeout=_TERMINATE_TIMEOUT_SECONDS) + except subprocess.TimeoutExpired: + log.warning( + "%s did not exit after %ds; sending SIGKILL", + self.name, + _TERMINATE_TIMEOUT_SECONDS, + ) + try: + os.killpg(pgid, signal.SIGKILL) + except ProcessLookupError: + pass + proc.wait() + finally: + # Always close the log handle, even if we returned early above or + # the wait/signal calls raised something unexpected. Idempotent: + # subsequent stop() calls (ex. via __exit__ after manual stop) are + # no-ops once the handle is None. + if self._log_fh is not None: + self._log_fh.close() + self._log_fh = None + + def __enter__(self) -> "SupervisedProcess": + """Context-manager entry: start the child and return ``self``.""" + self.start() + return self + + def __exit__(self, *exc: object) -> None: + """Context-manager exit: stop the child regardless of exception state.""" + self.stop() + + +@contextmanager +def supervised_processes(specs: list[SupervisedProcess]) -> Iterator[list[SupervisedProcess]]: + """Start every spec in order; stop them in reverse on exit. + + Backed by ``ExitStack`` so that if any spec's ``start()`` raises, every + already-started child gets ``stop()``-ed before the exception propagates. + On clean exit, children are torn down in LIFO order (i.e. NMP services + stop before the mock LLMs they depend on). + """ + with ExitStack() as stack: + for spec in specs: + stack.enter_context(spec) + if spec.health_url is not None: + wait_http( + spec.health_url, + timeout_seconds=spec.health_timeout_seconds, + label=spec.name, + ) + yield specs + + +def wait_http(url: str, *, timeout_seconds: float, label: str, poll_interval: float = 1.0) -> None: + """Poll ``url`` until it returns < 400 or ``timeout_seconds`` elapses. + + Used as a readiness gate between starting one supervised process and the + next. ``label`` is included in the ``TimeoutError`` message so it's clear + which dependency failed to come up. Raises ``TimeoutError`` with the last + HTTP error or transport exception attached. + """ + deadline = time.monotonic() + timeout_seconds + last_error: Exception | None = None + + while time.monotonic() < deadline: + try: + r = httpx.get(url, timeout=5.0) + if r.status_code < 400: + return + last_error = RuntimeError(f"HTTP {r.status_code}: {r.text[:200]}") + except httpx.HTTPError as exc: + last_error = exc + time.sleep(poll_interval) + + raise TimeoutError(f"Timed out waiting for {label} at {url}: {last_error}") diff --git a/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/run.py b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/run.py new file mode 100644 index 0000000000..47114e0195 --- /dev/null +++ b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/run.py @@ -0,0 +1,361 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Top-level entry point for the nemo-guardrails IGW benchmark harness. + +The harness orchestrates the entire benchmark run by running the following steps: + +1. Resolve paths and validate the upstream NeMo Guardrails checkout. +2. Write a per-run AIPerf config under ``runs//generated/``. +3. Start the two mock LLM servers from ``${NEMO_GUARDRAILS_REPO_ROOT}/benchmark``. +4. Start (or reuse) ``nemo services run``. +5. Wait for per-process health probes, seed NMP resources via the SDK, smoke-test the VirtualModel. +6. Invoke ``python -m benchmark.aiperf run --config-file ...`` for the sweep. +7. Collect per-sweep results and exit non-zero on any failure. + +Process supervision uses session-scoped subprocesses and an ``ExitStack`` so a +``SIGTERM`` from CI cleans up forked workers (e.g. ``uvicorn --workers 4``). +""" + +from __future__ import annotations + +import argparse +import logging +import os +import sys +import time +from contextlib import ExitStack +from pathlib import Path + +from nemo_guardrails_plugin.benchmarks.aiperf_runner import ( + collect_sweep_results, + prepare_runtime_aiperf_config, + run_aiperf_sweep, +) +from nemo_guardrails_plugin.benchmarks.bootstrap import ensure_aiperf_venv +from nemo_guardrails_plugin.benchmarks.constants import ( + AIPERF_SHIM_BASE_URL, + APP_PROVIDER_URL, + CS_PROVIDER_URL, + IGW_CHAT_PATH, + NMP_BASE_URL, + NMP_HEALTH_PATH, + WORKSPACE, +) +from nemo_guardrails_plugin.benchmarks.paths import ( + RunPaths, + build_run_paths, + default_nemoguardrails_repo_root, + discover_nmp_repo_root, +) +from nemo_guardrails_plugin.benchmarks.processes import ( + SupervisedProcess, + supervised_processes, + wait_http, +) +from nemo_guardrails_plugin.benchmarks.seeding import SeededResources, seed_benchmark +from nemo_platform import APIStatusError, NeMoPlatform + +log = logging.getLogger("nemo_guardrails_plugin.benchmarks") + +_MOCK_HEALTH_TIMEOUT_SECONDS = 60.0 +_NMP_HEALTH_TIMEOUT_SECONDS = 180.0 + + +_REQUIRED_NEMOGUARDRAILS_FILES = ( + Path("benchmark/aiperf/__main__.py"), + Path("benchmark/aiperf/run_aiperf.py"), + Path("benchmark/mock_llm_server/run_server.py"), + Path("benchmark/mock_llm_server/configs/meta-llama-3.3-70b-instruct.env"), + Path("benchmark/mock_llm_server/configs/nvidia-llama-3.1-nemoguard-8b-content-safety.env"), + Path("examples/configs/content_safety_local/config.yml"), + Path("examples/configs/content_safety_local/prompts.yml"), +) + + +def _configure_logging(verbose: bool) -> None: + logging.basicConfig( + level=logging.DEBUG if verbose else logging.INFO, + format="%(asctime)s [%(name)s] %(levelname)s %(message)s", + datefmt="%Y-%m-%dT%H:%M:%S", + ) + + +def _validate_nemoguardrails_repo(nemoguardrails_repo_root: Path) -> None: + """Fail fast if the upstream checkout is missing files the harness depends on.""" + missing = [p for p in _REQUIRED_NEMOGUARDRAILS_FILES if not (nemoguardrails_repo_root / p).is_file()] + if missing: + bullet = "\n - ".join(str(p) for p in missing) + raise FileNotFoundError( + f"NeMo Guardrails checkout at {nemoguardrails_repo_root} is missing required files:\n - {bullet}" + ) + + +def _build_mock_nim_processes(paths: RunPaths, workers: int) -> list[SupervisedProcess]: + """Spawn ``python -m benchmark.mock_llm_server.run_server`` for both mocks. + + Each child is given its own log file and a ``PYTHONPATH`` pointing at the + upstream checkout so its imports resolve. + """ + env = {"PYTHONPATH": str(paths.nemoguardrails_repo_root)} + workdir = paths.nemoguardrails_repo_root / "benchmark" + + # Helper to build a ``SupervisedProcess`` for one of the mock LLM servers. + def spec(name: str, port: int, env_file: Path, *, health_url: str) -> SupervisedProcess: + return SupervisedProcess( + name=name, + cmd=[ + sys.executable, + "-m", + "benchmark.mock_llm_server.run_server", + "--workers", + str(workers), + "--port", + str(port), + "--config-file", + str(env_file), + ], + log_path=paths.log_dir / f"{name}.log", + cwd=workdir, + env=env, + health_url=health_url, + health_timeout_seconds=_MOCK_HEALTH_TIMEOUT_SECONDS, + ) + + return [ + # Main LLM mock server + spec( + "mock-app-llm", + 8000, + paths.nemoguardrails_repo_root / "benchmark/mock_llm_server/configs/meta-llama-3.3-70b-instruct.env", + health_url=f"{APP_PROVIDER_URL}/health", + ), + # Content-safety LLM mock server + spec( + "mock-content-safety-llm", + 8001, + paths.nemoguardrails_repo_root + / "benchmark/mock_llm_server/configs/nvidia-llama-3.1-nemoguard-8b-content-safety.env", + health_url=f"{CS_PROVIDER_URL}/health", + ), + ] + + +def _build_nmp_process(paths: RunPaths) -> SupervisedProcess: + """Start ``nemo services run`` in a supervised process. + + The harness sets ``NMP_BASE_URL`` and ``NMP_DATA_DIR`` env vars so the + child process can talk to NMP over HTTP and write state to the per-run data dir. + """ + return SupervisedProcess( + name="nmp-services", + cmd=["nemo", "services", "run"], + log_path=paths.log_dir / "nmp-services.log", + cwd=paths.nmp_repo_root, + env={"NMP_BASE_URL": NMP_BASE_URL, "NMP_DATA_DIR": str(paths.nmp_data_dir)}, + health_url=f"{NMP_BASE_URL}{NMP_HEALTH_PATH}", + health_timeout_seconds=_NMP_HEALTH_TIMEOUT_SECONDS, + ) + + +def _build_aiperf_shim_process(paths: RunPaths) -> SupervisedProcess: + """Run the shim that satisfies AIPerf's `/v1/models` pre-check. + + Without this, AIPerf's hard-coded health check against the `/v1/models` + endpoint would 404, and the sweep would never start. See + `nemo_guardrails_plugin.benchmarks.shim` for details. + """ + return SupervisedProcess( + name="aiperf-shim", + cmd=[sys.executable, "-m", "nemo_guardrails_plugin.benchmarks.shim"], + log_path=paths.log_dir / "aiperf-shim.log", + cwd=paths.nmp_repo_root, + health_url=f"{AIPERF_SHIM_BASE_URL}/__shim/health", + health_timeout_seconds=_MOCK_HEALTH_TIMEOUT_SECONDS, + ) + + +def _smoke_test(client: NeMoPlatform, seeded: SeededResources) -> None: + """Verify the VirtualModel is reachable and returns a chat completion, + before running the AIPerf sweep. + """ + payload = { + "model": seeded.vm_ref, + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 16, + } + + last_error: str = "no attempts made" + + for attempt in range(60): + try: + body = client.inference.gateway.openai.post( + "v1/chat/completions", + workspace=WORKSPACE, + body=payload, + ) + if body.get("choices"): + return + last_error = f"response missing choices: {body}" + except APIStatusError as exc: + last_error = f"HTTP {exc.status_code}: {str(exc)[:500]}" + log.info("Smoke test attempt %d: %s; retrying", attempt + 1, last_error) + time.sleep(1.0) + + raise RuntimeError(f"Smoke test failed after 60 attempts: {last_error}") + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + prog="nemo-guardrails-benchmark", + description="Run the nemo-guardrails IGW benchmark sweep.", + ) + parser.add_argument( + "--nemo-guardrails-repo-root", + type=Path, + default=Path( + os.environ.get( + "NEMO_GUARDRAILS_REPO_ROOT", + str(default_nemoguardrails_repo_root(discover_nmp_repo_root())), + ) + ), + help="Path to a local NeMo Guardrails checkout (default: $NEMO_GUARDRAILS_REPO_ROOT or ../NeMo-Guardrails).", + ) + parser.add_argument( + "--reuse-services", + action="store_true", + default=os.environ.get("NMP_BENCHMARK_REUSE_SERVICES", "0") == "1", + help="Skip starting `nemo services run` and reuse an existing local NMP at :8080.", + ) + parser.add_argument( + "--keep-running", + action="store_true", + default=os.environ.get("NMP_BENCHMARK_KEEP_RUNNING", "0") == "1", + help="Leave started processes alive after the sweep (debugging).", + ) + parser.add_argument( + "--mock-workers", + type=int, + default=int(os.environ.get("NMP_BENCHMARK_MOCK_WORKERS", "4")), + help="uvicorn worker count for each mock LLM server.", + ) + parser.add_argument( + "--run-id", + default=None, + help="Override the per-run directory name (default: current timestamp).", + ) + parser.add_argument("--verbose", "-v", action="store_true") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + _configure_logging(args.verbose) + + # Validate the NeMo Guardrails checkout before attempting to run the benchmark. + nemoguardrails_repo_root = args.nemo_guardrails_repo_root.resolve() + _validate_nemoguardrails_repo(nemoguardrails_repo_root) + + log.info("Validated NeMo Guardrails local checkout at: %s", nemoguardrails_repo_root) + + # Build the directory structure that will contain the benchmark results. + nmp_repo_root = discover_nmp_repo_root() + paths = build_run_paths( + nmp_repo_root=nmp_repo_root, + nemoguardrails_repo_root=nemoguardrails_repo_root, + run_id=args.run_id, + ) + paths.ensure_directories() + + log.info("Created directory for benchmark results at: %s", paths.run_dir) + + sweep_config = prepare_runtime_aiperf_config( + template_path=paths.config_template, + runtime_config_path=paths.runtime_config, + aiperf_output_dir=paths.aiperf_output_dir, + ) + log.info( + "AIPerf sweep: concurrency=%s, duration=%ss", + sweep_config.get("sweeps", {}).get("concurrency"), + sweep_config.get("base_config", {}).get("benchmark_duration"), + ) + + # Ensure the dedicated aiperf venv exists *before* we start any supervised + # processes. + aiperf_python = ensure_aiperf_venv(paths.aiperf_venv_dir) + log.info("Using aiperf python at %s", aiperf_python) + + processes = _build_mock_nim_processes(paths, args.mock_workers) + if not args.reuse_services: + processes.append(_build_nmp_process(paths)) + + processes.append(_build_aiperf_shim_process(paths)) + + # Start the processes and wait for them to be ready before seeding NMP. + with ExitStack() as stack: + stack.enter_context(supervised_processes(processes)) + if args.keep_running: + # Pop the cleanup so processes outlive this script. + stack.pop_all() + + if args.reuse_services: + log.info("Waiting for existing NMP services at %s...", NMP_BASE_URL) + wait_http( + f"{NMP_BASE_URL}{NMP_HEALTH_PATH}", + timeout_seconds=_NMP_HEALTH_TIMEOUT_SECONDS, + label="nmp-services", + ) + + log.info(f"All services are ready. Seeding benchmark resources in workspace {WORKSPACE}...") + + client = NeMoPlatform(base_url=NMP_BASE_URL) + seeded = seed_benchmark( + client, + nemoguardrails_repo_root=paths.nemoguardrails_repo_root, + generated_dir=paths.generated_dir, + ) + + log.info("Waiting for VirtualModel %s to be ready...", seeded.vm_ref) + _smoke_test(client, seeded) + + log.info( + "Starting AIPerf sweep against %s -> shim -> %s%s", + AIPERF_SHIM_BASE_URL, + NMP_BASE_URL, + IGW_CHAT_PATH, + ) + aiperf_exit = run_aiperf_sweep( + nemoguardrails_repo_root=paths.nemoguardrails_repo_root, + runtime_config=paths.runtime_config, + log_path=paths.log_dir / "aiperf.log", + python_executable=str(aiperf_python), + venv_bin_path=paths.aiperf_venv_dir / "bin", + ) + + sweep_results = collect_sweep_results(paths.aiperf_output_dir) + failures = sum(1 for r in sweep_results if not r.passed) + + if not sweep_results: + # AIPerf exited before producing any per-sweep dirs. Surface that + # explicitly so the log isn't ambiguous about why we're failing. + log.error( + "aiperf exited with code %d and produced no per-sweep results in %s", + aiperf_exit, + paths.aiperf_output_dir, + ) + else: + log.info( + "Sweep summary: %d run(s), %d failure(s); per-sweep outputs under %s", + len(sweep_results), + failures, + paths.aiperf_output_dir, + ) + + if failures or aiperf_exit != 0 or not sweep_results: + return 1 + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/seeding.py b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/seeding.py new file mode 100644 index 0000000000..96cd152f0f --- /dev/null +++ b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/seeding.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Seed NMP with the workspace, providers, guardrail config, and VirtualModel +required by the IGW guardrails benchmark. + +Replaces the previous ``setup_nmp_guardrails_benchmark.sh`` flow with direct +NMP SDK calls (``client.workspaces``, ``client.inference``, ``client.guardrail``). +""" + +from __future__ import annotations + +import json +import logging +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import yaml +from nemo_guardrails_plugin.benchmarks.constants import ( + APP_MODEL_NAME, + APP_PROVIDER, + APP_PROVIDER_URL, + CS_MODEL_NAME, + CS_PROVIDER, + CS_PROVIDER_URL, + GUARDRAIL_CONFIG, + GUARDRAILS_MIDDLEWARE_CONFIG_TYPE, + GUARDRAILS_MIDDLEWARE_NAME, + VM_NAME, + WORKSPACE, +) +from nemo_platform import NeMoPlatform, NotFoundError +from nemo_platform.types.inference.middleware_call_param import MiddlewareCallParam +from nemo_platform.types.inference.virtual_model_inference_config_param import ( + VirtualModelInferenceConfigParam, +) + +log = logging.getLogger(__name__) + +_PROVIDER_WAIT_TIMEOUT_SECONDS = 60 +_PROVIDER_POLL_INTERVAL_SECONDS = 1.0 + + +@dataclass(frozen=True) +class SeededResources: + workspace: str + app_provider_name: str + cs_provider_name: str + app_model_entity: str + cs_model_entity: str + guardrail_config_name: str + vm_name: str + + @property + def vm_ref(self) -> str: + return f"{self.workspace}/{self.vm_name}" + + @property + def guardrail_config_ref(self) -> str: + return f"{self.workspace}/{self.guardrail_config_name}" + + +def seed_benchmark( + client: NeMoPlatform, + *, + nemoguardrails_repo_root: Path, + generated_dir: Path, + provider_wait_timeout: float = _PROVIDER_WAIT_TIMEOUT_SECONDS, +) -> SeededResources: + """Create workspace, providers, GuardrailConfig, and VirtualModel. + + All ``create`` calls are idempotent (``exist_ok=True``) so this is safe to + rerun against a reused NMP instance. + """ + generated_dir.mkdir(parents=True, exist_ok=True) + + log.info("Creating workspace %s", WORKSPACE) + client.workspaces.create( + name=WORKSPACE, + description="Local IGW guardrails benchmark workspace", + exist_ok=True, + ) + + log.info("Registering app mock provider %s", APP_PROVIDER) + client.inference.providers.create( + workspace=WORKSPACE, + name=APP_PROVIDER, + host_url=APP_PROVIDER_URL, + enabled_models=[APP_MODEL_NAME], + description=f"Benchmark mock app LLM on {APP_PROVIDER_URL}", + exist_ok=True, + ) + + log.info("Registering content-safety mock provider %s", CS_PROVIDER) + client.inference.providers.create( + workspace=WORKSPACE, + name=CS_PROVIDER, + host_url=CS_PROVIDER_URL, + enabled_models=[CS_MODEL_NAME], + description=f"Benchmark mock content-safety LLM on {CS_PROVIDER_URL}", + exist_ok=True, + ) + + log.info("Waiting for provider discovery") + app_provider = _wait_for_served_model( + client, + provider_name=APP_PROVIDER, + served_model_name=APP_MODEL_NAME, + timeout_seconds=provider_wait_timeout, + ) + cs_provider = _wait_for_served_model( + client, + provider_name=CS_PROVIDER, + served_model_name=CS_MODEL_NAME, + timeout_seconds=provider_wait_timeout, + ) + + app_entity = _extract_model_entity(app_provider, APP_MODEL_NAME, provider_name=APP_PROVIDER) + cs_entity = _extract_model_entity(cs_provider, CS_MODEL_NAME, provider_name=CS_PROVIDER) + + _dump_model(generated_dir / "app_provider.json", app_provider) + _dump_model(generated_dir / "content_safety_provider.json", cs_provider) + + log.info("Building GuardrailConfig payload from %s", nemoguardrails_repo_root) + config_data = build_guardrail_config_data( + source_config_dir=nemoguardrails_repo_root / "examples" / "configs" / "content_safety_local", + content_safety_model_entity=cs_entity, + ) + # Persist the same payload shape the old shell harness produced for debuggability. + (generated_dir / "content_safety_local_nmp_request.json").write_text( + json.dumps( + { + "name": GUARDRAIL_CONFIG, + "description": "Benchmark content_safety_local config routed through IGW", + "data": config_data, + "exist_ok": True, + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + + log.info("Creating GuardrailConfig %s", GUARDRAIL_CONFIG) + client.guardrail.configs.create( + workspace=WORKSPACE, + name=GUARDRAIL_CONFIG, + description="Benchmark content_safety_local config routed through IGW", + data=config_data, + exist_ok=True, + ) + + middleware_call: MiddlewareCallParam = { + "name": GUARDRAILS_MIDDLEWARE_NAME, + "config_type": GUARDRAILS_MIDDLEWARE_CONFIG_TYPE, + "config_id": f"{WORKSPACE}/{GUARDRAIL_CONFIG}", + } + vm_models: list[VirtualModelInferenceConfigParam] = [{"model": app_entity, "backend_format": "OPENAI_CHAT"}] + + log.info("Creating VirtualModel %s/%s", WORKSPACE, VM_NAME) + vm = client.inference.virtual_models.create( + workspace=WORKSPACE, + name=VM_NAME, + default_model_entity=app_entity, + models=vm_models, + request_middleware=[middleware_call], + response_middleware=[middleware_call], + exist_ok=True, + ) + _dump_model(generated_dir / "virtual_model.json", vm) + + return SeededResources( + workspace=WORKSPACE, + app_provider_name=APP_PROVIDER, + cs_provider_name=CS_PROVIDER, + app_model_entity=app_entity, + cs_model_entity=cs_entity, + guardrail_config_name=GUARDRAIL_CONFIG, + vm_name=VM_NAME, + ) + + +def _wait_for_served_model( + client: NeMoPlatform, + *, + provider_name: str, + served_model_name: str, + timeout_seconds: float, +) -> Any: + """Poll a provider until ``served_models`` lists the expected entry. + + Gateway readiness alone is not enough: VirtualModel creation needs the + discovered ``model_entity_id``, which only appears once the provider has + enumerated its models. + """ + deadline = time.monotonic() + timeout_seconds + last_error: Exception | None = None + while time.monotonic() < deadline: + try: + provider = client.inference.providers.retrieve(provider_name, workspace=WORKSPACE) + except NotFoundError as exc: + last_error = exc + time.sleep(_PROVIDER_POLL_INTERVAL_SECONDS) + continue + served_models = getattr(provider, "served_models", None) or [] + for m in served_models: + if getattr(m, "served_model_name", None) == served_model_name and getattr(m, "model_entity_id", None): + return provider + time.sleep(_PROVIDER_POLL_INTERVAL_SECONDS) + + raise TimeoutError( + f"Provider {provider_name!r} did not surface served model " + f"{served_model_name!r} within {timeout_seconds}s: {last_error}" + ) + + +def _extract_model_entity(provider: Any, served_model_name: str, *, provider_name: str) -> str: + for m in getattr(provider, "served_models", None) or []: + if getattr(m, "served_model_name", None) == served_model_name: + entity = getattr(m, "model_entity_id", None) + if entity: + return entity + raise RuntimeError(f"Provider {provider_name!r} does not expose served model {served_model_name!r}") + + +def build_guardrail_config_data( + *, + source_config_dir: Path, + content_safety_model_entity: str, +) -> dict[str, Any]: + """Read the upstream content_safety_local config and rewrite it for NMP. + + The upstream config references an HTTP base_url; in NMP we instead route by + ``model_entity_id`` resolved via the inference gateway. Prompts are inlined + from the sibling ``prompts.yml`` file. + """ + config_yaml = source_config_dir / "config.yml" + prompts_yaml = source_config_dir / "prompts.yml" + if not config_yaml.is_file(): + raise FileNotFoundError(f"Expected guardrails config at {config_yaml}") + if not prompts_yaml.is_file(): + raise FileNotFoundError(f"Expected guardrails prompts at {prompts_yaml}") + + config = yaml.safe_load(config_yaml.read_text(encoding="utf-8")) + if not isinstance(config, dict): + raise ValueError(f"Expected a YAML mapping at {config_yaml}, got {type(config).__name__}") + + prompts = yaml.safe_load(prompts_yaml.read_text(encoding="utf-8")) or {} + if not isinstance(prompts, dict): + raise ValueError(f"Expected a YAML mapping at {prompts_yaml}, got {type(prompts).__name__}") + + config["models"] = [ + { + "type": "content_safety", + "engine": "nim", + "model": content_safety_model_entity, + } + ] + config["prompts"] = prompts.get("prompts", []) + return config + + +def _dump_model(path: Path, model: Any) -> None: + """Best-effort serialize an SDK response model to JSON.""" + if hasattr(model, "model_dump"): + payload = model.model_dump(mode="json") + elif hasattr(model, "to_dict"): + payload = model.to_dict() + else: + payload = json.loads(json.dumps(model, default=str)) + path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") diff --git a/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/shim.py b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/shim.py new file mode 100644 index 0000000000..d304fcd8bc --- /dev/null +++ b/plugins/nemo-guardrails/src/nemo_guardrails_plugin/benchmarks/shim.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""HTTP shim that satisfies AIPerf's pre-check and proxies chat completions. + +AIPerf's ``_check_service`` issues ``GET urljoin(base_url, "/v1/models")`` and +expects a 200, but NMP doesn't serve ``/v1/models`` at the IGW workspace root +the way OpenAI does, and upstream AIPerf doesn't expose a knob to override the +probe path. To unblock the benchmark without patching AIPerf or NMP, we run +this tiny shim on a separate port: + +- ``GET /v1/models`` -> ``200 {"object":"list","data":[]}`` +- ``POST /v1/chat/completions`` -> reverse-proxy to NMP IGW +- ``GET /__shim/health`` -> ``200 {"status":"ok"}`` +- Any other path -> 404 + +This module is meant to be invoked as ``python -m +nemo_guardrails_plugin.benchmarks.shim`` so it can be supervised by the same +process-group machinery the harness uses for the mock LLMs and ``nemo +services run``. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import httpx +from nemo_guardrails_plugin.benchmarks.constants import ( + AIPERF_SHIM_HOST, + AIPERF_SHIM_PORT, + IGW_CHAT_PATH, + NMP_BASE_URL, +) + +log = logging.getLogger("nemo_guardrails_plugin.benchmarks.shim") + + +_MODELS_RESPONSE = json.dumps({"object": "list", "data": []}).encode("utf-8") +_HEALTH_RESPONSE = json.dumps({"status": "ok"}).encode("utf-8") + + +class _ShimHandler(BaseHTTPRequestHandler): + """Minimal handler that routes the two paths AIPerf actually touches.""" + + # Allow override via class attr so tests can swap in a mock httpx client. + upstream_url: str = f"{NMP_BASE_URL}{IGW_CHAT_PATH}" + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 + log.debug("shim: " + format, *args) + + def _send_json(self, status: int, body: bytes) -> None: + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self) -> None: # noqa: N802 - http.server API + if self.path == "/v1/models": + self._send_json(200, _MODELS_RESPONSE) + return + if self.path == "/__shim/health": + self._send_json(200, _HEALTH_RESPONSE) + return + self._send_json(404, json.dumps({"detail": "not found"}).encode("utf-8")) + + def do_POST(self) -> None: # noqa: N802 - http.server API + if self.path != "/v1/chat/completions": + self._send_json(404, json.dumps({"detail": "not found"}).encode("utf-8")) + return + + length = int(self.headers.get("Content-Length", "0") or "0") + body = self.rfile.read(length) if length else b"" + + forwarded_headers = { + k: v for k, v in self.headers.items() if k.lower() not in {"host", "content-length", "connection"} + } + + try: + response = httpx.post( + self.upstream_url, + content=body, + headers=forwarded_headers, + timeout=httpx.Timeout(connect=10.0, read=300.0, write=10.0, pool=10.0), + ) + except httpx.HTTPError as e: + log.warning("shim: upstream request failed: %s", e) + self._send_json( + 502, + json.dumps({"detail": f"upstream error: {e}"}).encode("utf-8"), + ) + return + + self.send_response(response.status_code) + for header, value in response.headers.items(): + if header.lower() in {"transfer-encoding", "connection", "content-length"}: + continue + self.send_header(header, value) + self.send_header("Content-Length", str(len(response.content))) + self.end_headers() + self.wfile.write(response.content) + + +def serve(host: str = AIPERF_SHIM_HOST, port: int = AIPERF_SHIM_PORT) -> None: + """Run the shim until the process is signalled.""" + httpd = ThreadingHTTPServer((host, port), _ShimHandler) + log.info("AIPerf shim listening on http://%s:%d", host, port) + try: + httpd.serve_forever() + finally: + httpd.server_close() + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog="nemo-guardrails-benchmark-shim", + description=__doc__, + ) + parser.add_argument("--host", default=AIPERF_SHIM_HOST) + parser.add_argument("--port", type=int, default=AIPERF_SHIM_PORT) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + ) + args = parser.parse_args(argv) + logging.basicConfig( + level=args.log_level, + format="%(asctime)s [%(name)s] %(levelname)s %(message)s", + ) + serve(host=args.host, port=args.port) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/plugins/nemo-guardrails/tests/unit/benchmarks/test_aiperf_runner.py b/plugins/nemo-guardrails/tests/unit/benchmarks/test_aiperf_runner.py new file mode 100644 index 0000000000..1c81252934 --- /dev/null +++ b/plugins/nemo-guardrails/tests/unit/benchmarks/test_aiperf_runner.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from pathlib import Path + +import pytest +import yaml +from nemo_guardrails_plugin.benchmarks.aiperf_runner import ( + SweepRunResult, + collect_sweep_results, + prepare_runtime_aiperf_config, +) + + +def _write_template(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + yaml.safe_dump( + { + "batch_name": "nmp_igw_guardrails_sweep_concurrency", + "output_base_dir": "plugins/nemo-guardrails/benchmarks/artifacts/aiperf_results", + "base_config": {"model": "benchmark/guardrails-vm"}, + "sweeps": {"concurrency": [1, 2, 4]}, + } + ), + encoding="utf-8", + ) + + +class TestPrepareRuntimeAiperfConfig: + def test_overrides_output_base_dir(self, tmp_path: Path) -> None: + template_path = tmp_path / "template.yaml" + _write_template(template_path) + runtime_config_path = tmp_path / "out" / "runtime.yaml" + aiperf_output_dir = tmp_path / "results" + + config = prepare_runtime_aiperf_config( + template_path=template_path, + runtime_config_path=runtime_config_path, + aiperf_output_dir=aiperf_output_dir, + ) + + assert config["output_base_dir"] == str(aiperf_output_dir) + written = yaml.safe_load(runtime_config_path.read_text(encoding="utf-8")) + assert written["output_base_dir"] == str(aiperf_output_dir) + assert written["base_config"]["model"] == "benchmark/guardrails-vm" + assert written["sweeps"]["concurrency"] == [1, 2, 4] + + def test_missing_template_raises(self, tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError): + prepare_runtime_aiperf_config( + template_path=tmp_path / "absent.yaml", + runtime_config_path=tmp_path / "out.yaml", + aiperf_output_dir=tmp_path / "results", + ) + + def test_non_mapping_template_raises(self, tmp_path: Path) -> None: + template_path = tmp_path / "bad.yaml" + template_path.write_text("- just\n- a\n- list\n", encoding="utf-8") + with pytest.raises(ValueError, match="mapping"): + prepare_runtime_aiperf_config( + template_path=template_path, + runtime_config_path=tmp_path / "out.yaml", + aiperf_output_dir=tmp_path / "results", + ) + + +def _make_sweep_dir(parent: Path, sweep_label: str, *, returncode: int | None, duration: float | None) -> Path: + sweep_dir = parent / sweep_label + sweep_dir.mkdir(parents=True) + if returncode is not None: + (sweep_dir / "process_result.json").write_text(json.dumps({"returncode": returncode}), encoding="utf-8") + if duration is not None: + (sweep_dir / "run_metadata.json").write_text(json.dumps({"duration_seconds": duration}), encoding="utf-8") + return sweep_dir + + +class TestCollectSweepResults: + def test_empty_dir_returns_empty(self, tmp_path: Path) -> None: + assert collect_sweep_results(tmp_path / "missing") == [] + + def test_collects_all_sweeps_with_status(self, tmp_path: Path) -> None: + batch = tmp_path / "nmp_igw_guardrails_sweep_concurrency" / "20260527_120000" + _make_sweep_dir(batch, "concurrency1", returncode=0, duration=70.5) + _make_sweep_dir(batch, "concurrency2", returncode=1, duration=70.5) + + results = collect_sweep_results(tmp_path) + + assert len(results) == 2 + results_by_label = {r.sweep_label: r for r in results} + assert results_by_label["concurrency1"].passed + assert results_by_label["concurrency1"].duration_seconds == 70.5 + assert not results_by_label["concurrency2"].passed + assert results_by_label["concurrency2"].return_code == 1 + + def test_missing_process_result_is_failure(self, tmp_path: Path) -> None: + batch = tmp_path / "batch" / "ts" + _make_sweep_dir(batch, "concurrency1", returncode=None, duration=10.0) + + results = collect_sweep_results(tmp_path) + + assert len(results) == 1 + assert not results[0].passed + + def test_malformed_json_treated_as_failure_without_crashing(self, tmp_path: Path) -> None: + batch = tmp_path / "batch" / "ts" + sweep = _make_sweep_dir(batch, "concurrency1", returncode=None, duration=None) + (sweep / "process_result.json").write_text("not-json", encoding="utf-8") + (sweep / "run_metadata.json").write_text("not-json", encoding="utf-8") + + results = collect_sweep_results(tmp_path) + assert results[0].return_code == 1 + assert results[0].duration_seconds == 0.0 + + +class TestSweepRunResult: + def test_passed_property(self) -> None: + passing = SweepRunResult( + sweep_label="x", + output_dir=Path("."), + return_code=0, + duration_seconds=1.0, + metadata_path=None, + process_result_path=None, + ) + assert passing.passed + assert not SweepRunResult( + sweep_label="x", + output_dir=Path("."), + return_code=2, + duration_seconds=1.0, + metadata_path=None, + process_result_path=None, + ).passed diff --git a/plugins/nemo-guardrails/tests/unit/benchmarks/test_paths.py b/plugins/nemo-guardrails/tests/unit/benchmarks/test_paths.py new file mode 100644 index 0000000000..53b2b42e36 --- /dev/null +++ b/plugins/nemo-guardrails/tests/unit/benchmarks/test_paths.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pytest +from nemo_guardrails_plugin.benchmarks.paths import ( + build_run_paths, + default_nemoguardrails_repo_root, + discover_nmp_repo_root, +) + + +def _make_fake_repo(root: Path) -> Path: + root.mkdir(parents=True, exist_ok=True) + (root / "pyproject.toml").write_text("", encoding="utf-8") + (root / "plugins").mkdir() + return root + + +class TestDiscoverNmpRepoRoot: + def test_finds_root_from_nested_path(self, tmp_path: Path) -> None: + repo = _make_fake_repo(tmp_path / "repo") + nested = repo / "plugins" / "foo" / "src" / "deep" + nested.mkdir(parents=True) + + assert discover_nmp_repo_root(nested) == repo + + def test_returns_repo_when_pointed_directly_at_it(self, tmp_path: Path) -> None: + repo = _make_fake_repo(tmp_path / "repo") + assert discover_nmp_repo_root(repo) == repo + + def test_raises_when_no_repo_in_ancestry(self, tmp_path: Path) -> None: + with pytest.raises(RuntimeError, match="repo root"): + discover_nmp_repo_root(tmp_path) + + +class TestDefaultNgRepoRoot: + def test_sibling_of_nmp_root(self, tmp_path: Path) -> None: + nmp = _make_fake_repo(tmp_path / "nemo-platform") + ng = tmp_path / "NeMo-Guardrails" + + assert default_nemoguardrails_repo_root(nmp) == ng.resolve() + + +class TestBuildRunPaths: + def test_layout_matches_documented_structure(self, tmp_path: Path) -> None: + nmp = _make_fake_repo(tmp_path / "nemo-platform") + ng = tmp_path / "NeMo-Guardrails" + ng.mkdir() + + paths = build_run_paths(nmp_repo_root=nmp, nemoguardrails_repo_root=ng, run_id="20260527_120000") + + assert paths.run_dir == nmp / "plugins/nemo-guardrails/benchmarks/artifacts/runs/20260527_120000" + assert paths.log_dir == paths.run_dir / "logs" + assert paths.generated_dir == paths.run_dir / "generated" + assert paths.aiperf_output_dir == paths.run_dir / "aiperf_results" + assert paths.nmp_data_dir == nmp / "plugins/nemo-guardrails/benchmarks/artifacts/nmp-data" + assert ( + paths.config_template + == nmp / "plugins/nemo-guardrails/benchmarks/configs/nmp_igw_guardrails_sweep_concurrency.yaml" + ) + assert paths.runtime_config == paths.generated_dir / "nmp_igw_guardrails_sweep_concurrency.yaml" + + def test_ensure_directories_creates_required_dirs(self, tmp_path: Path) -> None: + nmp = _make_fake_repo(tmp_path / "nemo-platform") + ng = tmp_path / "NeMo-Guardrails" + ng.mkdir() + + paths = build_run_paths(nmp_repo_root=nmp, nemoguardrails_repo_root=ng, run_id="x") + paths.ensure_directories() + + assert paths.log_dir.is_dir() + assert paths.generated_dir.is_dir() + assert paths.aiperf_output_dir.is_dir() + assert paths.nmp_data_dir.is_dir() + + def test_run_id_uses_timestamp_when_not_given(self, tmp_path: Path) -> None: + nmp = _make_fake_repo(tmp_path / "nemo-platform") + ng = tmp_path / "NeMo-Guardrails" + ng.mkdir() + + paths = build_run_paths(nmp_repo_root=nmp, nemoguardrails_repo_root=ng) + # Timestamp format: YYYYmmdd_HHMMSS = 15 chars including underscore. + assert len(paths.run_id) == 15 + assert paths.run_id[8] == "_" + + def test_aiperf_venv_dir_is_outside_run_dir(self, tmp_path: Path) -> None: + """The cached aiperf venv must be shared across runs, not under run_dir.""" + nmp = _make_fake_repo(tmp_path / "nemo-platform") + ng = tmp_path / "NeMo-Guardrails" + ng.mkdir() + + paths = build_run_paths(nmp_repo_root=nmp, nemoguardrails_repo_root=ng, run_id="x") + assert paths.run_dir not in paths.aiperf_venv_dir.parents + assert paths.aiperf_venv_dir.parent.name == "venvs" diff --git a/plugins/nemo-guardrails/tests/unit/benchmarks/test_seeding.py b/plugins/nemo-guardrails/tests/unit/benchmarks/test_seeding.py new file mode 100644 index 0000000000..f0ad57aeaa --- /dev/null +++ b/plugins/nemo-guardrails/tests/unit/benchmarks/test_seeding.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import yaml +from nemo_guardrails_plugin.benchmarks.constants import ( + APP_MODEL_NAME, + APP_PROVIDER, + CS_MODEL_NAME, + CS_PROVIDER, + GUARDRAIL_CONFIG, + VM_NAME, + WORKSPACE, +) +from nemo_guardrails_plugin.benchmarks.seeding import ( + build_guardrail_config_data, + seed_benchmark, +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _make_provider(*, provider_name: str, served_model_name: str, entity_suffix: str = "entity") -> SimpleNamespace: + return SimpleNamespace( + name=provider_name, + served_models=[ + SimpleNamespace( + served_model_name=served_model_name, + model_entity_id=f"{WORKSPACE}/{served_model_name.replace('/', '-')}-{entity_suffix}", + ) + ], + ) + + +def _write_upstream_configs(ng_root: Path) -> Path: + cs_dir = ng_root / "examples" / "configs" / "content_safety_local" + cs_dir.mkdir(parents=True) + (cs_dir / "config.yml").write_text( + yaml.safe_dump( + { + "models": [ + { + "type": "main", + "engine": "nim", + "model": "meta/llama-3.3-70b-instruct", + "parameters": {"base_url": "http://localhost:8000"}, + }, + ], + "rails": {"input": {"flows": ["content safety check input $model=content_safety"]}}, + } + ), + encoding="utf-8", + ) + (cs_dir / "prompts.yml").write_text( + yaml.safe_dump({"prompts": [{"task": "content_safety_check_input", "content": "..."}]}), + encoding="utf-8", + ) + return cs_dir + + +@pytest.fixture +def fake_client() -> MagicMock: + client = MagicMock() + client.inference.providers.create = MagicMock() + client.inference.providers.retrieve = MagicMock( + side_effect=lambda name, workspace=None: _make_provider( + provider_name=name, + served_model_name=APP_MODEL_NAME if name == APP_PROVIDER else CS_MODEL_NAME, + ) + ) + client.inference.virtual_models.create = MagicMock( + return_value=SimpleNamespace(name=VM_NAME, default_model_entity=f"{WORKSPACE}/app") + ) + client.guardrail.configs.create = MagicMock(return_value=SimpleNamespace(name=GUARDRAIL_CONFIG)) + client.workspaces.create = MagicMock(return_value=SimpleNamespace(name=WORKSPACE)) + return client + + +# --------------------------------------------------------------------------- +# build_guardrail_config_data +# --------------------------------------------------------------------------- + + +class TestBuildGuardrailConfigData: + def test_rewrites_models_and_inlines_prompts(self, tmp_path: Path) -> None: + cs_dir = _write_upstream_configs(tmp_path) + + data = build_guardrail_config_data( + source_config_dir=cs_dir, + content_safety_model_entity="benchmark/cs-entity", + ) + + assert data["models"] == [{"type": "content_safety", "engine": "nim", "model": "benchmark/cs-entity"}] + assert data["prompts"] == [{"task": "content_safety_check_input", "content": "..."}] + # Non-models fields preserved. + assert data["rails"]["input"]["flows"] == ["content safety check input $model=content_safety"] + + def test_missing_config_raises(self, tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError, match="config.yml"): + build_guardrail_config_data( + source_config_dir=tmp_path / "nope", + content_safety_model_entity="x/y", + ) + + +# --------------------------------------------------------------------------- +# seed_benchmark +# --------------------------------------------------------------------------- + + +class TestSeedBenchmark: + def test_calls_sdk_with_expected_payloads(self, fake_client: MagicMock, tmp_path: Path) -> None: + ng_root = tmp_path / "NeMo-Guardrails" + _write_upstream_configs(ng_root) + generated_dir = tmp_path / "generated" + + seeded = seed_benchmark( + fake_client, + nemoguardrails_repo_root=ng_root, + generated_dir=generated_dir, + provider_wait_timeout=1.0, + ) + + fake_client.workspaces.create.assert_called_once_with( + name=WORKSPACE, + description="Local IGW guardrails benchmark workspace", + exist_ok=True, + ) + # Both providers registered. + provider_create_names = [c.kwargs["name"] for c in fake_client.inference.providers.create.call_args_list] + assert sorted(provider_create_names) == sorted([APP_PROVIDER, CS_PROVIDER]) + + # Guardrail config payload uses the discovered content-safety entity. + gc_call = fake_client.guardrail.configs.create.call_args + assert gc_call.kwargs["name"] == GUARDRAIL_CONFIG + assert gc_call.kwargs["workspace"] == WORKSPACE + assert gc_call.kwargs["exist_ok"] is True + cs_entity = seeded.cs_model_entity + assert gc_call.kwargs["data"]["models"][0]["model"] == cs_entity + + # VirtualModel uses the discovered app entity and points middleware at the + # guardrail config we just created. + vm_call = fake_client.inference.virtual_models.create.call_args + assert vm_call.kwargs["name"] == VM_NAME + assert vm_call.kwargs["default_model_entity"] == seeded.app_model_entity + assert vm_call.kwargs["models"] == [{"model": seeded.app_model_entity, "backend_format": "OPENAI_CHAT"}] + expected_middleware = [ + { + "name": "nemo-guardrails", + "config_type": "guardrail_config", + "config_id": f"{WORKSPACE}/{GUARDRAIL_CONFIG}", + } + ] + assert vm_call.kwargs["request_middleware"] == expected_middleware + assert vm_call.kwargs["response_middleware"] == expected_middleware + + def test_generated_dir_contains_artifacts(self, fake_client: MagicMock, tmp_path: Path) -> None: + ng_root = tmp_path / "NeMo-Guardrails" + _write_upstream_configs(ng_root) + generated_dir = tmp_path / "generated" + + seed_benchmark( + fake_client, + nemoguardrails_repo_root=ng_root, + generated_dir=generated_dir, + provider_wait_timeout=1.0, + ) + + assert (generated_dir / "app_provider.json").is_file() + assert (generated_dir / "content_safety_provider.json").is_file() + assert (generated_dir / "virtual_model.json").is_file() + + request_payload = json.loads( + (generated_dir / "content_safety_local_nmp_request.json").read_text(encoding="utf-8") + ) + assert request_payload["name"] == GUARDRAIL_CONFIG + assert request_payload["exist_ok"] is True + assert request_payload["data"]["models"][0]["type"] == "content_safety" + + def test_returns_seeded_resources(self, fake_client: MagicMock, tmp_path: Path) -> None: + ng_root = tmp_path / "NeMo-Guardrails" + _write_upstream_configs(ng_root) + + seeded = seed_benchmark( + fake_client, + nemoguardrails_repo_root=ng_root, + generated_dir=tmp_path / "generated", + provider_wait_timeout=1.0, + ) + + assert seeded.workspace == WORKSPACE + assert seeded.vm_ref == f"{WORKSPACE}/{VM_NAME}" + assert seeded.guardrail_config_ref == f"{WORKSPACE}/{GUARDRAIL_CONFIG}" + + def test_raises_if_served_models_never_populated(self, tmp_path: Path) -> None: + ng_root = tmp_path / "NeMo-Guardrails" + _write_upstream_configs(ng_root) + + client = MagicMock() + client.workspaces.create = MagicMock() + client.inference.providers.create = MagicMock() + client.inference.providers.retrieve = MagicMock(return_value=SimpleNamespace(served_models=[])) + + with pytest.raises(TimeoutError, match="served model"): + seed_benchmark( + client, + nemoguardrails_repo_root=ng_root, + generated_dir=tmp_path / "generated", + provider_wait_timeout=0.1, + ) diff --git a/plugins/nemo-guardrails/tests/unit/benchmarks/test_shim.py b/plugins/nemo-guardrails/tests/unit/benchmarks/test_shim.py new file mode 100644 index 0000000000..f100d93b82 --- /dev/null +++ b/plugins/nemo-guardrails/tests/unit/benchmarks/test_shim.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Smoke tests for the AIPerf shim. + +We stand up the shim on a real local port and a stub upstream HTTP server, then +exercise the three routes AIPerf and the harness rely on. +""" + +from __future__ import annotations + +import json +import socket +import threading +from contextlib import closing +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Generator + +import httpx +import pytest +from nemo_guardrails_plugin.benchmarks import shim as shim_module + + +def _free_port() -> int: + with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +class _UpstreamHandler(BaseHTTPRequestHandler): + received_body: bytes = b"" + + def log_message(self, *args: object, **kwargs: object) -> None: # noqa: A002 + return + + def do_POST(self) -> None: # noqa: N802 + length = int(self.headers.get("Content-Length", "0") or "0") + type(self).received_body = self.rfile.read(length) if length else b"" + body = json.dumps({"choices": [{"message": {"content": "hi"}}]}).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + +@pytest.fixture +def upstream_server() -> Generator[str, None, None]: + port = _free_port() + server = ThreadingHTTPServer(("127.0.0.1", port), _UpstreamHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{port}/chat" + finally: + server.shutdown() + server.server_close() + + +@pytest.fixture +def shim_server(upstream_server: str) -> Generator[str, None, None]: + port = _free_port() + # Override the class attr so the shim proxies to our stub upstream. + previous = shim_module._ShimHandler.upstream_url + shim_module._ShimHandler.upstream_url = upstream_server + server = ThreadingHTTPServer(("127.0.0.1", port), shim_module._ShimHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{port}" + finally: + shim_module._ShimHandler.upstream_url = previous + server.shutdown() + server.server_close() + + +def test_models_returns_200_with_openai_shape(shim_server: str) -> None: + response = httpx.get(f"{shim_server}/v1/models", timeout=2.0) + assert response.status_code == 200 + body = response.json() + assert body == {"object": "list", "data": []} + + +def test_health_endpoint_returns_200(shim_server: str) -> None: + response = httpx.get(f"{shim_server}/__shim/health", timeout=2.0) + assert response.status_code == 200 + assert response.json() == {"status": "ok"} + + +def test_chat_completions_proxies_to_upstream(shim_server: str) -> None: + payload = {"model": "x", "messages": [{"role": "user", "content": "hi"}]} + response = httpx.post( + f"{shim_server}/v1/chat/completions", + json=payload, + timeout=5.0, + ) + assert response.status_code == 200 + assert response.json()["choices"][0]["message"]["content"] == "hi" + # Body should be forwarded byte-for-byte (ignoring possible whitespace). + assert json.loads(_UpstreamHandler.received_body) == payload + + +def test_unknown_path_returns_404(shim_server: str) -> None: + response = httpx.get(f"{shim_server}/something/else", timeout=2.0) + assert response.status_code == 404 diff --git a/uv.lock b/uv.lock index 9a43dcd91e..52c13d4806 100644 --- a/uv.lock +++ b/uv.lock @@ -3912,6 +3912,12 @@ dependencies = [ { name = "nemoguardrails", extra = ["tracing"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] +[package.optional-dependencies] +bench = [ + { name = "httpx", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-12-nemoplatform-cu128') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-12-nemoplatform-cu128' and extra == 'extra-20-nmp-safe-synthesizer-cpu') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-20-nmp-safe-synthesizer-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128') or (sys_platform == 'darwin' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-12-nemoplatform-cu128') or (sys_platform == 'darwin' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128') or (sys_platform == 'darwin' and extra == 'extra-12-nemoplatform-cu128' and extra == 'extra-20-nmp-safe-synthesizer-cpu') or (sys_platform == 'darwin' and extra == 'extra-20-nmp-safe-synthesizer-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128') or (sys_platform == 'linux' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-12-nemoplatform-cu128') or (sys_platform == 'linux' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128') or (sys_platform == 'linux' and extra == 'extra-12-nemoplatform-cu128' and extra == 'extra-20-nmp-safe-synthesizer-cpu') or (sys_platform == 'linux' and extra == 'extra-20-nmp-safe-synthesizer-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128')" }, + { name = "pyyaml", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-12-nemoplatform-cu128') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-12-nemoplatform-cu128' and extra == 'extra-20-nmp-safe-synthesizer-cpu') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-20-nmp-safe-synthesizer-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128') or (sys_platform == 'darwin' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-12-nemoplatform-cu128') or (sys_platform == 'darwin' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128') or (sys_platform == 'darwin' and extra == 'extra-12-nemoplatform-cu128' and extra == 'extra-20-nmp-safe-synthesizer-cpu') or (sys_platform == 'darwin' and extra == 'extra-20-nmp-safe-synthesizer-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128') or (sys_platform == 'linux' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-12-nemoplatform-cu128') or (sys_platform == 'linux' and extra == 'extra-12-nemoplatform-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128') or (sys_platform == 'linux' and extra == 'extra-12-nemoplatform-cu128' and extra == 'extra-20-nmp-safe-synthesizer-cpu') or (sys_platform == 'linux' and extra == 'extra-20-nmp-safe-synthesizer-cpu' and extra == 'extra-20-nmp-safe-synthesizer-cu128')" }, +] + [package.dev-dependencies] dev = [ { name = "pytest", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -3923,10 +3929,13 @@ dev = [ [package.metadata] requires-dist = [ { name = "dataclasses-json", specifier = ">=0.6.7" }, + { name = "httpx", marker = "extra == 'bench'", specifier = ">=0.27" }, { name = "nemo-platform", editable = "packages/nemo_platform" }, { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" }, { name = "nemoguardrails", extras = ["tracing"], specifier = "==0.21.0" }, + { name = "pyyaml", marker = "extra == 'bench'", specifier = ">=6.0" }, ] +provides-extras = ["bench"] [package.metadata.requires-dev] dev = [