From 4c36153ca76dafc44fe84e5119c864abb68713ef Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Wed, 27 May 2026 05:02:19 +0000 Subject: [PATCH 01/17] feat(generation): vllm benchmark harness data models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the pydantic schemas the benchmark harness uses for corpus loading, candidate description, and result serialization. Pure data-models commit; the runner + subprocess wrapper land in the next commit. ## Models - ``TraceHeader``, ``BenchmarkPrompt``, ``BenchmarkCorpus`` — the replayable input schema. ``BenchmarkCorpus.from_trace_jsonl(path)`` loads a header-then-records JSONL into a typed corpus. - ``BenchmarkEngineConfig`` — sparse pydantic model holding vLLM constructor kwargs the harness will forward (attention backend, prefix caching, scheduler caps, kv_cache_dtype, speculative_config, compilation_config, kv_cache_metrics, etc.). Every field is optional and falls through to vLLM defaults when unset. - ``BenchmarkCandidate`` — one configuration to benchmark: name + engine config + sparse sampling overrides + prompt-assembly / batch-dispatch modes. - ``CandidateMetrics`` — per-cell measurements. Cell-specific fields (throughput, acceptance, TTFT, finish reasons, startup) live directly on the model; observability primitives (peak VRAM, KV cache fraction, loadavg, engine runtime config) compose ``CellObservability`` from PR-A's ``vllm_observability`` module. - ``SkipRecord``, ``BenchmarkOutput`` — failure record + matrix result wrapper. ## Composition over duplication ``CandidateMetrics.observability: CellObservability`` is the key DRY mechanism. Instead of re-declaring ``peak_vram_gb`` / ``kv_cache_usage_perc`` / ``loadavg_pre`` / etc. on the benchmark schema (the harness-v2 shape), the benchmark schema *embeds* the production observability event. Downstream consumers can: - Read ``metrics.observability.peak_vram_gb`` for the GPU number. - Read ``metrics.observability.engine_runtime_config`` for the scheduler/cache state. - Call ``metrics.observability.to_wandb_payload()`` to get a flat namespaced dict for wandb logging. Adding a new observability primitive in PR-A automatically flows through to benchmark output without touching this module. ## Type aliases ``PromptAssemblyMode`` and ``BatchDispatchMode`` are exported as ``Literal[...]`` aliases so callers (presets, CLI) can reference them by name without re-deriving the allowed values. ## What's NOT in this commit - The runner (``run_benchmark``). It lands next. - The subprocess wrapper. Lands with the runner. - Presets (``vllm_benchmark_presets``). Separate commit. - Wandb integration (``vllm_benchmark_wandb``). Separate commit. - CLI driver (``tools/vllm_benchmark.py``). Separate commit. - Bracketed-A/B methodology (``condition_label``, ``bracket_position`` on the candidate/metrics schemas). Separate commit; the schema additions are small but the methodology change is conceptually distinct. - Cluster-conditioned analyzer + effect-size + CI. Separate commits. Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark.py | 386 ++++++++++++++++++ 1 file changed, 386 insertions(+) create mode 100644 src/nemo_safe_synthesizer/generation/vllm_benchmark.py diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py new file mode 100644 index 000000000..00209e350 --- /dev/null +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -0,0 +1,386 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark-harness data models for the vLLM backend. + +The harness replays a captured workload corpus (one ``GenerationTrace`` +JSONL) under varying engine + sampling configurations and reports +calibrated metrics per candidate. The models in this module stay +CPU-importable; the runner + subprocess wrapper live in sibling +modules and import vLLM lazily. + +Architecture: + +- :class:`BenchmarkCorpus` — the replayable input (one corpus per + dataset, captured once via the production ``VllmBackend`` trace + surface). Header carries the model reference + LoRA path + the + engine kwargs at capture time. Prompt records carry the original + sampling params so the harness can replay them faithfully. +- :class:`BenchmarkCandidate` — one configuration to benchmark + (engine kwargs overlay + sparse sampling overrides + per-cell + identity for sweep grouping). +- :class:`CandidateMetrics` — per-cell measured outputs: throughput, + acceptance, TTFT, etc. Composes :class:`CellObservability` from PR-A's + ``vllm_observability`` module so the benchmark schema doesn't + re-define observability primitives — it consumes them. +- :class:`BenchmarkOutput` — JSON-serialised result of one matrix + invocation, with skip records for candidates that failed. + +The runner (next commit) is in ``vllm_benchmark.py`` alongside these +models; the subprocess wrapper + single-run entry point are split into +``vllm_benchmark_single_run.py``. +""" + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field + +from ..generation.vllm_observability import CellObservability + +PromptAssemblyMode = Literal["multi_record", "per_record"] +"""Prompt-assembly regime — controls how max_tokens partitions the budget.""" + +BatchDispatchMode = Literal["replicate", "n_fanout"] +"""How corpus prompts get submitted to vLLM. + +- ``'replicate'``: every corpus prompt becomes an independent request + (``n=1``). The default; matches production ``VllmBackend.generate``. +- ``'n_fanout'``: dispatch a single prompt with + ``SamplingParams.n = num_prompts`` so vLLM amortises the shared + schema-prefix prefill across N samples. Only valid when every + corpus prompt is identical; the harness uses the first prompt as + the fanout payload. +""" + + +class TraceHeader(BaseModel): + """Header line at the top of a captured trace JSONL. + + Carries the engine + LoRA + dataset metadata needed to rebuild an + equivalent inference setup at replay time. + """ + + model_config = ConfigDict(extra="allow") + + run_id: str = Field(description="Capture identifier; opaque to the harness, surfaced in JSON.") + pretrained_model: str = Field(description="Model reference (HF name, local path, or ``ModelRef`` string).") + lora_path: Path | None = Field(default=None, description="LoRA adapter path; ``None`` for the bare-model case.") + dataset_schema: dict[str, Any] = Field(description="Tabular schema the processor validates against.") + engine_parameters: dict[str, Any] = Field( + default_factory=dict, + description="vLLM constructor kwargs at capture time (e.g. attention backend, structured-output backend).", + ) + max_tokens_per_example: int | None = Field( + default=None, + description="Resolver hint: per-example max-tokens budget at capture time. Drives ``max_model_len`` resolution.", + ) + + +class BenchmarkPrompt(BaseModel): + """One captured prompt + its original sampling params + reference output. + + The harness replays ``prompt`` through ``LLM.generate(...)`` using + a SamplingParams built from ``original_sampling_params`` + the + candidate's overrides. ``original_output_text`` is preserved for + qualitative diffing but isn't used by the runner directly. + """ + + model_config = ConfigDict(extra="forbid") + + row_index: int = Field(description="Index in the original capture, surfaced for forensics.") + prompt: str = Field(description="The exact prompt text submitted to vLLM at capture time.") + original_sampling_params: dict[str, Any] = Field( + default_factory=dict, + description="Sampling params recorded at capture (temperature, top_p, etc.).", + ) + expected_finish_reason: str | None = Field( + default=None, + description="The finish reason the engine returned at capture time; informational only.", + ) + original_output_text: str = Field( + default="", + description="The text the engine generated at capture time. Preserved for diffing; not consumed by the runner.", + ) + + +class BenchmarkCorpus(BaseModel): + """One captured workload corpus, loaded from a trace JSONL. + + Use :meth:`from_trace_jsonl` to load. The header line establishes + the model + LoRA + dataset_schema; every subsequent record line is + a :class:`BenchmarkPrompt`. + """ + + model_config = ConfigDict(extra="forbid") + + header: TraceHeader + prompts: list[BenchmarkPrompt] + + @classmethod + def from_trace_jsonl(cls, path: str | Path) -> BenchmarkCorpus: + """Load a corpus from a JSONL file containing one ``header`` + many ``record`` lines.""" + path = Path(path) + header: TraceHeader | None = None + prompts: list[BenchmarkPrompt] = [] + with path.open("r", encoding="utf-8") as fh: + for line_no, raw in enumerate(fh, start=1): + line = raw.strip() + if not line: + continue + payload = json.loads(line) + kind = payload.pop("kind", None) + if kind == "header": + if header is not None: + raise ValueError(f"{path}: duplicate header on line {line_no}") + header = TraceHeader.model_validate(payload) + elif kind == "record": + if header is None: + raise ValueError(f"{path}: record on line {line_no} before any header") + prompts.append( + BenchmarkPrompt( + row_index=int(payload["row_index"]), + prompt=str(payload["prompt"]), + original_sampling_params=dict(payload.get("sampling_params") or {}), + expected_finish_reason=payload.get("finish_reason"), + original_output_text=str(payload.get("output_text", "")), + ), + ) + else: + raise ValueError(f"{path}: unknown kind={kind!r} on line {line_no}") + if header is None: + raise ValueError(f"{path}: missing header line") + return cls(header=header, prompts=prompts) + + +class BenchmarkEngineConfig(BaseModel): + """Engine-construction kwargs the harness forwards to ``vllm.LLM(...)``. + + Sparse: every field is optional. Unset fields fall through to vLLM's + own defaults (or to the corpus header's engine_parameters when the + runner builds the engine for a candidate). The runner explicitly + drops ``None``-valued fields when assembling kwargs so vLLM treats + them as "not configured" rather than "override to None". + + This is a benchmark-side schema, not a production config. Production + construction lives in ``VllmBackend.initialize``. Eventually if PR-1's + ``vllm_engine_factory.VllmEngineParameters`` lands, this can compose + against that; for now it stands alone. + """ + + model_config = ConfigDict(extra="forbid") + + attention_backend: str | None = Field( + default=None, + description="vLLM attention backend (``FLASHINFER``, ``FLASH_ATTN``, ``TRITON_ATTN``, etc.). ``None`` or ``'auto'`` leaves it unset.", + ) + structured_generation_backend: str = Field( + default="xgrammar", + description="Structured-outputs backend. ``'xgrammar'`` is vLLM's current default; ``'outlines'`` and ``'guidance'`` are the alternatives.", + ) + max_model_len: int | None = Field( + default=None, + description="Context-window cap forwarded to ``vllm.LLM(max_model_len=...)``. ``None`` lets vLLM auto-resolve from the model.", + ) + enable_prefix_caching: bool | None = Field( + default=None, + description=( + "Forwarded to ``vllm.LLM(enable_prefix_caching=...)``. " + "On for shared-schema tabular workloads (the prefix amortises across the batch); " + "off when measuring per-cell cold-start behaviour. ``None`` keeps vLLM's default." + ), + ) + max_num_seqs: int | None = Field(default=None, description="vLLM scheduler ``max_num_seqs`` cap.") + max_num_batched_tokens: int | None = Field(default=None, description="vLLM scheduler ``max_num_batched_tokens`` cap.") + enable_chunked_prefill: bool | None = Field(default=None, description="Chunked-prefill engagement.") + kv_cache_dtype: str | None = Field( + default=None, + description="KV-cache dtype (``'auto'``, ``'fp8'``, etc.). Halves memory footprint at a per-token quality cost when set to ``'fp8'``.", + ) + seed: int | None = Field( + default=None, + description="Engine-level RNG seed (``vllm.LLM(seed=...)``). Does NOT pin per-request sampling RNG; that's ``sampling_overrides['seed']`` on :class:`BenchmarkCandidate`.", + ) + gpu_memory_utilization: float | None = Field( + default=None, + description="Fraction of GPU memory vLLM may allocate. ``None`` lets the harness pick a sensible default from the host.", + ) + speculative_config: dict[str, Any] | None = Field( + default=None, + description=( + "Speculative-decoding config dict forwarded verbatim to " + "``vllm.LLM(speculative_config=...)``. ``{'method': 'ngram', " + "'num_speculative_tokens': 4, 'prompt_lookup_max': 4}`` is the " + "n-gram preset; ``{'method': 'eagle', 'model': '...'}`` is the " + "draft-model preset." + ), + ) + compilation_config: dict[str, Any] | None = Field( + default=None, + description=( + "Forwarded to ``vllm.LLM(compilation_config=...)``. Controls " + "vLLM's torch.compile / cudagraph capture settings. Useful " + "for perf experiments that probe startup-vs-runtime tradeoffs " + "without an ``enforce_eager`` rebuild." + ), + ) + kv_cache_metrics: bool | None = Field( + default=None, + description=( + "Forwarded to ``vllm.LLM(kv_cache_metrics=...)`` when supported. " + "Enables vLLM's KV-cache residency / hit-rate metrics surface; " + "useful for additional observability beyond what " + "``LLM.get_metrics()`` already provides." + ), + ) + + +class BenchmarkCandidate(BaseModel): + """One configuration to benchmark — engine kwargs + sampling overrides + identity. + + ``sampling_overrides`` is sparse — only fields that differ from the + corpus default. The runner merges it on top of each prompt's + ``original_sampling_params``. Pass ``{"seed": int}`` to pin + ``SamplingParams.seed`` for reproducible acceptance-rate + measurements. + """ + + model_config = ConfigDict(extra="forbid") + + name: str = Field(description="Human-readable identifier, e.g. ``'baseline'`` or ``'prefix_caching=on'``.") + engine_config: BenchmarkEngineConfig = Field( + default_factory=BenchmarkEngineConfig, + description="Engine-construction snapshot for this candidate.", + ) + sampling_overrides: dict[str, Any] = Field( + default_factory=dict, + description="Sparse sampling-params overlay applied on top of the corpus default.", + ) + prompt_overrides: dict[str, Any] | None = Field( + default=None, + description="Reserved for future prompting-strategy variations; runner may not consume yet.", + ) + prompt_mode: PromptAssemblyMode = Field( + default="multi_record", + description=( + "Prompt-assembly regime. ``'multi_record'`` keeps the captured " + "sampling-params unchanged. ``'per_record'`` divides ``max_tokens`` " + "by the per-record hint so each prompt decodes a single record." + ), + ) + batch_dispatch_mode: BatchDispatchMode = Field( + default="replicate", + description="See :data:`BatchDispatchMode` module-level doc.", + ) + + +class CandidateMetrics(BaseModel): + """Measured outputs for one benchmark cell. + + Carries cell-specific bench measurements (throughput, acceptance, + TTFT, etc.) directly; observability primitives (peak VRAM, KV + cache usage, loadavg, engine_runtime_config) are composed from + PR-A's :class:`CellObservability` schema in the ``observability`` + field. This composition keeps the schema DRY — adding a new + observability primitive in PR-A automatically flows through to + benchmark output via ``model_dump()``. + """ + + model_config = ConfigDict(extra="forbid") + + name: str = Field(description="Matches :attr:`BenchmarkCandidate.name`.") + + # Cell-level throughput + acceptance. + raw_tok_s: float = Field(description="Output tokens / wall seconds, ignoring validity.") + acceptance_rate: float = Field(description="Fraction of generated records that passed validation (0..1).") + effective_tok_s: float = Field(description="``raw_tok_s * acceptance_rate``; the operator-relevant headline.") + + # Per-request latency stats. TTFT is queue-inclusive under batched + # submission (vLLM's ``first_token_latency`` = ``first_token_ts - + # arrival_time`` — so prompts later in the batch contribute their + # queue wait). Useful for spotting tail-of-batch wait time but not + # for per-candidate comparisons under varying batch shapes. + ttft_p50_ms: float = Field(description="Median time-to-first-token in milliseconds (queue-inclusive).") + ttft_p99_ms: float = Field(description="99th-percentile time-to-first-token in milliseconds (queue-inclusive).") + + # Counts. + prompts_attempted: int = Field(description="Number of corpus prompts replayed against this candidate.") + prompts_accepted: int = Field(description="Prompts that produced at least one valid record.") + total_output_tokens: int = Field(description="Sum of generated tokens across all replays.") + total_wall_seconds: float = Field(description="Total wall-clock seconds spent generating for this candidate.") + records_per_second: float = Field( + default=0.0, + description="Valid synthetic records produced per wall second of generation.", + ) + + # Startup / overlap accounting. + startup_seconds: float = Field( + default=0.0, + description="Wall seconds the runner blocked waiting on the async engine build.", + ) + simulate_training_overlap_seconds: float = Field( + default=0.0, + description="Seconds the runner slept after kicking off engine build to simulate concurrent training.", + ) + startup_overlap_savings_seconds: float = Field( + default=0.0, + description="Best-effort wall-time savings from overlapping engine init with simulated training.", + ) + + # Finish-reason distribution. + finish_reason_distribution: dict[str, int] = Field( + default_factory=dict, + description="Counts of vLLM finish reasons (``stop``, ``length``, etc.) across replays.", + ) + + # Composed observability (PR-A's schema). The runner builds this from + # ``NvmlPeakSampler.peak_gb`` + ``read_loadavg`` pre/post + + # ``probe_engine_runtime_config`` + ``read_vllm_runtime_metrics``, + # and sets ``flag_did_not_engage`` based on the candidate's intended + # engine config vs the probed runtime config. + observability: CellObservability = Field( + default_factory=CellObservability, + description=( + "Cell-level observability snapshot. Composed from PR-A's schema " + "so benchmark consumers can read e.g. ``metrics.observability." + "peak_vram_gb`` and ``metrics.observability.kv_cache_usage_perc`` " + "without the benchmark schema re-defining those fields." + ), + ) + + +class SkipRecord(BaseModel): + """Per-candidate failure record persisted alongside successful metrics. + + Lets a JSON consumer post-mortem which candidates failed without + parsing the run log. + """ + + model_config = ConfigDict(extra="forbid") + + name: str = Field(description="Matches the candidate's :attr:`BenchmarkCandidate.name`.") + error: str = Field(description="Truncated stderr captured from the child subprocess.") + error_class: str = Field(description="Best-effort exception class name parsed from stderr.") + attempted_at: datetime = Field(description="UTC timestamp when the subprocess was launched.") + + +class BenchmarkOutput(BaseModel): + """Full result of one matrix invocation — JSON-serialised for diffing.""" + + model_config = ConfigDict(extra="forbid") + + created_at: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc), + description="UTC timestamp at which the run was finalised.", + ) + corpus_run_id: str = Field(description="``run_id`` from the captured trace header.") + corpus_size: int = Field(description="Number of prompts in the corpus (informational).") + candidates: list[CandidateMetrics] = Field(description="Per-candidate metrics, in submission order.") + skipped_candidates: list[SkipRecord] = Field( + default_factory=list, + description="Candidates that exited non-zero or produced no result file.", + ) From 28cd644e1393574abaf03ed5e7cc543c2d1b4358 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Wed, 27 May 2026 05:05:21 +0000 Subject: [PATCH 02/17] =?UTF-8?q?feat(generation):=20vllm=20benchmark=20ru?= =?UTF-8?q?nner=20=E2=80=94=20replay=20corpus,=20emit=20metrics?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the core ``run_benchmark`` function + supporting helpers to ``vllm_benchmark.py``. The runner builds a fresh vLLM engine per candidate, replays the corpus prompts through one concurrent ``LLM.generate`` call, and returns a ``CandidateMetrics`` with the composed ``CellObservability`` populated. ## Helpers added - ``_merge_sampling_kwargs`` — composes vLLM ``SamplingParams`` kwargs from the corpus default + candidate overrides; strips capture-time metadata fields (``structured_outputs`` summary) that ``SamplingParams`` doesn't accept. - ``_percentile`` — linear-interpolation percentile (``0.0`` on empty input); used for TTFT p50/p99. - ``_extract_ttft_ms`` — reads per-request first-token latency off ``RequestOutput.metrics``. Tries ``first_token_latency`` (modern, seconds) first; falls back to ``first_token_time - arrival_time`` for older captures. Queue-inclusive under batched submission (the docstring on ``CandidateMetrics.ttft_p50_ms`` flags this explicitly). - ``_build_vllm_kwargs`` — composes ``vllm.LLM(...)`` kwargs from the corpus header (model, LoRA, engine_parameters at capture) + candidate overlay (sparse ``BenchmarkEngineConfig`` fields). Translates ``attention_backend`` → ``attention_config={"backend": ...}`` and ``structured_generation_backend`` → ``structured_outputs_config=StructuredOutputsConfig(backend=...)``. Drops ``None``-valued fields explicitly. - ``_build_engine_async`` — starts ``vllm.LLM(...)`` on a daemon thread, returns ``(thread, ready_event, result_dict)``. The runner can sleep for the simulated-training-overlap window before joining via ``ready_event.wait()``, measuring how much engine-init cost can be hidden behind concurrent training. ## ``run_benchmark`` shape Single function. Wraps the entire body in: - ``NvmlPeakSampler`` context (from PR-A) — peak device VRAM across the whole cell. - ``read_loadavg()`` pre + post (from PR-A) — bracketing host load capture. - ``probe_engine_runtime_config(llm)`` (from PR-A) — captured once after engine init, reused for the observability event. - ``flag_engagement_mismatches(intended, actual)`` (from PR-A) — cross-checks candidate's ``engine_config.model_dump(exclude_none=True)`` against the probe; sets ``flag_did_not_engage`` and logs a warning when mismatches fire. End-of-generation: - ``read_vllm_runtime_metrics(llm)`` (from PR-A) — KV cache, prefix cache, spec accept. - Build :class:`CellObservability` event with all measurements. - Return :class:`CandidateMetrics` with the event composed in. Cleanup runs in ``finally`` so the sampler shuts down + the partial observability event still builds if generation raised mid-batch. ## What's NOT in this commit - Subprocess isolation wrapper (``run_benchmark_in_subprocess``). Lands next so each candidate gets a fresh CUDA context. - The ``vllm_benchmark_single_run`` subprocess entry point. Lands with the wrapper. - ``n_fanout`` dispatch mode. The runner currently dispatches all prompts as independent requests (``n=1`` per request). n_fanout needs the SamplingParams ``n=N`` path; landing in a follow-up. - Streaming step-loop with mid-batch abort (the PR-3-style stopping conditions). Out of scope for now; the synchronous dispatch suffices for the standard benchmark-sliver pattern. - Presets (``vllm_benchmark_presets.py``). Separate commit. - Wandb integration. Separate commit. - CLI driver. Separate commit. ## Dependencies on PR-A This commit lifts the architectural smell I flagged earlier: observability primitives live in PR-A's ``vllm_observability``, the benchmark schema *composes* them rather than re-defining them. The runner imports five symbols from PR-A: - ``NvmlPeakSampler`` - ``CellObservability`` - ``read_loadavg`` - ``probe_engine_runtime_config`` - ``read_vllm_runtime_metrics`` - ``flag_engagement_mismatches`` No PR-1 dependency. The vLLM engine is constructed via direct kwargs (via ``_build_vllm_kwargs``), not via PR-1's ``build_vllm_engine`` factory — keeps PR-B mergeable independently of PR-1's review timing. Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark.py | 326 +++++++++++++++++- 1 file changed, 324 insertions(+), 2 deletions(-) diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index 00209e350..bbcf6f89c 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -34,13 +34,34 @@ from __future__ import annotations import json +import threading +import time from datetime import datetime, timezone from pathlib import Path -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal from pydantic import BaseModel, ConfigDict, Field -from ..generation.vllm_observability import CellObservability +from ..generation.vllm_observability import ( + CellObservability, + NvmlPeakSampler, + flag_engagement_mismatches, + probe_engine_runtime_config, + read_loadavg, + read_vllm_runtime_metrics, +) +from ..observability import get_logger + +if TYPE_CHECKING: + from vllm import LLM + +logger = get_logger(__name__) + +# Fields the corpus carries on each prompt's ``original_sampling_params`` +# that aren't valid kwargs to ``vllm.SamplingParams``. Stripped during +# the merge so the harness's SamplingParams constructor doesn't reject +# capture-time-only metadata. +_NON_SAMPLING_FIELDS: tuple[str, ...] = ("structured_outputs",) PromptAssemblyMode = Literal["multi_record", "per_record"] """Prompt-assembly regime — controls how max_tokens partitions the budget.""" @@ -384,3 +405,304 @@ class BenchmarkOutput(BaseModel): default_factory=list, description="Candidates that exited non-zero or produced no result file.", ) + + +# --------------------------------------------------------------------------- +# Sampling-params + percentile helpers +# --------------------------------------------------------------------------- + + +def _merge_sampling_kwargs(base: dict[str, Any], overrides: dict[str, Any]) -> dict[str, Any]: + """Compose ``SamplingParams`` kwargs from corpus default + candidate overrides. + + The corpus's captured ``original_sampling_params`` may carry fields + that ``vllm.SamplingParams`` doesn't accept (e.g. structured-output + presence summaries). Strip those unless the override explicitly + provides them. Overrides win on conflict. + """ + merged: dict[str, Any] = {**base, **overrides} + for field in _NON_SAMPLING_FIELDS: + if field not in overrides: + merged.pop(field, None) + return merged + + +def _percentile(values: list[float], pct: float) -> float: + """Linear-interpolation percentile (``pct`` in [0, 100]); ``0.0`` on empty input.""" + if not values: + return 0.0 + ordered = sorted(values) + rank = (len(ordered) - 1) * (pct / 100.0) + lo = int(rank) + hi = min(lo + 1, len(ordered) - 1) + if lo == hi: + return ordered[lo] + return ordered[lo] + (rank - lo) * (ordered[hi] - ordered[lo]) + + +def _extract_ttft_ms(output: Any) -> float | None: + """Read TTFT (ms) off a ``RequestOutput``; ``None`` when the engine omits metrics. + + vLLM 0.18+ exposes per-request first-token latency through + ``RequestOutput.metrics.first_token_latency`` (seconds, populated + only when the engine is built with ``disable_log_stats=False``). + Older paths populated separate ``first_token_time`` / ``arrival_time`` + timestamps; fall through to those when the modern field is absent + so corpus captures from older vLLM versions still produce TTFT. + """ + metrics = getattr(output, "metrics", None) + if metrics is None: + return None + latency_s = getattr(metrics, "first_token_latency", None) + if latency_s is not None: + return max(0.0, float(latency_s) * 1000.0) + first = getattr(metrics, "first_token_time", None) + arrival = getattr(metrics, "arrival_time", None) + if first is None or arrival is None: + return None + return max(0.0, (float(first) - float(arrival)) * 1000.0) + + +# --------------------------------------------------------------------------- +# Engine construction +# --------------------------------------------------------------------------- + + +def _build_vllm_kwargs(header: TraceHeader, engine_config: BenchmarkEngineConfig) -> dict[str, Any]: + """Compose ``vllm.LLM(...)`` kwargs from corpus header + candidate overlay. + + The corpus header pins the model + LoRA + dataset_schema. The + candidate's ``engine_config`` overlays sparse engine-side overrides + (attention backend, prefix caching, scheduler caps, etc.). Unset + candidate fields fall through to the header's ``engine_parameters`` + when those are populated, or to vLLM's own defaults otherwise. + + Drops ``None``-valued candidate fields explicitly so vLLM treats + them as "not configured" rather than "override to None". + """ + overlay = engine_config.model_dump(exclude_none=True) + base = dict(header.engine_parameters) + base.update(overlay) + # Required-positional kwargs that aren't in BenchmarkEngineConfig: + base["model"] = header.pretrained_model + base.setdefault("enable_lora", header.lora_path is not None) + # ``attention_backend`` → ``attention_config`` translation. vLLM's + # public API takes a config dict rather than a bare string. + attention_backend = base.pop("attention_backend", None) + if attention_backend not in (None, "auto"): + base["attention_config"] = {"backend": attention_backend} + # ``structured_generation_backend`` → ``structured_outputs_config``. + from vllm.config import StructuredOutputsConfig # noqa: PLC0415 — lazy, vLLM is heavy + + sg_backend = base.pop("structured_generation_backend", None) + if sg_backend is not None: + base["structured_outputs_config"] = StructuredOutputsConfig(backend=sg_backend) + return base + + +def _build_engine_async( + header: TraceHeader, + engine_config: BenchmarkEngineConfig, +) -> tuple[threading.Thread, threading.Event, dict[str, Any]]: + """Start ``vllm.LLM(...)`` in a daemon thread. + + Returns ``(thread, ready_event, result_dict)``. The result dict has + keys ``llm`` (set on success) and ``exception`` (set on failure). + The runner can sleep for the simulated-training-overlap window + before joining via ``ready_event.wait()``, measuring how much of + the engine-init cost can be hidden behind concurrent training. + """ + from vllm import LLM as vLLM # noqa: PLC0415 — lazy + + kwargs = _build_vllm_kwargs(header, engine_config) + ready = threading.Event() + result: dict[str, Any] = {"llm": None, "exception": None} + + def worker() -> None: + try: + result["llm"] = vLLM(**kwargs) + except BaseException as exc: # noqa: BLE001 — surface via the result dict + result["exception"] = exc + finally: + ready.set() + + thread = threading.Thread(target=worker, name="vllm-benchmark-engine-init", daemon=True) + thread.start() + return thread, ready, result + + +# --------------------------------------------------------------------------- +# The runner +# --------------------------------------------------------------------------- + + +def run_benchmark( + candidate: BenchmarkCandidate, + corpus: BenchmarkCorpus, + simulate_training_overlap_seconds: float = 0.0, +) -> CandidateMetrics: + """Replay ``corpus`` against one ``candidate`` and report measured metrics. + + Builds a fresh vLLM engine per candidate so engine-construction + knobs (attention backend, prefix caching, scheduler limits, ...) + actually take effect. The engine build runs in a background thread + so the runner can sleep for ``simulate_training_overlap_seconds`` + before joining; this measures how much of the cold-start cost can + be hidden inside a training phase. + + All corpus prompts are submitted in one concurrent ``LLM.generate`` + call (matching ``VllmBackend.generate()``'s production dispatch + shape) so the shared schema-prefix KV cache amortises across the + batch and prefix caching can actually fire. + + Wraps the entire body in PR-A's :class:`NvmlPeakSampler` context and + emits a composed :class:`CellObservability` on the returned + :class:`CandidateMetrics`. The ``flag_did_not_engage`` bit is set + when the engine's effective runtime config disagrees with the + candidate's intended ``engine_config`` on any checked field. + """ + # Lazy imports — keep this module CPU-importable. + from vllm.lora.request import LoRARequest # noqa: PLC0415 + from vllm.sampling_params import SamplingParams # noqa: PLC0415 + + from ..config.generate import ValidationParameters # noqa: PLC0415 + from .processors import TabularDataProcessor # noqa: PLC0415 + + loadavg_pre = read_loadavg() + vram_sampler = NvmlPeakSampler() + vram_sampler.__enter__() + try: + overlap = max(0.0, simulate_training_overlap_seconds) + _init_thread, engine_ready, init_result = _build_engine_async( + corpus.header, + candidate.engine_config, + ) + if overlap > 0.0: + time.sleep(overlap) + + wait_start = time.monotonic() + engine_ready.wait() + startup_seconds = max(0.0, time.monotonic() - wait_start) + if init_result["exception"] is not None: + raise init_result["exception"] + llm: LLM = init_result["llm"] + + startup_overlap_savings_seconds = ( + min(overlap, startup_seconds + overlap) if overlap > 0.0 else 0.0 + ) + + # Probe the engine's effective runtime config + check for + # candidate-intent / engine-actual disagreements. + engine_runtime_config = probe_engine_runtime_config(llm) + intended = candidate.engine_config.model_dump(exclude_none=True) + mismatches = flag_engagement_mismatches(intended, engine_runtime_config) + if mismatches: + logger.runtime.warning( + "vllm_benchmark.flag_did_not_engage", + extra={"ctx": {"candidate": candidate.name, "mismatches": mismatches}}, + ) + + processor = TabularDataProcessor( + corpus.header.dataset_schema, + config=ValidationParameters(), + tokenizer=None, + ) + lora_request = ( + LoRARequest("lora", 1, str(corpus.header.lora_path)) + if corpus.header.lora_path is not None + else None + ) + + # Build SamplingParams from corpus default + candidate overrides. + if corpus.prompts: + base_sampling = corpus.prompts[0].original_sampling_params + else: + base_sampling = {} + sampling_kwargs = _merge_sampling_kwargs(base_sampling, candidate.sampling_overrides) + # n=1 unless the caller's override sets it (n_fanout sets it explicitly). + sampling_kwargs.setdefault("n", 1) + sampling_params = SamplingParams(**sampling_kwargs) + + # Dispatch. + prompts: list[str] = [p.prompt for p in corpus.prompts] + gen_start = time.perf_counter() + outputs: list[Any] = list( + llm.generate(prompts=prompts, sampling_params=sampling_params, lora_request=lora_request) + ) if prompts else [] + total_wall = max(time.perf_counter() - gen_start, 0.0) + + # Process outputs. + ttft_ms_samples: list[float] = [] + finish_reasons: dict[str, int] = {} + total_output_tokens = 0 + total_valid = 0 + total_invalid = 0 + prompts_accepted = 0 + for output in outputs: + ttft = _extract_ttft_ms(output) + if ttft is not None: + ttft_ms_samples.append(ttft) + best = output.outputs[0] if getattr(output, "outputs", None) else None + if best is None: + continue + total_output_tokens += len(getattr(best, "token_ids", []) or []) + finish_reason = str(getattr(best, "finish_reason", None) or "unknown") + finish_reasons[finish_reason] = finish_reasons.get(finish_reason, 0) + 1 + text = getattr(best, "text", "") or "" + parsed = processor.process(text) + valid_records = getattr(parsed, "valid_records", None) or [] + invalid_records = getattr(parsed, "invalid_records", None) or [] + total_valid += len(valid_records) + total_invalid += len(invalid_records) + if valid_records: + prompts_accepted += 1 + + total_records = total_valid + total_invalid + acceptance = (total_valid / total_records) if total_records > 0 else 0.0 + raw_tok_s = (total_output_tokens / total_wall) if total_wall > 0 else 0.0 + records_per_second = (total_valid / total_wall) if total_wall > 0 else 0.0 + finally: + # Cleanup runs regardless of whether the body raised. Sampler + # always shuts down; observability event always builds (with + # whatever measurements were captured up to the failure). + vram_sampler.__exit__(None, None, None) + + # Read end-of-generation metrics + post-load and assemble the event. + try: + vllm_metrics = read_vllm_runtime_metrics(llm) + except Exception as exc: # noqa: BLE001 — degraded mode + logger.runtime.warning( + "vllm_benchmark.metrics_read_failed", + extra={"ctx": {"error": str(exc)}}, + ) + vllm_metrics = {"kv_cache_usage_perc": None, "prefix_cache_hit_rate": None, "spec_accept_rate": None} + + observability = CellObservability( + peak_vram_gb=vram_sampler.peak_gb, + kv_cache_usage_perc=vllm_metrics["kv_cache_usage_perc"], + prefix_cache_hit_rate=vllm_metrics["prefix_cache_hit_rate"], + spec_accept_rate=vllm_metrics["spec_accept_rate"], + loadavg_pre=loadavg_pre, + loadavg_post=read_loadavg(), + engine_runtime_config=engine_runtime_config, + flag_did_not_engage=bool(mismatches), + ) + + return CandidateMetrics( + name=candidate.name, + raw_tok_s=raw_tok_s, + acceptance_rate=acceptance, + effective_tok_s=raw_tok_s * acceptance, + ttft_p50_ms=_percentile(ttft_ms_samples, 50), + ttft_p99_ms=_percentile(ttft_ms_samples, 99), + prompts_attempted=len(corpus.prompts), + prompts_accepted=prompts_accepted, + total_output_tokens=total_output_tokens, + total_wall_seconds=total_wall, + records_per_second=records_per_second, + startup_seconds=startup_seconds, + simulate_training_overlap_seconds=overlap, + startup_overlap_savings_seconds=startup_overlap_savings_seconds, + finish_reason_distribution=finish_reasons, + observability=observability, + ) From 3ab65b754a33784006d673f64db57d2fbec739f6 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Wed, 27 May 2026 14:17:36 +0000 Subject: [PATCH 03/17] feat(generation): subprocess-isolated benchmark candidate runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each candidate runs in a fresh Python interpreter via ``python -m nemo_safe_synthesizer.generation.vllm_benchmark_single_run`` so vLLM's module-level CUDA + DRAM state is reclaimed by the OS between candidates. The in-process Python runtime can't clean up vLLM's globals reliably; subprocess isolation is the only way to run a multi-candidate matrix without state bleed. Adds to ``vllm_benchmark.py``: - ``SUBPROCESS_STDERR_LIMIT = 500`` — captured-stderr cap. - ``_truncate_stderr`` — trim with a tail-preserving sentinel. - ``_parse_error_class`` — best-effort exception class extraction from a Python traceback. - ``SubprocessRunResult`` — pydantic model carrying either ``metrics`` (on success) or ``error`` + ``error_class`` (on failure). - ``run_benchmark_in_subprocess(candidate, corpus_path, ...)`` — spawns the child, reads back the JSON-serialised ``CandidateMetrics``, returns ``SubprocessRunResult``. New module ``vllm_benchmark_single_run.py``: - argparse-based entry point: ``--candidate`` (JSON), ``--corpus`` (path), ``--result-out`` (path), ``--simulate-training-overlap-seconds`` (float). - Loads the candidate + corpus, calls ``run_benchmark``, writes the resulting ``CandidateMetrics`` JSON. Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark.py | 99 +++++++++++++++++++ .../generation/vllm_benchmark_single_run.py | 50 ++++++++++ 2 files changed, 149 insertions(+) create mode 100644 src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index bbcf6f89c..dbcc16039 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -34,6 +34,9 @@ from __future__ import annotations import json +import subprocess +import sys +import tempfile import threading import time from datetime import datetime, timezone @@ -63,6 +66,9 @@ # capture-time-only metadata. _NON_SAMPLING_FIELDS: tuple[str, ...] = ("structured_outputs",) +SUBPROCESS_STDERR_LIMIT: int = 500 +"""Maximum bytes of captured stderr to record on a subprocess failure.""" + PromptAssemblyMode = Literal["multi_record", "per_record"] """Prompt-assembly regime — controls how max_tokens partitions the budget.""" @@ -706,3 +712,96 @@ def run_benchmark( finish_reason_distribution=finish_reasons, observability=observability, ) + + +# --------------------------------------------------------------------------- +# Subprocess isolation +# --------------------------------------------------------------------------- + + +class SubprocessRunResult(BaseModel): + """Outcome of one subprocess-isolated candidate run.""" + + model_config = ConfigDict(extra="forbid") + + metrics: CandidateMetrics | None = Field(default=None, description="Populated when the child exited successfully.") + error: str | None = Field(default=None, description="Captured stderr summary; populated on non-zero exit.") + error_class: str | None = Field(default=None, description="Best-effort exception class name parsed from stderr.") + + +def _truncate_stderr(stderr: str, limit: int = SUBPROCESS_STDERR_LIMIT) -> str: + """Trim ``stderr`` to ``limit`` bytes, keeping the tail.""" + stderr = stderr.strip() + if len(stderr) <= limit: + return stderr + head = "...[truncated]..." + return head + stderr[-(limit - len(head)) :] + + +def _parse_error_class(stderr: str) -> str: + """Best-effort parse of the exception class name from a Python traceback.""" + for line in reversed(stderr.strip().splitlines()): + stripped = line.strip() + if not stripped: + continue + head = stripped.split(":", 1)[0] + if head.isidentifier() or "." in head: + return head + return "Error" + return "Error" + + +def run_benchmark_in_subprocess( + candidate: BenchmarkCandidate, + corpus_path: str | Path, + simulate_training_overlap_seconds: float = 0.0, +) -> SubprocessRunResult: + """Run one candidate in a child process the OS reclaims on exit. + + Spawns ``python -m nemo_safe_synthesizer.generation.vllm_benchmark_single_run`` + with the candidate JSON-encoded as argv. Child writes + ``CandidateMetrics`` JSON to a temp file; parent reads it back. + + Subprocess isolation is what makes a multi-candidate matrix + reliable on this stack: vLLM holds significant CUDA + DRAM state + in module-level globals that the in-process Python runtime can't + clean up between candidates. Each candidate runs in a fresh + interpreter; the OS reclaims everything on child exit. + """ + candidate_json = candidate.model_dump_json() + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False, encoding="utf-8") as result_fh: + result_path = Path(result_fh.name) + try: + completed = subprocess.run( + [ + sys.executable, + "-m", + "nemo_safe_synthesizer.generation.vllm_benchmark_single_run", + "--candidate", + candidate_json, + "--corpus", + str(corpus_path), + "--result-out", + str(result_path), + "--simulate-training-overlap-seconds", + str(simulate_training_overlap_seconds), + ], + capture_output=True, + text=True, + check=False, + ) + if completed.returncode != 0: + stderr = _truncate_stderr(completed.stderr or completed.stdout) + return SubprocessRunResult( + error=stderr or f"subprocess exit {completed.returncode}", + error_class=_parse_error_class(completed.stderr or completed.stdout), + ) + if not result_path.exists() or result_path.stat().st_size == 0: + return SubprocessRunResult( + error="subprocess exited 0 but produced no result file", + error_class="RuntimeError", + ) + metrics = CandidateMetrics.model_validate_json(result_path.read_text(encoding="utf-8")) + return SubprocessRunResult(metrics=metrics) + finally: + result_path.unlink(missing_ok=True) diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py new file mode 100644 index 000000000..bda272185 --- /dev/null +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Subprocess entry point for one isolated benchmark cell. + +Invoked by :func:`vllm_benchmark.run_benchmark_in_subprocess` via +``python -m nemo_safe_synthesizer.generation.vllm_benchmark_single_run``. +Loads the candidate from the ``--candidate`` JSON argv, loads the +corpus from ``--corpus``, runs :func:`vllm_benchmark.run_benchmark`, +and writes the resulting ``CandidateMetrics`` JSON to ``--result-out``. + +Subprocess isolation is what makes a multi-candidate matrix reliable +on this stack — vLLM holds significant CUDA + DRAM state in +module-level globals; running each candidate in a fresh interpreter +lets the OS reclaim everything on exit. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from .vllm_benchmark import BenchmarkCandidate, BenchmarkCorpus, run_benchmark + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run one benchmark candidate in isolation.") + parser.add_argument("--candidate", required=True, help="JSON-serialised BenchmarkCandidate.") + parser.add_argument("--corpus", required=True, help="Path to the corpus JSONL.") + parser.add_argument("--result-out", required=True, help="Path to write the CandidateMetrics JSON.") + parser.add_argument( + "--simulate-training-overlap-seconds", + type=float, + default=0.0, + help="Seconds to sleep after kicking off engine init (simulates concurrent training).", + ) + args = parser.parse_args() + + candidate = BenchmarkCandidate.model_validate_json(args.candidate) + corpus = BenchmarkCorpus.from_trace_jsonl(args.corpus) + metrics = run_benchmark( + candidate=candidate, + corpus=corpus, + simulate_training_overlap_seconds=args.simulate_training_overlap_seconds, + ) + Path(args.result_out).write_text(metrics.model_dump_json(indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() From 0e3d17edf8d53db1882c2fae5a5337a5d3f33a9d Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Wed, 27 May 2026 14:18:51 +0000 Subject: [PATCH 04/17] feat(generation): vllm benchmark preset matrices MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds ``vllm_benchmark_presets.py`` — the named preset matrices the CLI driver exposes via ``--candidates ``. Each preset takes the corpus's default ``BenchmarkEngineConfig`` and returns one or more ``BenchmarkCandidate``s with the appropriate engine overrides applied. ## Methodology constants - ``DEFAULT_BENCHMARK_SEED = 42`` — every preset-built candidate pins ``SamplingParams.seed`` to this. Prior triages observed acceptance-rate CoV up to ~14% from per-request RNG drift; seed-pinning collapses that to in-cluster noise (~3% CoV; ~5% pooled on long-output workloads). - ``DEFAULT_MAX_MODEL_LEN = 4096`` — fallback context-window cap when the corpus header doesn't specify. Stops vLLM from over-provisioning the KV cache to the tokenizer's reported max (32768 on Mistral-7B). ## Sweep matrices - ``baseline`` — pass-through. - ``attention_backend_sweep`` — ``FLASHINFER``, ``FLASH_ATTN``, ``TRITON_ATTN``. - ``prefix_caching_sweep`` — single ``prefix_caching=on`` candidate (the implicit baseline is off). - ``batching_sweep`` — (max_num_seqs, max_num_batched_tokens) = (128, 4096), (256, 8192), (512, 16384). - ``structured_backend_sweep`` — ``xgrammar``, ``outlines``, ``guidance``. - ``max_model_len_sweep`` — 2048, 4096, 8192. - ``default_matrix`` — concatenation of all sweeps, deduplicated by (engine_config, sampling_overrides) — two candidates with identical config but different names are treated as duplicates because the engine + sampling combo is what the runner actually measures. ## What's NOT in this commit - ``bracketed_ab`` family (the methodology presets that interleave baseline + candidate cells with seed-pinned bracketing). Separate commit; depends on the schema additions for ``condition_label`` + ``bracket_position`` on ``BenchmarkCandidate`` / ``CandidateMetrics``. Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark_presets.py | 197 ++++++++++++++++++ 1 file changed, 197 insertions(+) create mode 100644 src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py new file mode 100644 index 000000000..066564cb9 --- /dev/null +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark preset matrices for the vLLM harness. + +A "preset" is a function that takes the corpus's default +:class:`BenchmarkEngineConfig` (extracted from the trace header) and +returns one or more :class:`BenchmarkCandidate` instances. The CLI +driver exposes these by name via ``--candidates ``. + +Every preset-built candidate pins ``SamplingParams.seed`` to +:data:`DEFAULT_BENCHMARK_SEED` (via :func:`_seeded_overrides`) so the +per-request RNG portion of acceptance-rate variance is eliminated. +Residual variance is structural non-RNG (numerical / scheduler +interleaving) and is handled post-hoc by the cluster-conditioned +analyzer (separate commit). +""" + +from __future__ import annotations + +import json +from typing import Any + +from .vllm_benchmark import BenchmarkCandidate, BenchmarkEngineConfig + +DEFAULT_BENCHMARK_SEED: int = 42 +"""Default ``SamplingParams.seed`` value for preset-built candidates. + +Prior triages observed acceptance-rate CoV up to ~14% on long-output +workloads from per-request JSON-schema sampling RNG drift; pinning the +seed collapses that variance to in-cluster noise (~3% CoV; ~5% pooled +on long-output workloads per the 2026-05-26 seed-pin verification). +42 is conventional; any fixed int works. +""" + +DEFAULT_MAX_MODEL_LEN: int = 4096 +"""Fallback ``max_model_len`` for presets when the corpus header doesn't have a resolver hint. + +vLLM 0.20 defaults this to whatever the model's tokenizer reports, which +for Mistral-7B is 32768 — over-provisions the KV cache budget for the +tabular workloads we benchmark. +""" + +ATTENTION_BACKENDS: tuple[str, ...] = ( + "FLASHINFER", + "FLASH_ATTN", + "TRITON_ATTN", +) +"""CUDA attention backends the sweep covers; excludes ROCm/XPU/MLA variants.""" + +STRUCTURED_BACKENDS: tuple[str, ...] = ("xgrammar", "outlines", "guidance") +"""Structured-output backends the sweep covers.""" + +BATCHING_STEPS: tuple[tuple[int, int], ...] = ( + (128, 4096), + (256, 8192), + (512, 16384), +) +"""(max_num_seqs, max_num_batched_tokens) steps for the batching sweep.""" + +MAX_MODEL_LEN_STEPS: tuple[int, ...] = (2048, 4096, 8192) +"""``max_model_len`` steps for the max-model-len sweep.""" + + +def _seeded_overrides(seed: int | None = DEFAULT_BENCHMARK_SEED, *, extra: dict[str, Any] | None = None) -> dict[str, Any]: + """Build a ``sampling_overrides`` dict that pins seed + optional extras. + + Returns a fresh dict so callers can mutate without affecting other + candidates. ``seed=None`` opts out of seed pinning. + """ + overrides: dict[str, Any] = {} + if seed is not None: + overrides["seed"] = seed + if extra: + overrides.update(extra) + return overrides + + +def _with_default_max_model_len(base: BenchmarkEngineConfig) -> BenchmarkEngineConfig: + """Set ``max_model_len`` to :data:`DEFAULT_MAX_MODEL_LEN` when unset. + + Saves the KV-cache budget from over-provisioning to the model's + tokenizer-reported maximum. + """ + if base.max_model_len is not None: + return base + return base.model_copy(update={"max_model_len": DEFAULT_MAX_MODEL_LEN}) + + +def _named_copy(base: BenchmarkEngineConfig, name: str, **updates: Any) -> BenchmarkCandidate: + """Build a candidate that overrides exactly ``updates`` on top of ``base``. + + Seed-pinned via :data:`DEFAULT_BENCHMARK_SEED`. Callers needing + un-pinned RNG should construct ``BenchmarkCandidate`` directly. + """ + return BenchmarkCandidate( + name=name, + engine_config=_with_default_max_model_len(base).model_copy(update=updates), + sampling_overrides=_seeded_overrides(), + ) + + +def baseline(base: BenchmarkEngineConfig) -> BenchmarkCandidate: + """Pass-through candidate — runs the corpus's default config unchanged. + + Still seed-pinned for reproducibility. See :func:`_named_copy`. + """ + return BenchmarkCandidate( + name="baseline", + engine_config=_with_default_max_model_len(base), + sampling_overrides=_seeded_overrides(), + ) + + +def attention_backend_sweep(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: + """One candidate per attention backend in :data:`ATTENTION_BACKENDS`.""" + return [ + _named_copy(base, f"attention_backend={backend}", attention_backend=backend) + for backend in ATTENTION_BACKENDS + ] + + +def prefix_caching_sweep(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: + """Probe ``enable_prefix_caching=True`` against the implicit-off baseline. + + The single emitted candidate enables prefix caching; the off case + is the implicit baseline because vLLM defaults + ``enable_prefix_caching=False`` and the harness leaves the field + unset in :func:`baseline`. + """ + return [_named_copy(base, "prefix_caching=on", enable_prefix_caching=True)] + + +def batching_sweep(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: + """Vary ``max_num_seqs`` and ``max_num_batched_tokens`` across :data:`BATCHING_STEPS`.""" + return [ + _named_copy( + base, + f"batch_seqs={seqs}_tokens={tokens}", + max_num_seqs=seqs, + max_num_batched_tokens=tokens, + ) + for seqs, tokens in BATCHING_STEPS + ] + + +def structured_backend_sweep(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: + """One candidate per structured-output backend in :data:`STRUCTURED_BACKENDS`.""" + return [ + _named_copy(base, f"structured_backend={backend}", structured_generation_backend=backend) + for backend in STRUCTURED_BACKENDS + ] + + +def max_model_len_sweep(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: + """One candidate per ``max_model_len`` step in :data:`MAX_MODEL_LEN_STEPS`.""" + return [_named_copy(base, f"max_model_len={length}", max_model_len=length) for length in MAX_MODEL_LEN_STEPS] + + +def default_matrix(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: + """Concatenation of every sweep, deduplicated by (engine config, overrides). + + Dedup is conservative — two candidates with identical engine config + but different ``name`` are still considered duplicates because the + resulting engine + sampling combination is what the runner measures. + """ + seen: set[tuple[Any, ...]] = set() + out: list[BenchmarkCandidate] = [] + for builder in ( + attention_backend_sweep, + prefix_caching_sweep, + batching_sweep, + structured_backend_sweep, + max_model_len_sweep, + ): + for cand in builder(base): + key = ( + cand.engine_config.model_dump_json(), + json.dumps(cand.sampling_overrides, sort_keys=True), + ) + if key in seen: + continue + seen.add(key) + out.append(cand) + return out + + +# Map of preset name → callable used by the CLI to resolve ``--candidates``. +PRESETS: dict[str, object] = { + "baseline": lambda base: [baseline(base)], + "attention_backend_sweep": attention_backend_sweep, + "prefix_caching_sweep": prefix_caching_sweep, + "batching_sweep": batching_sweep, + "structured_backend_sweep": structured_backend_sweep, + "max_model_len_sweep": max_model_len_sweep, + "default_matrix": default_matrix, +} From 4ab4f6ab7c1386c859159f07ceefa5efc689711d Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Wed, 27 May 2026 14:19:51 +0000 Subject: [PATCH 05/17] feat(generation): wandb integration for benchmark cells (per-cell new-run) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds ``vllm_benchmark_wandb.py`` — the benchmark-side wandb wiring. Each benchmark cell becomes one wandb run, grouped by sweep ID and tagged by condition_label + dataset + 'benchmark' job_type. ## Distinct from PR-A's production wandb pattern PR-A's ``log_cell_observability(event)`` logs to the *currently active* wandb run (production logs cells as a time-series within one ``safe-synthesizer run generate`` invocation). This module's ``init_cell_run`` opens a *new* wandb run per benchmark cell so each candidate is its own row in the wandb UI — appropriate for sweep-style isolation where you want per-condition aggregates. Both patterns coexist because they're the right shape for different use cases. They reuse PR-A's ``WandbSettings`` so env-var handling (``WANDB_MODE``, ``WANDB_PROJECT`` / ``NSS_WANDB_PROJECT``) is consistent. ## API - ``resolve_sweep_id() -> str`` — ``WANDB_RUN_GROUP`` env when set; auto-generated timestamp otherwise. - ``init_cell_run(*, candidate_name, candidate_idx, total, corpus_run_id, corpus_size, sweep_id, candidate_condition_label='', candidate_bracket_position=0)`` — opens a wandb run; returns None when disabled/unavailable. Precedence for label + position: candidate-carried > env var > positional fallback. - ``_flatten_metrics(metrics)`` — flattens ``CandidateMetrics`` for ``wandb.log``. Direct fields land top-level; ``finish_reason_distribution`` flattens to ``finish_reason/`` scalars; ``observability`` is delegated to ``CellObservability.to_wandb_payload()`` with the ``vllm_cell`` prefix so production + benchmark agree on observability key namespacing. - ``log_and_finish(run, metrics, exit_code=0)`` — flattens, logs, finishes. No-op when ``run is None``. ## Auth Via ``$HOME/.netrc`` (set up by ``wandb login``). ``WANDB_API_KEY`` is NOT read — passing the key via env leaks it via ``ps auxe``. ## Soft dependency Missing wandb / netrc / init exception → warning + None return, harness continues. The benchmark JSON output is still written. Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark_wandb.py | 166 ++++++++++++++++++ 1 file changed, 166 insertions(+) create mode 100644 src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py new file mode 100644 index 000000000..0aa739ad3 --- /dev/null +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark-side wandb integration — per-cell new-run mode. + +Each benchmark cell becomes one wandb run (grouped by sweep ID, tagged +by condition_label + dataset). Distinct from PR-A's production +``log_cell_observability`` pattern, which logs to the *currently +active* wandb run — production logs cells as a time-series within one +run, benchmark logs each cell as its own run for per-condition +isolation in the wandb UI. + +Reuses :class:`WandbSettings` from ``nemo_safe_synthesizer.cli.wandb_setup`` +so env-var handling (``WANDB_MODE`` defaults to ``disabled``, +``WANDB_PROJECT`` / ``NSS_WANDB_PROJECT`` precedence) matches the rest +of the pipeline. + +Soft dependency: wandb is observability, not a hard requirement. Any +failure (missing package, missing netrc, ``wandb.init`` exception) +logs a warning and the harness continues — the benchmark JSON output +is still written. +""" + +from __future__ import annotations + +import os +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any + +from ..cli.wandb_setup import WandbMode, WandbPhase, WandbSettings +from ..observability import get_logger + +if TYPE_CHECKING: + from .vllm_benchmark import CandidateMetrics + +logger = get_logger(__name__) + +# A benchmark cell is structurally a GENERATE invocation that's +# measured rather than consumed. Phase aligns with production GENERATE +# runs; job_type distinguishes benchmark from production at the +# wandb-UI level. +_BENCHMARK_JOB_TYPE: str = "benchmark" + + +def resolve_sweep_id() -> str: + """Resolve the wandb group identifier for the current sweep. + + Reads ``WANDB_RUN_GROUP`` when set (the orchestrator sets this once + per sweep to group all cells). Falls back to an auto-generated + timestamp so single-cell invocations don't all collapse into one + bucket. + """ + return os.environ.get("WANDB_RUN_GROUP") or f"sweep-{datetime.now(timezone.utc):%Y%m%d-%H%M%S}" + + +def init_cell_run( + *, + candidate_name: str, + candidate_idx: int, + total: int, + corpus_run_id: str, + corpus_size: int, + sweep_id: str, + candidate_condition_label: str = "", + candidate_bracket_position: int = 0, +) -> Any: + """Open one wandb run for a single benchmark cell. + + Returns the run object on success, ``None`` when wandb is disabled, + misconfigured, or the package is unavailable. Callers must treat + ``None`` as "no wandb"; subsequent log/finish helpers are no-ops + on ``None``. + + Precedence for ``condition_label`` / ``bracket_position``: + candidate-carried values win when non-empty/non-zero; otherwise + falls back to ``BENCHMARK_CONDITION_LABEL`` / + ``BENCHMARK_BRACKET_POSITION`` env vars; finally falls back to + ``candidate_name`` / ``candidate_idx - 1``. The candidate-carried + path is set by the ``bracketed_ab`` preset family. + + Auth via ``~/.netrc`` (set up by ``wandb login``). + ``WANDB_API_KEY`` is NOT read — passing the key via env leaks it + via ``ps auxe``. + """ + settings = WandbSettings() + if settings.wandb_mode == WandbMode.DISABLED: + return None + try: + import wandb # noqa: PLC0415 — soft dependency + except ImportError: + logger.warning("wandb not installed; skipping wandb integration for this cell") + return None + + condition_label = ( + candidate_condition_label + or os.environ.get("BENCHMARK_CONDITION_LABEL") + or candidate_name + ) + dataset = os.environ.get("BENCHMARK_DATASET", "unknown") + bracket_position = ( + candidate_bracket_position + if candidate_bracket_position > 0 + else int(os.environ.get("BENCHMARK_BRACKET_POSITION", str(candidate_idx - 1))) + ) + tags = [t for t in (condition_label, dataset, _BENCHMARK_JOB_TYPE) if t and t != "unknown"] + try: + return wandb.init( + project=settings.effective_wandb_project, + name=candidate_name, + group=sweep_id, + job_type=_BENCHMARK_JOB_TYPE, + tags=tags, + mode=settings.wandb_mode.value, + reinit=True, + config={ + "candidate_name": candidate_name, + "candidate_idx": candidate_idx, + "total_candidates": total, + "corpus_run_id": corpus_run_id, + "corpus_size": corpus_size, + "condition_label": condition_label, + "dataset": dataset, + "bracket_position": bracket_position, + "sweep_id": sweep_id, + "phase": WandbPhase.GENERATE.value, + }, + ) + except Exception as exc: # noqa: BLE001 — degraded mode by design + logger.warning("wandb.init failed; continuing without wandb", exc_info=exc) + return None + + +def _flatten_metrics(metrics: CandidateMetrics) -> dict[str, Any]: + """Project ``CandidateMetrics`` into a flat dict for ``wandb.log``. + + Direct fields land at the top level. ``finish_reason_distribution`` + is flattened to ``finish_reason/`` scalars. The composed + ``observability`` field is delegated to + :meth:`CellObservability.to_wandb_payload` with the + ``vllm_cell`` prefix so production + benchmark agree on + observability key namespacing. + """ + payload = metrics.model_dump(exclude={"observability", "finish_reason_distribution"}) + fr_dist = metrics.finish_reason_distribution or {} + payload.update({f"finish_reason/{k}": v for k, v in fr_dist.items()}) + payload.update(metrics.observability.to_wandb_payload()) + return payload + + +def log_and_finish(run: Any, metrics: CandidateMetrics | None, exit_code: int = 0) -> None: + """Log metrics (if any) and close the wandb run. Swallows wandb errors. + + Called on both success (``metrics`` populated, ``exit_code=0``) and + skip paths (``metrics=None``, ``exit_code=1``). No-op when ``run`` + is ``None`` so callers don't need to branch. + """ + if run is None: + return + try: + if metrics is not None: + import wandb # noqa: PLC0415 + + wandb.log(_flatten_metrics(metrics)) + run.finish(exit_code=exit_code) + except Exception as exc: # noqa: BLE001 — degraded mode + logger.warning("wandb finish failed; continuing", exc_info=exc) From 94492712e455b8c4c056107bd9e6aba3d0390749 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Wed, 27 May 2026 14:21:28 +0000 Subject: [PATCH 06/17] feat(generation): vllm benchmark CLI driver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds ``tools/vllm_benchmark.py`` — the click-based CLI exposing ``list`` / ``run`` / ``compare`` subcommands. ## Subcommands - ``list`` — prints available preset names. - ``run CORPUS --output PATH [--candidates PRESET | --candidates-file PATH] [--simulate-training-overlap-seconds N]`` — loads the corpus, resolves candidates, runs each in subprocess isolation (``run_benchmark_in_subprocess``), logs to wandb (via the benchmark-side per-cell wiring), writes the BenchmarkOutput JSON. - ``compare PATH`` — renders a saved BenchmarkOutput as a candidate-by-metric table, sorted by effective_tok_s. Reads ``peak_vram_gb`` from ``metrics.observability`` (composed from PR-A's schema). ## Bug fix in this commit ``BenchmarkEngineConfig`` switched from ``extra='forbid'`` to ``extra='ignore'`` so the CLI can validate the corpus header's raw ``engine_parameters`` dict into the typed config without raising on capture-time-only fields the model doesn't expose. The typed fields get extracted; everything else flows through ``_build_vllm_kwargs``'s base layer regardless. ## Wandb wiring Each candidate gets its own wandb run via ``init_cell_run``, grouped by ``sweep_id`` (resolved from ``WANDB_RUN_GROUP`` env or auto-timestamp). ``log_and_finish`` flattens ``CandidateMetrics`` (including the composed ``CellObservability``) and closes the run with the right exit code. Skip records also finish the run (exit_code=1) so failed candidates appear in the wandb UI rather than silently disappearing. Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark.py | 7 +- tools/vllm_benchmark.py | 215 ++++++++++++++++++ 2 files changed, 221 insertions(+), 1 deletion(-) create mode 100644 tools/vllm_benchmark.py diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index dbcc16039..c2315a4f8 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -197,9 +197,14 @@ class BenchmarkEngineConfig(BaseModel): construction lives in ``VllmBackend.initialize``. Eventually if PR-1's ``vllm_engine_factory.VllmEngineParameters`` lands, this can compose against that; for now it stands alone. + + ``extra='ignore'`` so the CLI can validate a corpus header's raw + ``engine_parameters`` dict into this model — header dicts may carry + capture-time kwargs we don't expose as typed fields (those flow + through ``_build_vllm_kwargs`` as the base layer regardless). """ - model_config = ConfigDict(extra="forbid") + model_config = ConfigDict(extra="ignore") attention_backend: str | None = Field( default=None, diff --git a/tools/vllm_benchmark.py b/tools/vllm_benchmark.py new file mode 100644 index 000000000..c7547442a --- /dev/null +++ b/tools/vllm_benchmark.py @@ -0,0 +1,215 @@ +#!/usr/bin/env -S uv run +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +r"""vllm-benchmark: drive the vLLM benchmark harness from the command line. + +Subcommands: + list Print the available preset matrices. + run CORPUS --output PATH ... Replay CORPUS against the chosen candidates + and persist a BenchmarkOutput JSON. + compare PATH Render a previously-saved BenchmarkOutput + as a candidate-by-metric table. + +Invocation:: + + uv run --frozen --extra cu129 --extra engine --group dev \ + python tools/vllm_benchmark.py list + + uv run --frozen --extra cu129 --extra engine --group dev \ + python tools/vllm_benchmark.py run \ + /path/to/trace.jsonl \ + --output /path/to/benchmark.json \ + --candidates default_matrix +""" + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path + +import click +from rich.console import Console +from rich.table import Table + +from nemo_safe_synthesizer.generation.vllm_benchmark import ( + BenchmarkCandidate, + BenchmarkCorpus, + BenchmarkEngineConfig, + BenchmarkOutput, + CandidateMetrics, + SkipRecord, +) +from nemo_safe_synthesizer.generation.vllm_benchmark_presets import PRESETS +from nemo_safe_synthesizer.generation.vllm_benchmark_wandb import ( + init_cell_run, + log_and_finish, + resolve_sweep_id, +) + +console = Console() + + +@click.group() +def cli() -> None: + """vLLM benchmark harness CLI.""" + + +@cli.command("list") +def list_cmd() -> None: + """Print the available preset names.""" + for name in sorted(PRESETS): + console.print(name) + + +def _resolve_candidates( + base: object, + preset_name: str | None, + candidates_file: str | None, +) -> list[BenchmarkCandidate]: + """Resolve the candidate list from a preset name or a JSON file. + + Exactly one of ``preset_name`` / ``candidates_file`` must be set. + The candidates-file shape is ``{"candidates": [BenchmarkCandidate, ...]}``. + """ + if preset_name and candidates_file: + raise click.UsageError("--candidates and --candidates-file are mutually exclusive.") + if preset_name: + if preset_name not in PRESETS: + raise click.UsageError(f"Unknown preset {preset_name!r}; available: {sorted(PRESETS)}") + resolved = PRESETS[preset_name](base) # ty: ignore[invalid-argument-type] + return resolved if isinstance(resolved, list) else [resolved] + if candidates_file: + doc = json.loads(Path(candidates_file).read_text(encoding="utf-8")) + return [BenchmarkCandidate.model_validate(c) for c in doc["candidates"]] + raise click.UsageError("One of --candidates or --candidates-file is required.") + + +@cli.command("run") +@click.argument("corpus_path", type=click.Path(exists=True, dir_okay=False, path_type=Path)) +@click.option( + "--output", + "output_path", + required=True, + type=click.Path(dir_okay=False, path_type=Path), + help="Where to write the BenchmarkOutput JSON.", +) +@click.option("--candidates", "preset_name", default=None, help=f"Preset name. One of: {sorted(PRESETS)}.") +@click.option( + "--candidates-file", + "candidates_file", + default=None, + type=click.Path(exists=True, dir_okay=False), + help="Custom JSON file with a 'candidates' list of BenchmarkCandidate objects.", +) +@click.option( + "--simulate-training-overlap-seconds", + "simulate_training_overlap_seconds", + default=0.0, + type=float, + show_default=True, + help="Per-candidate seconds to sleep after engine-init kickoff (simulates concurrent training).", +) +def run_cmd( + corpus_path: Path, + output_path: Path, + preset_name: str | None, + candidates_file: str | None, + simulate_training_overlap_seconds: float, +) -> None: + """Replay CORPUS_PATH against the chosen candidates and persist results.""" + # Lazy import so ``list`` and ``compare`` work without spinning up vLLM. + from nemo_safe_synthesizer.generation.vllm_benchmark import run_benchmark_in_subprocess + + corpus = BenchmarkCorpus.from_trace_jsonl(corpus_path) + base = BenchmarkEngineConfig.model_validate(corpus.header.engine_parameters or {}) + candidates = _resolve_candidates(base, preset_name, candidates_file) + if not candidates: + raise click.UsageError("Resolved candidate list is empty.") + + results: list[CandidateMetrics] = [] + skipped: list[SkipRecord] = [] + sweep_id = resolve_sweep_id() + for idx, candidate in enumerate(candidates, start=1): + console.print(f"[{idx}/{len(candidates)}] running candidate {candidate.name!r}") + wandb_run = init_cell_run( + candidate_name=candidate.name, + candidate_idx=idx, + total=len(candidates), + corpus_run_id=corpus.header.run_id, + corpus_size=len(corpus.prompts), + sweep_id=sweep_id, + candidate_condition_label=getattr(candidate, "condition_label", ""), + candidate_bracket_position=getattr(candidate, "bracket_position", 0), + ) + result = run_benchmark_in_subprocess( + candidate, + corpus_path, + simulate_training_overlap_seconds=simulate_training_overlap_seconds, + ) + if result.metrics is None: + console.print(f" [yellow]skipped:[/yellow] {result.error_class or 'Error'}: {result.error}") + skipped.append( + SkipRecord( + name=candidate.name, + error=result.error or "subprocess failed", + error_class=result.error_class or "Error", + attempted_at=datetime.now(timezone.utc), + ), + ) + log_and_finish(wandb_run, metrics=None, exit_code=1) + continue + results.append(result.metrics) + log_and_finish(wandb_run, metrics=result.metrics, exit_code=0) + console.print( + f" raw={result.metrics.raw_tok_s:.1f} tok/s " + f"accept={result.metrics.acceptance_rate:.3f} " + f"effective={result.metrics.effective_tok_s:.1f} tok/s", + ) + + output = BenchmarkOutput( + corpus_run_id=corpus.header.run_id, + corpus_size=len(corpus.prompts), + candidates=results, + skipped_candidates=skipped, + ) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(output.model_dump_json(indent=2), encoding="utf-8") + console.print(f"[green]wrote[/green] {output_path} ({len(results)}/{len(candidates)} ok, {len(skipped)} skipped)") + + +@cli.command("compare") +@click.argument("output_path", type=click.Path(exists=True, dir_okay=False, path_type=Path)) +def compare_cmd(output_path: Path) -> None: + """Render OUTPUT_PATH as a candidate-by-metric table, sorted by effective_tok_s.""" + output = BenchmarkOutput.model_validate_json(output_path.read_text(encoding="utf-8")) + sorted_candidates = sorted(output.candidates, key=lambda c: c.effective_tok_s, reverse=True) + table = Table(title=f"BenchmarkOutput ({output.corpus_run_id}, n={output.corpus_size})") + for col in ("candidate", "eff tok/s", "raw tok/s", "accept", "ttft p50 ms", "ttft p99 ms", "peak vram GiB", "startup s", "ok/tried"): + table.add_column(col, justify="right" if col != "candidate" else "left", overflow="fold") + for m in sorted_candidates: + table.add_row( + m.name, + f"{m.effective_tok_s:.1f}", + f"{m.raw_tok_s:.1f}", + f"{m.acceptance_rate:.3f}", + f"{m.ttft_p50_ms:.1f}", + f"{m.ttft_p99_ms:.1f}", + f"{m.observability.peak_vram_gb:.2f}" if m.observability.peak_vram_gb is not None else "—", + f"{m.startup_seconds:.1f}", + f"{m.prompts_accepted}/{m.prompts_attempted}", + ) + console.print(table) + if output.skipped_candidates: + skip_table = Table(title="Skipped") + skip_table.add_column("candidate", overflow="fold") + skip_table.add_column("error class") + skip_table.add_column("error", overflow="fold") + for skip in output.skipped_candidates: + skip_table.add_row(skip.name, skip.error_class, skip.error) + console.print(skip_table) + + +if __name__ == "__main__": + cli() From 53569e41096f297a6c6d20a902ef9525a507adcd Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Wed, 27 May 2026 14:22:46 +0000 Subject: [PATCH 07/17] feat(benchmark): bracketed_ab preset + condition_label/bracket_position MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the bracketed-A/B methodology presets that interleave baselines with candidates for drift detection, plus the supporting schema fields on ``BenchmarkCandidate`` and ``CandidateMetrics``. ## Schema additions - ``BenchmarkCandidate.condition_label: str = ''`` — sweep-level condition (``'baseline'`` / ``'n_fanout'`` / ``'spec_ngram'`` / ``'fp8'``). Set by the preset. - ``BenchmarkCandidate.bracket_position: int = 0`` — sequence index within a bracketed_ab stream. - ``CandidateMetrics.condition_label`` + ``bracket_position`` — copied from the candidate by the runner so the analyzer can read them directly off CandidateMetrics without joining back to the BenchmarkCandidate by name. ## bracketed_ab function ``bracketed_ab(base, *, candidate_engine_overrides=..., candidate_sampling_overrides=..., candidate_batch_dispatch_mode=..., condition_label, n_samples_per_condition=DEFAULT_BRACKETED_AB_N)`` returns 2N cells: N baselines interleaved with N candidate cells. Both pin ``SamplingParams.seed=DEFAULT_BENCHMARK_SEED``. N=6 default. The 2026-05-26 seed-pin verification established that pinning collapses the per-request RNG portion of variance but ~5% pooled CoV remains as structural non-RNG variance — that's what the cluster-conditioned analyzer (next commits) handles. ## Matrix-condition wrappers Registered in PRESETS for CLI consumption: - ``bracketed_ab_baseline_pool`` — N baselines (the shared pool). - ``bracketed_ab_n_fanout`` — uses ``batch_dispatch_mode='n_fanout'``. - ``bracketed_ab_spec_ngram`` — uses ``speculative_config={ 'method':'ngram', 'num_speculative_tokens':4, 'prompt_lookup_max':4 }``. - ``bracketed_ab_fp8`` — uses ``kv_cache_dtype='fp8'``. CLI invocation: ``--candidates bracketed_ab_spec_ngram`` etc. Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark.py | 35 +++++ .../generation/vllm_benchmark_presets.py | 129 ++++++++++++++++++ 2 files changed, 164 insertions(+) diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index c2315a4f8..0707e34a5 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -308,6 +308,26 @@ class BenchmarkCandidate(BaseModel): default="replicate", description="See :data:`BatchDispatchMode` module-level doc.", ) + condition_label: str = Field( + default="", + description=( + "Sweep-level condition this candidate measures, e.g. " + "``'baseline'``, ``'n_fanout'``, ``'spec_ngram'``, ``'fp8'``. " + "Set by the ``bracketed_ab`` preset family. Used by the " + "cluster-conditioned analyzer to group cells by condition " + "regardless of per-cell name suffixes." + ), + ) + bracket_position: int = Field( + default=0, + ge=0, + description=( + "Sequence index within a ``bracketed_ab`` cell stream " + "(baseline_0=0, candidate_0=1, baseline_1=2, candidate_1=3, " + "etc.). Used by the analyzer to align candidate cells with " + "their bracketing baselines for drift detection." + ), + ) class CandidateMetrics(BaseModel): @@ -369,6 +389,19 @@ class CandidateMetrics(BaseModel): description="Counts of vLLM finish reasons (``stop``, ``length``, etc.) across replays.", ) + # Sweep grouping — copied from BenchmarkCandidate so the analyzer + # can read them off CandidateMetrics without re-joining to the + # BenchmarkCandidate by name. + condition_label: str = Field( + default="", + description="Copied from :attr:`BenchmarkCandidate.condition_label`.", + ) + bracket_position: int = Field( + default=0, + ge=0, + description="Copied from :attr:`BenchmarkCandidate.bracket_position`.", + ) + # Composed observability (PR-A's schema). The runner builds this from # ``NvmlPeakSampler.peak_gb`` + ``read_loadavg`` pre/post + # ``probe_engine_runtime_config`` + ``read_vllm_runtime_metrics``, @@ -715,6 +748,8 @@ def run_benchmark( simulate_training_overlap_seconds=overlap, startup_overlap_savings_seconds=startup_overlap_savings_seconds, finish_reason_distribution=finish_reasons, + condition_label=candidate.condition_label, + bracket_position=candidate.bracket_position, observability=observability, ) diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py index 066564cb9..91b6581d4 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py @@ -185,6 +185,131 @@ def default_matrix(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: return out +DEFAULT_BRACKETED_AB_N: int = 6 +"""Cells per condition in :func:`bracketed_ab`. + +N=6 is the methodology-critic-recommended compromise between N=4 (~20% +statistical power for detecting +5% effects under the stack's observed +~3% in-cluster CoV) and N=8 (busts the 24h × $350 envelope on a +4-dataset × 4-condition matrix). +""" + + +def bracketed_ab( + base: BenchmarkEngineConfig, + *, + candidate_engine_overrides: dict[str, Any] | None = None, + candidate_sampling_overrides: dict[str, Any] | None = None, + candidate_batch_dispatch_mode: str | None = None, + condition_label: str, + n_samples_per_condition: int = DEFAULT_BRACKETED_AB_N, +) -> list[BenchmarkCandidate]: + """Emit an interleaved baseline-candidate cell sequence for bracketed A/B. + + Returns ``2 * n_samples_per_condition`` cells: ``n`` baselines + interleaved with ``n`` candidate cells. ``bracket_position`` is set + on each cell so the cluster-conditioned analyzer can align + candidate cells with their bracketing baselines for drift detection. + + Both baselines and candidates pin + ``SamplingParams.seed=DEFAULT_BENCHMARK_SEED``. The seed-pin + verification (2026-05-26) showed this does NOT collapse acceptance + variance to <0.5% CoV on long-output workloads but DOES eliminate + the per-request RNG portion — residual ~5% pooled CoV is structural + non-RNG variance that the cluster-conditioned analyzer partitions + out post-hoc. + """ + engine_overrides = candidate_engine_overrides or {} + sampling_extra = candidate_sampling_overrides or {} + cells: list[BenchmarkCandidate] = [] + for i in range(n_samples_per_condition): + cells.append( + BenchmarkCandidate( + name=f"bracket_baseline_{i}", + engine_config=_with_default_max_model_len(base), + sampling_overrides=_seeded_overrides(), + condition_label="baseline", + bracket_position=2 * i, + ), + ) + cand_kwargs: dict[str, Any] = { + "name": f"bracket_{condition_label}_{i}", + "engine_config": _with_default_max_model_len(base).model_copy(update=engine_overrides), + "sampling_overrides": _seeded_overrides(extra=sampling_extra), + "condition_label": condition_label, + "bracket_position": 2 * i + 1, + } + if candidate_batch_dispatch_mode is not None: + cand_kwargs["batch_dispatch_mode"] = candidate_batch_dispatch_mode + cells.append(BenchmarkCandidate(**cand_kwargs)) + return cells + + +# Phase B matrix-condition wrappers. Each yields a 2N-cell sequence +# (N baselines + N condition-specific candidates, interleaved). + + +def bracketed_ab_baseline_pool(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: + """N baseline cells with bracket_position labels — the shared baseline pool.""" + return [ + BenchmarkCandidate( + name=f"bracket_baseline_pool_{i}", + engine_config=_with_default_max_model_len(base), + sampling_overrides=_seeded_overrides(), + condition_label="baseline", + bracket_position=i, + ) + for i in range(DEFAULT_BRACKETED_AB_N) + ] + + +def bracketed_ab_n_fanout(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: + """N baselines + N n_fanout candidates, interleaved. + + n_fanout uses vLLM's ``SamplingParams.n=N`` so a single prompt + forks the decode state across N samples sharing the prefill KV + cache. Speculative win only when num_prompts > max_num_seqs. + """ + return bracketed_ab( + base, + candidate_batch_dispatch_mode="n_fanout", + condition_label="n_fanout", + ) + + +def bracketed_ab_spec_ngram(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: + """N baselines + N spec_ngram candidates, interleaved. + + Speculative-decoding overlay: ``speculative_config={'method': 'ngram', + 'num_speculative_tokens': 4, 'prompt_lookup_max': 4}``. Prior triage + measured +8.1% effective throughput on Mistral-7B + tabular. + """ + return bracketed_ab( + base, + candidate_engine_overrides={ + "speculative_config": { + "method": "ngram", + "num_speculative_tokens": 4, + "prompt_lookup_max": 4, + }, + }, + condition_label="spec_ngram", + ) + + +def bracketed_ab_fp8(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: + """N baselines + N fp8 KV-cache candidates, interleaved. + + Forces ``kv_cache_dtype='fp8'`` to halve KV cache memory footprint + at a per-token quality cost vLLM characterises as small. + """ + return bracketed_ab( + base, + candidate_engine_overrides={"kv_cache_dtype": "fp8"}, + condition_label="fp8", + ) + + # Map of preset name → callable used by the CLI to resolve ``--candidates``. PRESETS: dict[str, object] = { "baseline": lambda base: [baseline(base)], @@ -194,4 +319,8 @@ def default_matrix(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: "structured_backend_sweep": structured_backend_sweep, "max_model_len_sweep": max_model_len_sweep, "default_matrix": default_matrix, + "bracketed_ab_baseline_pool": bracketed_ab_baseline_pool, + "bracketed_ab_n_fanout": bracketed_ab_n_fanout, + "bracketed_ab_spec_ngram": bracketed_ab_spec_ngram, + "bracketed_ab_fp8": bracketed_ab_fp8, } From 9e4fe73f6dadc5da79934465101a040d7b3f2392 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Wed, 27 May 2026 14:24:42 +0000 Subject: [PATCH 08/17] feat(benchmark): cluster_conditioned analyzer with effect-size + 95% CI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds ``vllm_benchmark_analysis.py`` — the post-hoc analyzer that partitions cells by cluster signal, computes per-condition + per-cluster aggregates, and emits Welch-CI effect sizes. Combined #7+#8 from the original brief because they're tightly coupled (effect size lives on the same condition aggregate as the cluster stats). ## Pipeline ``analyze(output_dir, cluster_signal='auto', min_cells_per_condition=6) -> AnalysisReport``: 1. Load all BenchmarkOutput JSONs in dir, flatten candidate cells. 2. Auto-pick cluster signal (or use the operator's choice). 3. Silhouette-scored ``k`` in range [2, 4]; ``random_state=42``. 4. Remap labels so cluster 0 has the lowest signal mean (stable ordering). 5. Per-condition: pooled mean/stddev/CoV on effective_tok_s + acceptance_rate. 6. Per-(condition × cluster): mean/stddev/CoV within cluster. 7. Effect size: Welch's-t Δ ± 95% CI vs the 'baseline' condition, both pooled and per-cluster. 8. Refuse aggregates for conditions with <6 cells (brief mandate). ## Schema Pydantic models: ``AnalysisReport``, ``ConditionAggregate``, ``ConditionClusterAggregate``, ``ClusterStats``, ``ClusterAssignment``, ``EffectSize``. Each ``extra='forbid'``. ``AnalysisReport.to_markdown_summary()`` renders a human-readable table-formatted summary. ## CLI ``tools/vllm_benchmark.py analyze [--cluster-signal auto|wall_seconds|acceptance_rate] [--min-cells-per-condition N] [--json-out PATH]`` — runs the pipeline + prints markdown + optionally writes the full report JSON. ## Welch CI math ``_welch_ttest_ci(cand, base, alpha=0.05) -> (mean_diff, ci_low, ci_high, welch_df) | None``. Welch–Satterthwaite degrees of freedom computed explicitly. Returns None for underdetermined (n<2 either side or both stddevs zero). Validated synthetically: clear-difference inputs (cand mean ~1708, base mean ~1502, n=6 each) → Δ=+205, CI=[+193, +216], df=7.3, CI excludes 0; identical inputs → Δ=0, CI brackets 0 symmetrically. Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark_analysis.py | 521 ++++++++++++++++++ tools/vllm_benchmark.py | 46 ++ 2 files changed, 567 insertions(+) create mode 100644 src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py new file mode 100644 index 000000000..674e0a1be --- /dev/null +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py @@ -0,0 +1,521 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cluster-conditioned analysis for benchmark output. + +This stack does not produce trustworthy single-cell measurements. +Pooled cross-cell CoV runs ~5-10%; in-cluster CoV is ~3%. This module +partitions cells by a per-dataset cluster signal, then reports +per-condition aggregates both pooled and within-cluster, so the +operator can see which clusters contained the candidate's samples +and trust the in-cluster Δ rather than the noisier pooled Δ. + +Cluster signal per workload shape: + +- ``wall_seconds`` for short-context workloads (bike_sales-shape) — + bimodality is load-driven; partitioning on wall_seconds separates + the fast vs normal-load clusters. +- ``acceptance_rate`` for long-output workloads (call_transcripts- + shape) — bimodality is RNG/scheduler driven; partitioning on + acceptance_rate separates the high vs low cluster the seed-pin + validation found persists even with seed=42. +- ``auto`` — picks whichever signal has higher pooled CoV. + +Cluster count is selected via silhouette score (post-hoc, k ∈ [2, 4]). +Refuses to compute Δ-style aggregates when a condition has fewer than +:data:`MIN_CELLS_PER_CONDITION` cells — single-cell measurements +should never drive promote/reject decisions on this stack. + +Effect-size + 95% CI reporting lives in this module too — see +:class:`EffectSize`, computed via Welch's t-test on the difference of +means + Welch–Satterthwaite degrees of freedom. +""" + +from __future__ import annotations + +import json +import statistics +from pathlib import Path +from typing import Literal + +import numpy as np +from pydantic import BaseModel, ConfigDict, Field + +from .vllm_benchmark import BenchmarkOutput, CandidateMetrics + +ClusterSignal = Literal["wall_seconds", "acceptance_rate", "auto"] + +MIN_CELLS_PER_CONDITION: int = 6 +"""Minimum cells per condition before Δ-style aggregates are computed. + +Matches :data:`DEFAULT_BRACKETED_AB_N`. Below this threshold the +analyzer records a refusal in the report rather than emitting +under-powered aggregates that could mislead the operator. +""" + +_DEFAULT_K_MAX: int = 4 +"""Upper bound on the silhouette-score sweep for ``n_clusters`` selection.""" + + +# --------------------------------------------------------------------------- +# Report schema +# --------------------------------------------------------------------------- + + +class ClusterAssignment(BaseModel): + """One cell's assignment to a cluster + its raw signal value.""" + + model_config = ConfigDict(extra="forbid") + + candidate_name: str + condition_label: str + bracket_position: int + cluster_id: int + signal_value: float + + +class ClusterStats(BaseModel): + """Summary statistics for one cluster across all conditions.""" + + model_config = ConfigDict(extra="forbid") + + cluster_id: int + n_cells: int + signal_mean: float + signal_stddev: float + signal_cov: float + + +class EffectSize(BaseModel): + """Welch's-t Δ ± 95% CI for one condition vs a baseline reference. + + ``cluster_id=None`` is a pooled effect across all clusters; + integer values mean within-cluster (the brief-mandated headline + comparison). + """ + + model_config = ConfigDict(extra="forbid") + + metric: str + condition_label: str + baseline_condition_label: str + cluster_id: int | None + n_candidate: int + n_baseline: int + candidate_mean: float + baseline_mean: float + delta_absolute: float + delta_pct: float + ci95_low: float + ci95_high: float + welch_df: float + + +class ConditionClusterAggregate(BaseModel): + """Per-(condition × cluster) aggregate: in-cluster condition stats.""" + + model_config = ConfigDict(extra="forbid") + + condition_label: str + cluster_id: int + n_cells: int + mean_effective_tok_s: float + stddev_effective_tok_s: float + cov_effective_tok_s: float + mean_acceptance_rate: float + mean_raw_tok_s: float + effect_size_vs_baseline: EffectSize | None = None + + +class ConditionAggregate(BaseModel): + """Per-condition aggregate: pooled + in-cluster breakdowns.""" + + model_config = ConfigDict(extra="forbid") + + condition_label: str + n_cells: int + pooled_mean_effective_tok_s: float + pooled_stddev_effective_tok_s: float + pooled_cov_effective_tok_s: float + pooled_mean_acceptance_rate: float + pooled_stddev_acceptance_rate: float + pooled_cov_acceptance_rate: float + in_cluster: list[ConditionClusterAggregate] + pooled_effect_size_vs_baseline: EffectSize | None = None + + +class AnalysisReport(BaseModel): + """Top-level analysis report ready for serialization or rendering.""" + + model_config = ConfigDict(extra="forbid") + + cluster_signal: str + n_clusters: int + n_cells: int + cluster_assignments: list[ClusterAssignment] + cluster_stats: list[ClusterStats] + condition_aggregates: list[ConditionAggregate] + refusals: list[str] = Field(default_factory=list) + + def to_markdown_summary(self) -> str: + """Render a human-readable summary.""" + lines: list[str] = [ + f"# Cluster-conditioned analysis ({self.n_cells} cells, k={self.n_clusters})", + "", + f"**Cluster signal**: `{self.cluster_signal}`", + "", + "## Clusters", + "", + "| cluster | n_cells | signal_mean | signal_stddev | signal_cov |", + "|--------:|--------:|------------:|--------------:|-----------:|", + ] + for cs in self.cluster_stats: + lines.append( + f"| {cs.cluster_id} | {cs.n_cells} | {cs.signal_mean:.4f} | " + f"{cs.signal_stddev:.4f} | {cs.signal_cov * 100:.2f}% |" + ) + lines.extend(("", "## Per-condition aggregates", "")) + for agg in self.condition_aggregates: + lines.append(f"### `{agg.condition_label}` (n={agg.n_cells})") + lines.append("") + lines.append( + f"- **Pooled**: eff_tok_s={agg.pooled_mean_effective_tok_s:.1f} " + f"± {agg.pooled_stddev_effective_tok_s:.1f} " + f"(CoV {agg.pooled_cov_effective_tok_s * 100:.2f}%); " + f"accept={agg.pooled_mean_acceptance_rate:.4f} " + f"(CoV {agg.pooled_cov_acceptance_rate * 100:.2f}%)" + ) + if agg.pooled_effect_size_vs_baseline is not None: + es = agg.pooled_effect_size_vs_baseline + lines.append( + f" - Δ vs baseline (pooled): {es.delta_absolute:+.1f} tok/s " + f"({es.delta_pct:+.2f}%) [95% CI: {es.ci95_low:+.1f}, {es.ci95_high:+.1f}; " + f"n={es.n_candidate}+{es.n_baseline}, Welch df={es.welch_df:.1f}]" + ) + if agg.in_cluster: + lines.append("- **In-cluster**:") + for ic in agg.in_cluster: + lines.append( + f" - cluster {ic.cluster_id}: n={ic.n_cells}, " + f"eff_tok_s={ic.mean_effective_tok_s:.1f} " + f"± {ic.stddev_effective_tok_s:.1f} " + f"(CoV {ic.cov_effective_tok_s * 100:.2f}%); " + f"accept={ic.mean_acceptance_rate:.4f}" + ) + if ic.effect_size_vs_baseline is not None: + es = ic.effect_size_vs_baseline + lines.append( + f" - Δ vs baseline cluster {es.cluster_id}: " + f"{es.delta_absolute:+.1f} tok/s ({es.delta_pct:+.2f}%) " + f"[95% CI: {es.ci95_low:+.1f}, {es.ci95_high:+.1f}; " + f"n={es.n_candidate}+{es.n_baseline}, Welch df={es.welch_df:.1f}]" + ) + lines.append("") + if self.refusals: + lines.extend(("## Refusals (insufficient sample size)", "")) + for r in self.refusals: + lines.append(f"- {r}") + lines.append("") + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _signal_value(cell: CandidateMetrics, signal: str) -> float: + """Extract the signal value for clustering; raises on unknown signal.""" + if signal == "wall_seconds": + return cell.total_wall_seconds + if signal == "acceptance_rate": + return cell.acceptance_rate + raise ValueError(f"unknown cluster signal: {signal!r}") + + +def _pooled_cov(values: list[float]) -> tuple[float, float, float]: + """Return ``(mean, stddev, cov)`` for a list; CoV is 0 when mean=0.""" + if not values: + return 0.0, 0.0, 0.0 + mean = statistics.fmean(values) + stddev = statistics.stdev(values) if len(values) > 1 else 0.0 + cov = (stddev / mean) if mean > 0 else 0.0 + return mean, stddev, cov + + +def _auto_select_signal(cells: list[CandidateMetrics]) -> str: + """Pick whichever signal has higher pooled CoV across the cells. + + Defaults to ``wall_seconds`` on ties (short-context bimodality is + the more common workload shape). + """ + if not cells: + return "wall_seconds" + _, _, cov_wall = _pooled_cov([c.total_wall_seconds for c in cells]) + _, _, cov_acc = _pooled_cov([c.acceptance_rate for c in cells]) + return "acceptance_rate" if cov_acc > cov_wall else "wall_seconds" + + +def _select_n_clusters(values: np.ndarray, k_max: int = _DEFAULT_K_MAX) -> int: + """Silhouette-score-best ``k`` in range ``[2, min(k_max, n-1)]``. + + Returns 1 when there are too few cells (<4) to cluster meaningfully. + """ + from sklearn.cluster import KMeans # noqa: PLC0415 + from sklearn.metrics import silhouette_score # noqa: PLC0415 + + n = len(values) + if n < 4: + return 1 + X = values.reshape(-1, 1) + best_k = 2 + best_score = -2.0 + for k in range(2, min(k_max, n - 1) + 1): + km = KMeans(n_clusters=k, n_init=10, random_state=42) + labels = km.fit_predict(X) + if len(set(labels.tolist())) < 2: + continue + score = float(silhouette_score(X, labels)) + if score > best_score: + best_score = score + best_k = k + return best_k + + +def _assign_clusters(values: np.ndarray, n_clusters: int) -> np.ndarray: + """KMeans-assign labels, then remap so cluster 0 has the lowest mean.""" + if n_clusters <= 1: + return np.zeros(len(values), dtype=int) + from sklearn.cluster import KMeans # noqa: PLC0415 + + X = values.reshape(-1, 1) + km = KMeans(n_clusters=n_clusters, n_init=10, random_state=42) + raw_labels = km.fit_predict(X) + centers = [(float(km.cluster_centers_[i][0]), i) for i in range(n_clusters)] + centers.sort() + remap = {orig: new for new, (_, orig) in enumerate(centers)} + return np.array([remap[int(lbl)] for lbl in raw_labels], dtype=int) + + +def _welch_ttest_ci( + candidate_values: list[float], + baseline_values: list[float], + alpha: float = 0.05, +) -> tuple[float, float, float, float] | None: + """Return ``(mean_diff, ci_low, ci_high, welch_df)`` or ``None`` when underdetermined. + + Uses Welch's unequal-variance t-test for the CI on the difference + of means. ``None`` when either input has fewer than 2 observations + or both stddevs are zero. + """ + if len(candidate_values) < 2 or len(baseline_values) < 2: + return None + from scipy import stats # noqa: PLC0415 + + cand = np.array(candidate_values, dtype=float) + base = np.array(baseline_values, dtype=float) + cand_var = float(np.var(cand, ddof=1)) + base_var = float(np.var(base, ddof=1)) + cand_n = len(cand) + base_n = len(base) + mean_diff = float(np.mean(cand) - np.mean(base)) + se_diff_sq = cand_var / cand_n + base_var / base_n + if se_diff_sq <= 0: + return None + se_diff = float(np.sqrt(se_diff_sq)) + df_num = (cand_var / cand_n + base_var / base_n) ** 2 + df_den = (cand_var / cand_n) ** 2 / max(cand_n - 1, 1) + (base_var / base_n) ** 2 / max(base_n - 1, 1) + if df_den <= 0: + return None + welch_df = float(df_num / df_den) + t_crit = float(stats.t.ppf(1.0 - alpha / 2.0, welch_df)) + half_width = t_crit * se_diff + return mean_diff, mean_diff - half_width, mean_diff + half_width, welch_df + + +def _effect_size( + candidate_values: list[float], + baseline_values: list[float], + *, + metric: str, + condition_label: str, + baseline_condition_label: str, + cluster_id: int | None, +) -> EffectSize | None: + """Compute an :class:`EffectSize` or return ``None`` when underdetermined.""" + res = _welch_ttest_ci(candidate_values, baseline_values) + if res is None: + return None + mean_diff, ci_low, ci_high, welch_df = res + baseline_mean = float(np.mean(baseline_values)) + candidate_mean = float(np.mean(candidate_values)) + delta_pct = (mean_diff / baseline_mean * 100.0) if baseline_mean != 0 else 0.0 + return EffectSize( + metric=metric, + condition_label=condition_label, + baseline_condition_label=baseline_condition_label, + cluster_id=cluster_id, + n_candidate=len(candidate_values), + n_baseline=len(baseline_values), + candidate_mean=candidate_mean, + baseline_mean=baseline_mean, + delta_absolute=mean_diff, + delta_pct=delta_pct, + ci95_low=ci_low, + ci95_high=ci_high, + welch_df=welch_df, + ) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def load_cells(output_dir: Path) -> list[CandidateMetrics]: + """Read every ``BenchmarkOutput`` JSON in ``output_dir``, flatten candidates. + + Subdirectories are NOT recursed. Callers wanting cross-dataset + analysis should invoke once per dataset dir. + """ + cells: list[CandidateMetrics] = [] + for path in sorted(output_dir.glob("*.json")): + try: + doc = BenchmarkOutput.model_validate_json(path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, ValueError): + continue + cells.extend(doc.candidates) + return cells + + +def analyze( + output_dir: Path, + cluster_signal: ClusterSignal = "auto", + min_cells_per_condition: int = MIN_CELLS_PER_CONDITION, +) -> AnalysisReport: + """Full pipeline: load → cluster → per-condition aggregate → effect-size → report.""" + cells = load_cells(output_dir) + if not cells: + raise ValueError(f"No BenchmarkOutput JSONs found under {output_dir}") + + resolved_signal = _auto_select_signal(cells) if cluster_signal == "auto" else cluster_signal + values = np.array([_signal_value(c, resolved_signal) for c in cells], dtype=float) + n_clusters = _select_n_clusters(values) + labels = _assign_clusters(values, n_clusters) + + assignments = [ + ClusterAssignment( + candidate_name=cell.name, + condition_label=cell.condition_label, + bracket_position=cell.bracket_position, + cluster_id=int(lbl), + signal_value=float(val), + ) + for cell, lbl, val in zip(cells, labels, values, strict=True) + ] + + cluster_stats: list[ClusterStats] = [] + for cid in range(n_clusters): + cluster_values = [float(values[i]) for i, lbl in enumerate(labels) if int(lbl) == cid] + mean, stddev, cov = _pooled_cov(cluster_values) + cluster_stats.append( + ClusterStats( + cluster_id=cid, + n_cells=len(cluster_values), + signal_mean=mean, + signal_stddev=stddev, + signal_cov=cov, + ), + ) + + by_condition: dict[str, list[tuple[int, CandidateMetrics]]] = {} + for cell, lbl in zip(cells, labels, strict=True): + by_condition.setdefault(cell.condition_label, []).append((int(lbl), cell)) + + baseline_pooled_eff: list[float] = [] + baseline_per_cluster_eff: dict[int, list[float]] = {} + for lbl, cell in by_condition.get("baseline", []): + baseline_pooled_eff.append(cell.effective_tok_s) + baseline_per_cluster_eff.setdefault(lbl, []).append(cell.effective_tok_s) + + aggregates: list[ConditionAggregate] = [] + refusals: list[str] = [] + for condition in sorted(by_condition): + labeled = by_condition[condition] + if len(labeled) < min_cells_per_condition: + refusals.append( + f"condition {condition!r} has only {len(labeled)} cells; " + f"need ≥{min_cells_per_condition} — refusing aggregate" + ) + continue + pooled_eff = [c.effective_tok_s for _, c in labeled] + pooled_acc = [c.acceptance_rate for _, c in labeled] + eff_mean, eff_stddev, eff_cov = _pooled_cov(pooled_eff) + acc_mean, acc_stddev, acc_cov = _pooled_cov(pooled_acc) + + in_cluster: list[ConditionClusterAggregate] = [] + for cid in range(n_clusters): + cluster_cells = [c for lbl, c in labeled if lbl == cid] + if not cluster_cells: + continue + ic_eff = [c.effective_tok_s for c in cluster_cells] + ic_acc = [c.acceptance_rate for c in cluster_cells] + ic_raw = [c.raw_tok_s for c in cluster_cells] + mean_eff, stddev_eff, cov_eff = _pooled_cov(ic_eff) + ic_effect: EffectSize | None = None + if condition != "baseline" and cid in baseline_per_cluster_eff: + ic_effect = _effect_size( + ic_eff, + baseline_per_cluster_eff[cid], + metric="effective_tok_s", + condition_label=condition, + baseline_condition_label="baseline", + cluster_id=cid, + ) + in_cluster.append( + ConditionClusterAggregate( + condition_label=condition, + cluster_id=cid, + n_cells=len(cluster_cells), + mean_effective_tok_s=mean_eff, + stddev_effective_tok_s=stddev_eff, + cov_effective_tok_s=cov_eff, + mean_acceptance_rate=statistics.fmean(ic_acc), + mean_raw_tok_s=statistics.fmean(ic_raw), + effect_size_vs_baseline=ic_effect, + ), + ) + pooled_effect: EffectSize | None = None + if condition != "baseline" and baseline_pooled_eff: + pooled_effect = _effect_size( + pooled_eff, + baseline_pooled_eff, + metric="effective_tok_s", + condition_label=condition, + baseline_condition_label="baseline", + cluster_id=None, + ) + aggregates.append( + ConditionAggregate( + condition_label=condition, + n_cells=len(labeled), + pooled_mean_effective_tok_s=eff_mean, + pooled_stddev_effective_tok_s=eff_stddev, + pooled_cov_effective_tok_s=eff_cov, + pooled_mean_acceptance_rate=acc_mean, + pooled_stddev_acceptance_rate=acc_stddev, + pooled_cov_acceptance_rate=acc_cov, + in_cluster=in_cluster, + pooled_effect_size_vs_baseline=pooled_effect, + ), + ) + + return AnalysisReport( + cluster_signal=resolved_signal, + n_clusters=n_clusters, + n_cells=len(cells), + cluster_assignments=assignments, + cluster_stats=cluster_stats, + condition_aggregates=aggregates, + refusals=refusals, + ) diff --git a/tools/vllm_benchmark.py b/tools/vllm_benchmark.py index c7547442a..0f6f1ca42 100644 --- a/tools/vllm_benchmark.py +++ b/tools/vllm_benchmark.py @@ -211,5 +211,51 @@ def compare_cmd(output_path: Path) -> None: console.print(skip_table) +@cli.command("analyze") +@click.argument("output_dir", type=click.Path(exists=True, file_okay=False, path_type=Path)) +@click.option( + "--cluster-signal", + type=click.Choice(["wall_seconds", "acceptance_rate", "auto"]), + default="auto", + show_default=True, + help="Which per-cell metric to partition cells on. Use 'wall_seconds' for short-context (load-driven bimodality), 'acceptance_rate' for long-output (RNG/scheduler driven), 'auto' to pick whichever has higher pooled CoV.", +) +@click.option( + "--min-cells-per-condition", + type=int, + default=None, + show_default="MIN_CELLS_PER_CONDITION (6)", + help="Refuse aggregates for conditions below this N. Brief mandates N≥6.", +) +@click.option( + "--json-out", + type=click.Path(dir_okay=False, path_type=Path), + default=None, + help="Optional path to write the full AnalysisReport as JSON.", +) +def analyze_cmd( + output_dir: Path, + cluster_signal: str, + min_cells_per_condition: int | None, + json_out: Path | None, +) -> None: + """Cluster-conditioned analysis across every BenchmarkOutput JSON in OUTPUT_DIR.""" + from nemo_safe_synthesizer.generation.vllm_benchmark_analysis import ( + MIN_CELLS_PER_CONDITION, + analyze, + ) + + report = analyze( + output_dir, + cluster_signal=cluster_signal, # ty: ignore[invalid-argument-type] + min_cells_per_condition=MIN_CELLS_PER_CONDITION if min_cells_per_condition is None else min_cells_per_condition, + ) + console.print(report.to_markdown_summary()) + if json_out is not None: + json_out.parent.mkdir(parents=True, exist_ok=True) + json_out.write_text(report.model_dump_json(indent=2), encoding="utf-8") + console.print(f"[green]wrote[/green] {json_out}") + + if __name__ == "__main__": cli() From 9ed068c3b2d253d2fea8f1d0c4f944bfc9167037 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Wed, 27 May 2026 14:27:44 +0000 Subject: [PATCH 09/17] test(generation): benchmark harness + analyzer contract tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 37 tests across two files (~55s wall). Covers consumer-facing contracts; skips runner end-to-end (requires vLLM spin-up). ## test_vllm_benchmark.py (29 tests) - **Schema contracts** — ``CandidateMetrics`` composition of ``CellObservability``, JSON round-trip (nested observability survives), ``BenchmarkEngineConfig.extra='ignore'`` tolerates capture-time kwargs, ``BenchmarkCandidate.extra='forbid'`` blocks silent drift. - **Corpus loader** — header + records JSONL parses correctly; missing header raises with a clear message. - **Helpers** — ``_percentile`` linear-interp + empty-list edge, ``_merge_sampling_kwargs`` strips capture-time-only metadata fields, overrides win on conflict, ``_extract_ttft_ms`` parametrized matrix (modern ``first_token_latency``, legacy ``first_token_time - arrival_time`` fallback, None on missing data), ``_truncate_stderr`` tail-preserving, ``_parse_error_class`` exception-name extraction. - **``_build_vllm_kwargs``** — header + candidate config overlay, attention_backend translation, None-valued field exclusion, attention='auto' treated as unset. - **Presets** — ``baseline`` pins seed, ``default_matrix`` dedupes by (engine_config, sampling_overrides), ``attention_backend_sweep`` covers known backends. - **bracketed_ab** — 2N interleaved cells with correct positions and labels, ``bracketed_ab_spec_ngram`` injects ``speculative_config`` on candidates only, all cells seed-pinned. - **SubprocessRunResult** — success + failure shapes. ## test_vllm_benchmark_analysis.py (8 tests) - **Welch CI math** — clear-difference CI excludes 0; identical-means CI brackets 0 symmetrically; underdetermined inputs (n<2 either side or both stddevs zero) return None. - **Full analyze pipeline** — synthetic 6+6 sweep (baselines around ~1500 tok/s, spec_ngram around ~1700) produces: - Correct n_cells + per-condition aggregates with pooled mean within tolerance. - Effect size emitted for spec_ngram, NOT for baseline (no self-comparison). - Refusal entries (not aggregates) for conditions with N<6. - AnalysisReport JSON round-trips losslessly. - Empty output dir raises with a clear message. ## DRY mechanisms - Module-level ``_metric`` helper builds CandidateMetrics with the necessary scaffolding; reused across all analyzer tests. - ``_write_output_dir`` writes a single BenchmarkOutput JSON to a tmp dir; used by the synthetic-sweep fixture. - Parametrized matrices for ``_extract_ttft_ms`` (4 cases) and ``_welch_ttest_ci`` underdetermined inputs (3 cases). - Shared fixtures: ``header`` (TraceHeader), ``empty_base`` (default BenchmarkEngineConfig), and ``synthetic_sweep_dir`` (12-cell baseline+spec_ngram fixture). Signed-off-by: Aaron Gonzales --- tests/generation/test_vllm_benchmark.py | 299 ++++++++++++++++++ .../test_vllm_benchmark_analysis.py | 154 +++++++++ 2 files changed, 453 insertions(+) create mode 100644 tests/generation/test_vllm_benchmark.py create mode 100644 tests/generation/test_vllm_benchmark_analysis.py diff --git a/tests/generation/test_vllm_benchmark.py b/tests/generation/test_vllm_benchmark.py new file mode 100644 index 000000000..040a0c782 --- /dev/null +++ b/tests/generation/test_vllm_benchmark.py @@ -0,0 +1,299 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for benchmark harness data models, helpers, and presets. + +Scope: consumer-facing contracts. Skip the runner itself (requires +spinning up vLLM); covered by the actual production cell smokes. +""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest +from pydantic import ValidationError + +from nemo_safe_synthesizer.generation.vllm_benchmark import ( + BenchmarkCandidate, + BenchmarkCorpus, + BenchmarkEngineConfig, + BenchmarkOutput, + BenchmarkPrompt, + CandidateMetrics, + SubprocessRunResult, + TraceHeader, + _build_vllm_kwargs, + _extract_ttft_ms, + _merge_sampling_kwargs, + _parse_error_class, + _percentile, + _truncate_stderr, +) +from nemo_safe_synthesizer.generation.vllm_benchmark_presets import ( + DEFAULT_BENCHMARK_SEED, + DEFAULT_BRACKETED_AB_N, + PRESETS, + attention_backend_sweep, + baseline, + bracketed_ab, + bracketed_ab_spec_ngram, + default_matrix, +) +from nemo_safe_synthesizer.generation.vllm_observability import CellObservability + + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def header() -> TraceHeader: + return TraceHeader( + run_id="test-run", + pretrained_model="mistralai/Mistral-7B-Instruct-v0.3", + dataset_schema={"col": "string"}, + engine_parameters={"max_lora_rank": 32, "structured_generation_backend": "outlines"}, + ) + + +@pytest.fixture +def empty_base() -> BenchmarkEngineConfig: + return BenchmarkEngineConfig() + + +# --------------------------------------------------------------------------- +# Data model contracts +# --------------------------------------------------------------------------- + + +class TestSchemaContracts: + def test_candidate_metrics_composes_observability(self) -> None: + """``CandidateMetrics`` embeds ``CellObservability`` rather than re-declaring its fields.""" + m = CandidateMetrics( + name="t", raw_tok_s=1.0, acceptance_rate=0.9, effective_tok_s=0.9, + ttft_p50_ms=0.0, ttft_p99_ms=0.0, + prompts_attempted=10, prompts_accepted=9, total_output_tokens=100, + total_wall_seconds=1.0, + observability=CellObservability(peak_vram_gb=64.5), + ) + assert m.observability.peak_vram_gb == 64.5 + + def test_benchmark_output_roundtrips_through_json(self) -> None: + """Consumers depend on the JSON serialization being lossless.""" + out = BenchmarkOutput( + corpus_run_id="r1", + corpus_size=10, + candidates=[ + CandidateMetrics( + name="t", raw_tok_s=1.0, acceptance_rate=0.9, effective_tok_s=0.9, + ttft_p50_ms=0.0, ttft_p99_ms=0.0, + prompts_attempted=10, prompts_accepted=9, total_output_tokens=100, + total_wall_seconds=1.0, + observability=CellObservability(peak_vram_gb=64.5, loadavg_pre=(1.0, 2.0, 3.0)), + ), + ], + ) + rt = BenchmarkOutput.model_validate_json(out.model_dump_json()) + assert rt == out + # Specifically: the nested observability survives the round-trip. + assert rt.candidates[0].observability.peak_vram_gb == 64.5 + assert rt.candidates[0].observability.loadavg_pre == (1.0, 2.0, 3.0) + + def test_engine_config_tolerates_unknown_kwargs(self) -> None: + """``BenchmarkEngineConfig.extra='ignore'`` so the CLI can validate raw header dicts.""" + cfg = BenchmarkEngineConfig.model_validate({"attention_backend": "FLASHINFER", "unknown_kwarg": 42}) + assert cfg.attention_backend == "FLASHINFER" + + def test_candidate_extra_fields_forbidden(self) -> None: + """``BenchmarkCandidate.extra='forbid'`` — adding a field must update the schema.""" + with pytest.raises(ValidationError): + BenchmarkCandidate.model_validate({"name": "t", "unknown_field": 42}) + + +# --------------------------------------------------------------------------- +# Corpus loader +# --------------------------------------------------------------------------- + + +class TestBenchmarkCorpus: + def test_loads_jsonl_with_header_and_records(self, tmp_path: Any) -> None: + path = tmp_path / "trace.jsonl" + lines = [ + json.dumps({"kind": "header", "run_id": "r", "pretrained_model": "m", "dataset_schema": {}}), + json.dumps({"kind": "record", "row_index": 0, "prompt": "p0", "sampling_params": {"temperature": 0.7}}), + json.dumps({"kind": "record", "row_index": 1, "prompt": "p1", "sampling_params": {"temperature": 0.7}}), + ] + path.write_text("\n".join(lines), encoding="utf-8") + corpus = BenchmarkCorpus.from_trace_jsonl(path) + assert corpus.header.run_id == "r" + assert len(corpus.prompts) == 2 + assert corpus.prompts[0].original_sampling_params["temperature"] == 0.7 + + def test_rejects_missing_header(self, tmp_path: Any) -> None: + path = tmp_path / "no_header.jsonl" + path.write_text(json.dumps({"kind": "record", "row_index": 0, "prompt": "x"}), encoding="utf-8") + with pytest.raises(ValueError, match="record on line 1 before any header"): + BenchmarkCorpus.from_trace_jsonl(path) + + +# --------------------------------------------------------------------------- +# Helper contracts +# --------------------------------------------------------------------------- + + +class TestHelpers: + def test_percentile_linear_interp(self) -> None: + assert _percentile([1.0, 2.0, 3.0, 4.0], 50) == 2.5 + assert _percentile([1.0], 50) == 1.0 + assert _percentile([], 50) == 0.0 # empty returns 0, not raises + + def test_merge_sampling_strips_non_sampling_fields(self) -> None: + """``structured_outputs`` is capture-time metadata, not a SamplingParams kwarg.""" + merged = _merge_sampling_kwargs( + {"temperature": 0.7, "top_p": 0.9, "structured_outputs": "json"}, + {"seed": 42}, + ) + assert merged == {"temperature": 0.7, "top_p": 0.9, "seed": 42} + + def test_merge_overrides_win_on_conflict(self) -> None: + merged = _merge_sampling_kwargs({"temperature": 0.7}, {"temperature": 0.0}) + assert merged["temperature"] == 0.0 + + @pytest.mark.parametrize( + ("metrics_obj", "expected"), + [ + (None, None), + (type("M", (), {"first_token_latency": 0.12})(), 120.0), + (type("M", (), {"first_token_latency": None, "first_token_time": 1.0, "arrival_time": 0.5})(), 500.0), + (type("M", (), {"first_token_latency": None, "first_token_time": None, "arrival_time": 0.5})(), None), + ], + ) + def test_extract_ttft_ms(self, metrics_obj: Any, expected: float | None) -> None: + """Tries modern ``first_token_latency`` first, falls back to ``first_token_time - arrival_time``.""" + output = type("Out", (), {"metrics": metrics_obj})() + assert _extract_ttft_ms(output) == expected + + def test_truncate_stderr_keeps_tail(self) -> None: + long = "x" * 600 + truncated = _truncate_stderr(long, limit=100) + assert len(truncated) <= 100 + assert truncated.endswith("xxxxx") # tail-preserving + + def test_parse_error_class_finds_exception_name(self) -> None: + assert _parse_error_class("Traceback (most recent call last):\n...\nValueError: bad") == "ValueError" + assert _parse_error_class("") == "Error" + + +# --------------------------------------------------------------------------- +# Engine kwargs builder +# --------------------------------------------------------------------------- + + +class TestBuildVllmKwargs: + def test_overlays_engine_config_on_header(self, header: TraceHeader) -> None: + """Candidate engine_config overlays the header's engine_parameters.""" + cfg = BenchmarkEngineConfig(attention_backend="FLASHINFER", max_model_len=4096) + kwargs = _build_vllm_kwargs(header, cfg) + assert kwargs["model"] == "mistralai/Mistral-7B-Instruct-v0.3" + assert kwargs["max_lora_rank"] == 32 # from header + assert kwargs["max_model_len"] == 4096 # from cfg + + def test_translates_attention_backend_to_attention_config(self, header: TraceHeader, empty_base: BenchmarkEngineConfig) -> None: + """vLLM's public API takes an ``attention_config`` dict, not a bare backend string.""" + cfg = empty_base.model_copy(update={"attention_backend": "FLASHINFER"}) + kwargs = _build_vllm_kwargs(header, cfg) + assert kwargs["attention_config"] == {"backend": "FLASHINFER"} + assert "attention_backend" not in kwargs # translated, not forwarded + + def test_drops_none_valued_overrides(self, header: TraceHeader, empty_base: BenchmarkEngineConfig) -> None: + """Unset candidate fields stay out of kwargs so vLLM uses its defaults.""" + kwargs = _build_vllm_kwargs(header, empty_base) + assert "enable_prefix_caching" not in kwargs + assert "max_num_seqs" not in kwargs + + def test_auto_attention_backend_is_treated_as_unset(self, header: TraceHeader, empty_base: BenchmarkEngineConfig) -> None: + """``attention_backend='auto'`` means "let vLLM pick" — no attention_config kwarg should appear.""" + cfg = empty_base.model_copy(update={"attention_backend": "auto"}) + kwargs = _build_vllm_kwargs(header, cfg) + assert "attention_config" not in kwargs + + +# --------------------------------------------------------------------------- +# Preset contracts +# --------------------------------------------------------------------------- + + +class TestPresets: + def test_baseline_pins_seed(self, empty_base: BenchmarkEngineConfig) -> None: + cand = baseline(empty_base) + assert cand.sampling_overrides == {"seed": DEFAULT_BENCHMARK_SEED} + + def test_default_matrix_dedupes(self, empty_base: BenchmarkEngineConfig) -> None: + """Concatenated sweeps that produce identical (engine, sampling) tuples are collapsed.""" + cands = default_matrix(empty_base) + keys = {(c.engine_config.model_dump_json(), json.dumps(c.sampling_overrides, sort_keys=True)) for c in cands} + assert len(keys) == len(cands) + + def test_attention_backend_sweep_covers_known_backends(self, empty_base: BenchmarkEngineConfig) -> None: + cands = attention_backend_sweep(empty_base) + backends = {c.engine_config.attention_backend for c in cands} + assert {"FLASHINFER", "FLASH_ATTN", "TRITON_ATTN"}.issubset(backends) + + +class TestBracketedAb: + def test_emits_2n_cells_interleaved(self, empty_base: BenchmarkEngineConfig) -> None: + """N baselines + N candidate cells, interleaved by bracket_position.""" + cells = bracketed_ab( + empty_base, + candidate_engine_overrides={"enable_prefix_caching": True}, + condition_label="prefix_on", + n_samples_per_condition=3, + ) + assert len(cells) == 6 + # Even positions are baseline; odd positions are the candidate. + for i, cell in enumerate(cells): + expected_label = "baseline" if i % 2 == 0 else "prefix_on" + assert cell.condition_label == expected_label + assert cell.bracket_position == i + + def test_spec_ngram_wrapper_applies_speculative_config(self, empty_base: BenchmarkEngineConfig) -> None: + cells = bracketed_ab_spec_ngram(empty_base) + # Candidate cells have speculative_config; baselines don't. + candidates = [c for c in cells if c.condition_label == "spec_ngram"] + baselines = [c for c in cells if c.condition_label == "baseline"] + assert len(candidates) == DEFAULT_BRACKETED_AB_N + assert len(baselines) == DEFAULT_BRACKETED_AB_N + for c in candidates: + assert c.engine_config.speculative_config is not None + assert c.engine_config.speculative_config["method"] == "ngram" + for b in baselines: + assert b.engine_config.speculative_config is None + + def test_all_cells_seed_pinned(self, empty_base: BenchmarkEngineConfig) -> None: + cells = bracketed_ab_spec_ngram(empty_base) + for c in cells: + assert c.sampling_overrides.get("seed") == DEFAULT_BENCHMARK_SEED + + +# --------------------------------------------------------------------------- +# SubprocessRunResult +# --------------------------------------------------------------------------- + + +class TestSubprocessRunResult: + def test_success_shape(self) -> None: + m = CandidateMetrics( + name="t", raw_tok_s=1.0, acceptance_rate=0.9, effective_tok_s=0.9, + ttft_p50_ms=0.0, ttft_p99_ms=0.0, + prompts_attempted=10, prompts_accepted=9, total_output_tokens=100, + total_wall_seconds=1.0, + ) + r = SubprocessRunResult(metrics=m) + assert r.metrics is not None and r.error is None + + def test_failure_shape(self) -> None: + r = SubprocessRunResult(error="exit 1", error_class="RuntimeError") + assert r.metrics is None and r.error_class == "RuntimeError" diff --git a/tests/generation/test_vllm_benchmark_analysis.py b/tests/generation/test_vllm_benchmark_analysis.py new file mode 100644 index 000000000..357c8708f --- /dev/null +++ b/tests/generation/test_vllm_benchmark_analysis.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the cluster-conditioned analyzer. + +Scope: contracts consumers depend on — cluster partitioning produces +the expected shape, refusals fire when sample size is too small, +effect-size CI brackets behave correctly, JSON output round-trips. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest + +from nemo_safe_synthesizer.generation.vllm_benchmark import ( + BenchmarkOutput, + CandidateMetrics, +) +from nemo_safe_synthesizer.generation.vllm_benchmark_analysis import ( + MIN_CELLS_PER_CONDITION, + AnalysisReport, + _effect_size, + _welch_ttest_ci, + analyze, +) + + +def _metric(name: str, condition: str, *, eff: float, accept: float, wall: float, bracket: int = 0) -> CandidateMetrics: + """Build a CandidateMetrics with minimal scaffolding — used to seed the analyzer.""" + return CandidateMetrics( + name=name, + raw_tok_s=eff / max(accept, 0.01), + acceptance_rate=accept, + effective_tok_s=eff, + ttft_p50_ms=0.0, + ttft_p99_ms=0.0, + prompts_attempted=143, + prompts_accepted=int(143 * accept), + total_output_tokens=100000, + total_wall_seconds=wall, + condition_label=condition, + bracket_position=bracket, + ) + + +def _write_output_dir(tmp_path: Path, cells: list[CandidateMetrics]) -> Path: + """Write a single BenchmarkOutput JSON to ``tmp_path / out.json``.""" + out = BenchmarkOutput(corpus_run_id="r1", corpus_size=143, candidates=cells) + (tmp_path / "out.json").write_text(out.model_dump_json(), encoding="utf-8") + return tmp_path + + +@pytest.fixture +def synthetic_sweep_dir(tmp_path: Path) -> Path: + """6 baselines + 6 spec_ngram cells, both with realistic-noise spread.""" + cells = [ + # Baselines: ~1500 eff_tok_s, ~0.99 acceptance. + *(_metric(f"baseline_{i}", "baseline", eff=1500 + i * 5, accept=0.99, wall=130.0, bracket=2 * i) for i in range(6)), + # spec_ngram: ~1700 eff_tok_s (+~13%), similar acceptance. + *(_metric(f"spec_{i}", "spec_ngram", eff=1700 + i * 5, accept=0.99, wall=115.0, bracket=2 * i + 1) for i in range(6)), + ] + return _write_output_dir(tmp_path, cells) + + +# --------------------------------------------------------------------------- +# Welch CI math +# --------------------------------------------------------------------------- + + +class TestWelchTtestCi: + def test_clear_difference_excludes_zero(self) -> None: + """Genuine difference in means → CI doesn't bracket 0.""" + res = _welch_ttest_ci([1700, 1720, 1690, 1710, 1705, 1715], [1500, 1510, 1495, 1505, 1502, 1498]) + assert res is not None + mean_diff, ci_low, ci_high, _df = res + assert mean_diff > 0 + assert ci_low > 0 # CI excludes 0 + + def test_identical_means_brackets_zero(self) -> None: + """Same means + nonzero variance → CI brackets 0.""" + res = _welch_ttest_ci([1500, 1510, 1495, 1505, 1502, 1498], [1500, 1510, 1495, 1505, 1502, 1498]) + assert res is not None + mean_diff, ci_low, ci_high, _df = res + assert mean_diff == 0 + assert ci_low < 0 < ci_high + + @pytest.mark.parametrize( + ("cand", "base"), + [ + ([1500], [1502]), # too few candidate observations + ([1500, 1500], [1502]), # too few baseline observations + ([1500, 1500], [1500, 1500]), # both stddevs zero + ], + ) + def test_underdetermined_returns_none(self, cand: list[float], base: list[float]) -> None: + """Degraded mode: returns None instead of nan/inf/raising.""" + assert _welch_ttest_ci(cand, base) is None + + +# --------------------------------------------------------------------------- +# Full analyze pipeline +# --------------------------------------------------------------------------- + + +class TestAnalyze: + def test_partitions_and_aggregates(self, synthetic_sweep_dir: Path) -> None: + report = analyze(synthetic_sweep_dir, cluster_signal="wall_seconds") + assert report.n_cells == 12 + # Two conditions present. + labels = {agg.condition_label for agg in report.condition_aggregates} + assert labels == {"baseline", "spec_ngram"} + # Each condition has the expected pooled aggregate. + spec_agg = next(agg for agg in report.condition_aggregates if agg.condition_label == "spec_ngram") + assert spec_agg.n_cells == 6 + assert spec_agg.pooled_mean_effective_tok_s == pytest.approx(1712.5, abs=0.1) + + def test_emits_effect_size_for_non_baseline_conditions(self, synthetic_sweep_dir: Path) -> None: + """spec_ngram vs baseline gets an effect size; baseline itself doesn't (no self-comparison).""" + report = analyze(synthetic_sweep_dir) + spec_agg = next(agg for agg in report.condition_aggregates if agg.condition_label == "spec_ngram") + baseline_agg = next(agg for agg in report.condition_aggregates if agg.condition_label == "baseline") + assert spec_agg.pooled_effect_size_vs_baseline is not None + assert baseline_agg.pooled_effect_size_vs_baseline is None + # The Δ should be roughly +200 tok/s, CI excludes 0. + es = spec_agg.pooled_effect_size_vs_baseline + assert es.delta_absolute > 0 + assert es.ci95_low > 0 + + def test_refuses_aggregates_below_min_cells(self, tmp_path: Path) -> None: + """Conditions with N<6 land in refusals, not aggregates.""" + # 6 baselines + only 3 candidate cells → spec_ngram should be refused. + cells = [ + *(_metric(f"baseline_{i}", "baseline", eff=1500.0, accept=0.99, wall=130.0) for i in range(6)), + *(_metric(f"spec_{i}", "spec_ngram", eff=1700.0, accept=0.99, wall=115.0) for i in range(3)), + ] + out_dir = _write_output_dir(tmp_path, cells) + report = analyze(out_dir) + labels = {agg.condition_label for agg in report.condition_aggregates} + assert "baseline" in labels + assert "spec_ngram" not in labels + assert any("spec_ngram" in r for r in report.refusals) + + def test_report_round_trips_through_json(self, synthetic_sweep_dir: Path) -> None: + report = analyze(synthetic_sweep_dir) + rt = AnalysisReport.model_validate_json(report.model_dump_json()) + assert rt == report + + def test_load_failure_raises(self, tmp_path: Path) -> None: + """Empty output dir raises with a clear message.""" + with pytest.raises(ValueError, match="No BenchmarkOutput JSONs"): + analyze(tmp_path) From 7ef702ac95a0d9cdbf10321b49a195bab8e83075 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Wed, 27 May 2026 15:15:05 +0000 Subject: [PATCH 10/17] style+fix: apply ruff format + correct TabularDataProcessor API call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ``make format`` output across 8 files (mechanical: line breaks, import sort, blank lines). - Runtime fix: ``processor.process(text)`` was wrong — the actual API is ``Processor.__call__(prompt_number, text) -> ParsedResponse`` (per ``processors.py``). Fixed to ``processor(prompt_idx, text)`` with the correct property access on ``parsed.valid_records`` / ``parsed.invalid_records``. Both surfaced by ``make check`` (format-check + typecheck). Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark.py | 31 ++++++------ .../generation/vllm_benchmark_presets.py | 7 +-- .../generation/vllm_benchmark_wandb.py | 6 +-- tests/generation/test_vllm_benchmark.py | 49 +++++++++++++------ .../test_vllm_benchmark_analysis.py | 13 +++-- tools/vllm_benchmark.py | 14 +++++- 6 files changed, 76 insertions(+), 44 deletions(-) diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index 0707e34a5..21f40ae28 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -227,7 +227,9 @@ class BenchmarkEngineConfig(BaseModel): ), ) max_num_seqs: int | None = Field(default=None, description="vLLM scheduler ``max_num_seqs`` cap.") - max_num_batched_tokens: int | None = Field(default=None, description="vLLM scheduler ``max_num_batched_tokens`` cap.") + max_num_batched_tokens: int | None = Field( + default=None, description="vLLM scheduler ``max_num_batched_tokens`` cap." + ) enable_chunked_prefill: bool | None = Field(default=None, description="Chunked-prefill engagement.") kv_cache_dtype: str | None = Field( default=None, @@ -631,9 +633,7 @@ def run_benchmark( raise init_result["exception"] llm: LLM = init_result["llm"] - startup_overlap_savings_seconds = ( - min(overlap, startup_seconds + overlap) if overlap > 0.0 else 0.0 - ) + startup_overlap_savings_seconds = min(overlap, startup_seconds + overlap) if overlap > 0.0 else 0.0 # Probe the engine's effective runtime config + check for # candidate-intent / engine-actual disagreements. @@ -652,9 +652,7 @@ def run_benchmark( tokenizer=None, ) lora_request = ( - LoRARequest("lora", 1, str(corpus.header.lora_path)) - if corpus.header.lora_path is not None - else None + LoRARequest("lora", 1, str(corpus.header.lora_path)) if corpus.header.lora_path is not None else None ) # Build SamplingParams from corpus default + candidate overrides. @@ -670,9 +668,11 @@ def run_benchmark( # Dispatch. prompts: list[str] = [p.prompt for p in corpus.prompts] gen_start = time.perf_counter() - outputs: list[Any] = list( - llm.generate(prompts=prompts, sampling_params=sampling_params, lora_request=lora_request) - ) if prompts else [] + outputs: list[Any] = ( + list(llm.generate(prompts=prompts, sampling_params=sampling_params, lora_request=lora_request)) + if prompts + else [] + ) total_wall = max(time.perf_counter() - gen_start, 0.0) # Process outputs. @@ -682,7 +682,7 @@ def run_benchmark( total_valid = 0 total_invalid = 0 prompts_accepted = 0 - for output in outputs: + for prompt_idx, output in enumerate(outputs): ttft = _extract_ttft_ms(output) if ttft is not None: ttft_ms_samples.append(ttft) @@ -693,9 +693,12 @@ def run_benchmark( finish_reason = str(getattr(best, "finish_reason", None) or "unknown") finish_reasons[finish_reason] = finish_reasons.get(finish_reason, 0) + 1 text = getattr(best, "text", "") or "" - parsed = processor.process(text) - valid_records = getattr(parsed, "valid_records", None) or [] - invalid_records = getattr(parsed, "invalid_records", None) or [] + # ``Processor.__call__(prompt_number, text) -> ParsedResponse`` is + # the actual interface — see ``processors.py``. ``valid_records`` + # and ``invalid_records`` are properties on ``ParsedResponse``. + parsed = processor(prompt_idx, text) + valid_records = parsed.valid_records + invalid_records = parsed.invalid_records total_valid += len(valid_records) total_invalid += len(invalid_records) if valid_records: diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py index 91b6581d4..0601b5e13 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py @@ -62,7 +62,9 @@ """``max_model_len`` steps for the max-model-len sweep.""" -def _seeded_overrides(seed: int | None = DEFAULT_BENCHMARK_SEED, *, extra: dict[str, Any] | None = None) -> dict[str, Any]: +def _seeded_overrides( + seed: int | None = DEFAULT_BENCHMARK_SEED, *, extra: dict[str, Any] | None = None +) -> dict[str, Any]: """Build a ``sampling_overrides`` dict that pins seed + optional extras. Returns a fresh dict so callers can mutate without affecting other @@ -115,8 +117,7 @@ def baseline(base: BenchmarkEngineConfig) -> BenchmarkCandidate: def attention_backend_sweep(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: """One candidate per attention backend in :data:`ATTENTION_BACKENDS`.""" return [ - _named_copy(base, f"attention_backend={backend}", attention_backend=backend) - for backend in ATTENTION_BACKENDS + _named_copy(base, f"attention_backend={backend}", attention_backend=backend) for backend in ATTENTION_BACKENDS ] diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py index 0aa739ad3..179e6694f 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py @@ -91,11 +91,7 @@ def init_cell_run( logger.warning("wandb not installed; skipping wandb integration for this cell") return None - condition_label = ( - candidate_condition_label - or os.environ.get("BENCHMARK_CONDITION_LABEL") - or candidate_name - ) + condition_label = candidate_condition_label or os.environ.get("BENCHMARK_CONDITION_LABEL") or candidate_name dataset = os.environ.get("BENCHMARK_DATASET", "unknown") bracket_position = ( candidate_bracket_position diff --git a/tests/generation/test_vllm_benchmark.py b/tests/generation/test_vllm_benchmark.py index 040a0c782..e838c5e2c 100644 --- a/tests/generation/test_vllm_benchmark.py +++ b/tests/generation/test_vllm_benchmark.py @@ -20,7 +20,6 @@ BenchmarkCorpus, BenchmarkEngineConfig, BenchmarkOutput, - BenchmarkPrompt, CandidateMetrics, SubprocessRunResult, TraceHeader, @@ -34,7 +33,6 @@ from nemo_safe_synthesizer.generation.vllm_benchmark_presets import ( DEFAULT_BENCHMARK_SEED, DEFAULT_BRACKETED_AB_N, - PRESETS, attention_backend_sweep, baseline, bracketed_ab, @@ -43,7 +41,6 @@ ) from nemo_safe_synthesizer.generation.vllm_observability import CellObservability - # --------------------------------------------------------------------------- # Shared fixtures # --------------------------------------------------------------------------- @@ -73,9 +70,15 @@ class TestSchemaContracts: def test_candidate_metrics_composes_observability(self) -> None: """``CandidateMetrics`` embeds ``CellObservability`` rather than re-declaring its fields.""" m = CandidateMetrics( - name="t", raw_tok_s=1.0, acceptance_rate=0.9, effective_tok_s=0.9, - ttft_p50_ms=0.0, ttft_p99_ms=0.0, - prompts_attempted=10, prompts_accepted=9, total_output_tokens=100, + name="t", + raw_tok_s=1.0, + acceptance_rate=0.9, + effective_tok_s=0.9, + ttft_p50_ms=0.0, + ttft_p99_ms=0.0, + prompts_attempted=10, + prompts_accepted=9, + total_output_tokens=100, total_wall_seconds=1.0, observability=CellObservability(peak_vram_gb=64.5), ) @@ -88,9 +91,15 @@ def test_benchmark_output_roundtrips_through_json(self) -> None: corpus_size=10, candidates=[ CandidateMetrics( - name="t", raw_tok_s=1.0, acceptance_rate=0.9, effective_tok_s=0.9, - ttft_p50_ms=0.0, ttft_p99_ms=0.0, - prompts_attempted=10, prompts_accepted=9, total_output_tokens=100, + name="t", + raw_tok_s=1.0, + acceptance_rate=0.9, + effective_tok_s=0.9, + ttft_p50_ms=0.0, + ttft_p99_ms=0.0, + prompts_attempted=10, + prompts_accepted=9, + total_output_tokens=100, total_wall_seconds=1.0, observability=CellObservability(peak_vram_gb=64.5, loadavg_pre=(1.0, 2.0, 3.0)), ), @@ -201,8 +210,10 @@ def test_overlays_engine_config_on_header(self, header: TraceHeader) -> None: assert kwargs["max_lora_rank"] == 32 # from header assert kwargs["max_model_len"] == 4096 # from cfg - def test_translates_attention_backend_to_attention_config(self, header: TraceHeader, empty_base: BenchmarkEngineConfig) -> None: - """vLLM's public API takes an ``attention_config`` dict, not a bare backend string.""" + def test_translates_attention_backend_to_attention_config( + self, header: TraceHeader, empty_base: BenchmarkEngineConfig + ) -> None: + """VLLM's public API takes an ``attention_config`` dict, not a bare backend string.""" cfg = empty_base.model_copy(update={"attention_backend": "FLASHINFER"}) kwargs = _build_vllm_kwargs(header, cfg) assert kwargs["attention_config"] == {"backend": "FLASHINFER"} @@ -214,7 +225,9 @@ def test_drops_none_valued_overrides(self, header: TraceHeader, empty_base: Benc assert "enable_prefix_caching" not in kwargs assert "max_num_seqs" not in kwargs - def test_auto_attention_backend_is_treated_as_unset(self, header: TraceHeader, empty_base: BenchmarkEngineConfig) -> None: + def test_auto_attention_backend_is_treated_as_unset( + self, header: TraceHeader, empty_base: BenchmarkEngineConfig + ) -> None: """``attention_backend='auto'`` means "let vLLM pick" — no attention_config kwarg should appear.""" cfg = empty_base.model_copy(update={"attention_backend": "auto"}) kwargs = _build_vllm_kwargs(header, cfg) @@ -286,9 +299,15 @@ def test_all_cells_seed_pinned(self, empty_base: BenchmarkEngineConfig) -> None: class TestSubprocessRunResult: def test_success_shape(self) -> None: m = CandidateMetrics( - name="t", raw_tok_s=1.0, acceptance_rate=0.9, effective_tok_s=0.9, - ttft_p50_ms=0.0, ttft_p99_ms=0.0, - prompts_attempted=10, prompts_accepted=9, total_output_tokens=100, + name="t", + raw_tok_s=1.0, + acceptance_rate=0.9, + effective_tok_s=0.9, + ttft_p50_ms=0.0, + ttft_p99_ms=0.0, + prompts_attempted=10, + prompts_accepted=9, + total_output_tokens=100, total_wall_seconds=1.0, ) r = SubprocessRunResult(metrics=m) diff --git a/tests/generation/test_vllm_benchmark_analysis.py b/tests/generation/test_vllm_benchmark_analysis.py index 357c8708f..66e25be8b 100644 --- a/tests/generation/test_vllm_benchmark_analysis.py +++ b/tests/generation/test_vllm_benchmark_analysis.py @@ -11,7 +11,6 @@ from __future__ import annotations from pathlib import Path -from typing import Any import pytest @@ -20,9 +19,7 @@ CandidateMetrics, ) from nemo_safe_synthesizer.generation.vllm_benchmark_analysis import ( - MIN_CELLS_PER_CONDITION, AnalysisReport, - _effect_size, _welch_ttest_ci, analyze, ) @@ -58,9 +55,15 @@ def synthetic_sweep_dir(tmp_path: Path) -> Path: """6 baselines + 6 spec_ngram cells, both with realistic-noise spread.""" cells = [ # Baselines: ~1500 eff_tok_s, ~0.99 acceptance. - *(_metric(f"baseline_{i}", "baseline", eff=1500 + i * 5, accept=0.99, wall=130.0, bracket=2 * i) for i in range(6)), + *( + _metric(f"baseline_{i}", "baseline", eff=1500 + i * 5, accept=0.99, wall=130.0, bracket=2 * i) + for i in range(6) + ), # spec_ngram: ~1700 eff_tok_s (+~13%), similar acceptance. - *(_metric(f"spec_{i}", "spec_ngram", eff=1700 + i * 5, accept=0.99, wall=115.0, bracket=2 * i + 1) for i in range(6)), + *( + _metric(f"spec_{i}", "spec_ngram", eff=1700 + i * 5, accept=0.99, wall=115.0, bracket=2 * i + 1) + for i in range(6) + ), ] return _write_output_dir(tmp_path, cells) diff --git a/tools/vllm_benchmark.py b/tools/vllm_benchmark.py index 0f6f1ca42..ac9fc34ce 100644 --- a/tools/vllm_benchmark.py +++ b/tools/vllm_benchmark.py @@ -53,7 +53,7 @@ @click.group() def cli() -> None: - """vLLM benchmark harness CLI.""" + """VLLM benchmark harness CLI.""" @cli.command("list") @@ -186,7 +186,17 @@ def compare_cmd(output_path: Path) -> None: output = BenchmarkOutput.model_validate_json(output_path.read_text(encoding="utf-8")) sorted_candidates = sorted(output.candidates, key=lambda c: c.effective_tok_s, reverse=True) table = Table(title=f"BenchmarkOutput ({output.corpus_run_id}, n={output.corpus_size})") - for col in ("candidate", "eff tok/s", "raw tok/s", "accept", "ttft p50 ms", "ttft p99 ms", "peak vram GiB", "startup s", "ok/tried"): + for col in ( + "candidate", + "eff tok/s", + "raw tok/s", + "accept", + "ttft p50 ms", + "ttft p99 ms", + "peak vram GiB", + "startup s", + "ok/tried", + ): table.add_column(col, justify="right" if col != "candidate" else "left", overflow="fold") for m in sorted_candidates: table.add_row( From a40b2bb296549a462602b77df4e0b83590caab81 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Thu, 4 Jun 2026 15:14:22 +0000 Subject: [PATCH 11/17] refactor(benchmark): tighten types + exceptions in vLLM harness Apply the base observability branch's elegance pass to the benchmark-specific modules: - type PRESETS as dict[str, PresetFn] and _resolve_candidates(base: BenchmarkEngineConfig); drop the ty:ignore + isinstance(list) guard - replace the engine-init thread's dict[str, Any] result channel with a typed _EngineInitResult dataclass; raise InternalError if the thread finishes without an LLM or exception (narrows init_result.llm to LLM) - convert manual NvmlPeakSampler __enter__/__exit__ to a `with` block - drop the redundant try/except around read_vllm_runtime_metrics (the hardened primitive never raises and returns a stable VllmRuntimeMetrics) - replace getattr(candidate, ...) defensive reads with direct attribute access; swap the analyze cluster_signal ty:ignore for a localized cast(ClusterSignal, ...) at the click boundary No behavior change; benchmark + analyzer contract tests still pass. Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark.py | 71 +++++++++++-------- .../generation/vllm_benchmark_presets.py | 6 +- tools/vllm_benchmark.py | 24 ++++--- 3 files changed, 64 insertions(+), 37 deletions(-) diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index 21f40ae28..ff8c68ef3 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -39,12 +39,14 @@ import tempfile import threading import time +from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import TYPE_CHECKING, Any, Literal from pydantic import BaseModel, ConfigDict, Field +from ..errors import InternalError from ..generation.vllm_observability import ( CellObservability, NvmlPeakSampler, @@ -538,7 +540,9 @@ def _build_vllm_kwargs(header: TraceHeader, engine_config: BenchmarkEngineConfig if attention_backend not in (None, "auto"): base["attention_config"] = {"backend": attention_backend} # ``structured_generation_backend`` → ``structured_outputs_config``. - from vllm.config import StructuredOutputsConfig # noqa: PLC0415 — lazy, vLLM is heavy + from vllm.config import ( + StructuredOutputsConfig, # noqa: PLC0415 — lazy, vLLM is heavy + ) sg_backend = base.pop("structured_generation_backend", None) if sg_backend is not None: @@ -546,14 +550,27 @@ def _build_vllm_kwargs(header: TraceHeader, engine_config: BenchmarkEngineConfig return base +@dataclass +class _EngineInitResult: + """Typed channel for the async engine-init thread's outcome. + + Exactly one field is populated once the ``ready`` event fires: + ``llm`` on success, ``exception`` on failure. Mutated by the worker + thread; read by the runner after the join. + """ + + llm: LLM | None = None + exception: BaseException | None = None + + def _build_engine_async( header: TraceHeader, engine_config: BenchmarkEngineConfig, -) -> tuple[threading.Thread, threading.Event, dict[str, Any]]: +) -> tuple[threading.Thread, threading.Event, _EngineInitResult]: """Start ``vllm.LLM(...)`` in a daemon thread. - Returns ``(thread, ready_event, result_dict)``. The result dict has - keys ``llm`` (set on success) and ``exception`` (set on failure). + Returns ``(thread, ready_event, result)``. The :class:`_EngineInitResult` + carries ``llm`` (set on success) or ``exception`` (set on failure). The runner can sleep for the simulated-training-overlap window before joining via ``ready_event.wait()``, measuring how much of the engine-init cost can be hidden behind concurrent training. @@ -562,13 +579,13 @@ def _build_engine_async( kwargs = _build_vllm_kwargs(header, engine_config) ready = threading.Event() - result: dict[str, Any] = {"llm": None, "exception": None} + result = _EngineInitResult() def worker() -> None: try: - result["llm"] = vLLM(**kwargs) - except BaseException as exc: # noqa: BLE001 — surface via the result dict - result["exception"] = exc + result.llm = vLLM(**kwargs) + except BaseException as exc: # noqa: BLE001 — surface via the result object + result.exception = exc finally: ready.set() @@ -616,8 +633,7 @@ def run_benchmark( loadavg_pre = read_loadavg() vram_sampler = NvmlPeakSampler() - vram_sampler.__enter__() - try: + with vram_sampler: overlap = max(0.0, simulate_training_overlap_seconds) _init_thread, engine_ready, init_result = _build_engine_async( corpus.header, @@ -629,9 +645,11 @@ def run_benchmark( wait_start = time.monotonic() engine_ready.wait() startup_seconds = max(0.0, time.monotonic() - wait_start) - if init_result["exception"] is not None: - raise init_result["exception"] - llm: LLM = init_result["llm"] + if init_result.exception is not None: + raise init_result.exception + if init_result.llm is None: + raise InternalError("engine init thread finished without an LLM or an exception") + llm = init_result.llm startup_overlap_savings_seconds = min(overlap, startup_seconds + overlap) if overlap > 0.0 else 0.0 @@ -669,7 +687,13 @@ def run_benchmark( prompts: list[str] = [p.prompt for p in corpus.prompts] gen_start = time.perf_counter() outputs: list[Any] = ( - list(llm.generate(prompts=prompts, sampling_params=sampling_params, lora_request=lora_request)) + list( + llm.generate( + prompts=prompts, + sampling_params=sampling_params, + lora_request=lora_request, + ) + ) if prompts else [] ) @@ -708,21 +732,12 @@ def run_benchmark( acceptance = (total_valid / total_records) if total_records > 0 else 0.0 raw_tok_s = (total_output_tokens / total_wall) if total_wall > 0 else 0.0 records_per_second = (total_valid / total_wall) if total_wall > 0 else 0.0 - finally: - # Cleanup runs regardless of whether the body raised. Sampler - # always shuts down; observability event always builds (with - # whatever measurements were captured up to the failure). - vram_sampler.__exit__(None, None, None) - # Read end-of-generation metrics + post-load and assemble the event. - try: - vllm_metrics = read_vllm_runtime_metrics(llm) - except Exception as exc: # noqa: BLE001 — degraded mode - logger.runtime.warning( - "vllm_benchmark.metrics_read_failed", - extra={"ctx": {"error": str(exc)}}, - ) - vllm_metrics = {"kv_cache_usage_perc": None, "prefix_cache_hit_rate": None, "spec_accept_rate": None} + # The ``with`` block above shut the sampler down, so ``peak_gb`` is now + # final. ``read_vllm_runtime_metrics`` is hardened to never raise and + # always returns a stable ``VllmRuntimeMetrics`` (None-valued fields on + # degrade), so no guard is needed here. + vllm_metrics = read_vllm_runtime_metrics(llm) observability = CellObservability( peak_vram_gb=vram_sampler.peak_gb, diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py index 0601b5e13..fd93d6fc7 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py @@ -19,10 +19,14 @@ from __future__ import annotations import json +from collections.abc import Callable from typing import Any from .vllm_benchmark import BenchmarkCandidate, BenchmarkEngineConfig +PresetFn = Callable[[BenchmarkEngineConfig], list[BenchmarkCandidate]] +"""A preset resolves the corpus-default engine config into candidate cells.""" + DEFAULT_BENCHMARK_SEED: int = 42 """Default ``SamplingParams.seed`` value for preset-built candidates. @@ -312,7 +316,7 @@ def bracketed_ab_fp8(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: # Map of preset name → callable used by the CLI to resolve ``--candidates``. -PRESETS: dict[str, object] = { +PRESETS: dict[str, PresetFn] = { "baseline": lambda base: [baseline(base)], "attention_backend_sweep": attention_backend_sweep, "prefix_caching_sweep": prefix_caching_sweep, diff --git a/tools/vllm_benchmark.py b/tools/vllm_benchmark.py index ac9fc34ce..f9be5fe82 100644 --- a/tools/vllm_benchmark.py +++ b/tools/vllm_benchmark.py @@ -28,6 +28,7 @@ import json from datetime import datetime, timezone from pathlib import Path +from typing import cast import click from rich.console import Console @@ -64,7 +65,7 @@ def list_cmd() -> None: def _resolve_candidates( - base: object, + base: BenchmarkEngineConfig, preset_name: str | None, candidates_file: str | None, ) -> list[BenchmarkCandidate]: @@ -78,8 +79,7 @@ def _resolve_candidates( if preset_name: if preset_name not in PRESETS: raise click.UsageError(f"Unknown preset {preset_name!r}; available: {sorted(PRESETS)}") - resolved = PRESETS[preset_name](base) # ty: ignore[invalid-argument-type] - return resolved if isinstance(resolved, list) else [resolved] + return PRESETS[preset_name](base) if candidates_file: doc = json.loads(Path(candidates_file).read_text(encoding="utf-8")) return [BenchmarkCandidate.model_validate(c) for c in doc["candidates"]] @@ -95,7 +95,12 @@ def _resolve_candidates( type=click.Path(dir_okay=False, path_type=Path), help="Where to write the BenchmarkOutput JSON.", ) -@click.option("--candidates", "preset_name", default=None, help=f"Preset name. One of: {sorted(PRESETS)}.") +@click.option( + "--candidates", + "preset_name", + default=None, + help=f"Preset name. One of: {sorted(PRESETS)}.", +) @click.option( "--candidates-file", "candidates_file", @@ -120,7 +125,9 @@ def run_cmd( ) -> None: """Replay CORPUS_PATH against the chosen candidates and persist results.""" # Lazy import so ``list`` and ``compare`` work without spinning up vLLM. - from nemo_safe_synthesizer.generation.vllm_benchmark import run_benchmark_in_subprocess + from nemo_safe_synthesizer.generation.vllm_benchmark import ( + run_benchmark_in_subprocess, + ) corpus = BenchmarkCorpus.from_trace_jsonl(corpus_path) base = BenchmarkEngineConfig.model_validate(corpus.header.engine_parameters or {}) @@ -140,8 +147,8 @@ def run_cmd( corpus_run_id=corpus.header.run_id, corpus_size=len(corpus.prompts), sweep_id=sweep_id, - candidate_condition_label=getattr(candidate, "condition_label", ""), - candidate_bracket_position=getattr(candidate, "bracket_position", 0), + candidate_condition_label=candidate.condition_label, + candidate_bracket_position=candidate.bracket_position, ) result = run_benchmark_in_subprocess( candidate, @@ -252,12 +259,13 @@ def analyze_cmd( """Cluster-conditioned analysis across every BenchmarkOutput JSON in OUTPUT_DIR.""" from nemo_safe_synthesizer.generation.vllm_benchmark_analysis import ( MIN_CELLS_PER_CONDITION, + ClusterSignal, analyze, ) report = analyze( output_dir, - cluster_signal=cluster_signal, # ty: ignore[invalid-argument-type] + cluster_signal=cast(ClusterSignal, cluster_signal), min_cells_per_condition=MIN_CELLS_PER_CONDITION if min_cells_per_condition is None else min_cells_per_condition, ) console.print(report.to_markdown_summary()) From ad4b4dc0169ddcffa9e643ebfab51a30ad4a6ade Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Thu, 4 Jun 2026 17:39:45 +0000 Subject: [PATCH 12/17] chore(benchmark): adopt cell -> generation observability rename Follow the base branch rename of the shared observability primitive (CellObservability -> GenerationObservability, log_cell_observability -> log_generation_observability, vllm_cell -> vllm_gen wandb prefix). Update the benchmark consumers: CandidateMetrics.observability field/default, the _flatten_metrics docstring, and the data-model tests. The benchmark's own grid "cell" vocabulary (cells, n_cells, init_cell_run, per-cell aggregation) is intentionally unchanged. Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark.py | 14 +++++++------- .../generation/vllm_benchmark_wandb.py | 8 ++++---- tests/generation/test_vllm_benchmark.py | 8 ++++---- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index ff8c68ef3..b060210fc 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -20,7 +20,7 @@ (engine kwargs overlay + sparse sampling overrides + per-cell identity for sweep grouping). - :class:`CandidateMetrics` — per-cell measured outputs: throughput, - acceptance, TTFT, etc. Composes :class:`CellObservability` from PR-A's + acceptance, TTFT, etc. Composes :class:`GenerationObservability` from PR-A's ``vllm_observability`` module so the benchmark schema doesn't re-define observability primitives — it consumes them. - :class:`BenchmarkOutput` — JSON-serialised result of one matrix @@ -48,7 +48,7 @@ from ..errors import InternalError from ..generation.vllm_observability import ( - CellObservability, + GenerationObservability, NvmlPeakSampler, flag_engagement_mismatches, probe_engine_runtime_config, @@ -340,7 +340,7 @@ class CandidateMetrics(BaseModel): Carries cell-specific bench measurements (throughput, acceptance, TTFT, etc.) directly; observability primitives (peak VRAM, KV cache usage, loadavg, engine_runtime_config) are composed from - PR-A's :class:`CellObservability` schema in the ``observability`` + PR-A's :class:`GenerationObservability` schema in the ``observability`` field. This composition keeps the schema DRY — adding a new observability primitive in PR-A automatically flows through to benchmark output via ``model_dump()``. @@ -411,8 +411,8 @@ class CandidateMetrics(BaseModel): # ``probe_engine_runtime_config`` + ``read_vllm_runtime_metrics``, # and sets ``flag_did_not_engage`` based on the candidate's intended # engine config vs the probed runtime config. - observability: CellObservability = Field( - default_factory=CellObservability, + observability: GenerationObservability = Field( + default_factory=GenerationObservability, description=( "Cell-level observability snapshot. Composed from PR-A's schema " "so benchmark consumers can read e.g. ``metrics.observability." @@ -619,7 +619,7 @@ def run_benchmark( batch and prefix caching can actually fire. Wraps the entire body in PR-A's :class:`NvmlPeakSampler` context and - emits a composed :class:`CellObservability` on the returned + emits a composed :class:`GenerationObservability` on the returned :class:`CandidateMetrics`. The ``flag_did_not_engage`` bit is set when the engine's effective runtime config disagrees with the candidate's intended ``engine_config`` on any checked field. @@ -739,7 +739,7 @@ def run_benchmark( # degrade), so no guard is needed here. vllm_metrics = read_vllm_runtime_metrics(llm) - observability = CellObservability( + observability = GenerationObservability( peak_vram_gb=vram_sampler.peak_gb, kv_cache_usage_perc=vllm_metrics["kv_cache_usage_perc"], prefix_cache_hit_rate=vllm_metrics["prefix_cache_hit_rate"], diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py index 179e6694f..09cb2959f 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py @@ -5,8 +5,8 @@ Each benchmark cell becomes one wandb run (grouped by sweep ID, tagged by condition_label + dataset). Distinct from PR-A's production -``log_cell_observability`` pattern, which logs to the *currently -active* wandb run — production logs cells as a time-series within one +``log_observability_event`` pattern, which logs to the *currently +active* wandb run — production logs generations as a time-series within one run, benchmark logs each cell as its own run for per-condition isolation in the wandb UI. @@ -132,8 +132,8 @@ def _flatten_metrics(metrics: CandidateMetrics) -> dict[str, Any]: Direct fields land at the top level. ``finish_reason_distribution`` is flattened to ``finish_reason/`` scalars. The composed ``observability`` field is delegated to - :meth:`CellObservability.to_wandb_payload` with the - ``vllm_cell`` prefix so production + benchmark agree on + :meth:`GenerationObservability.to_wandb_payload` with the + ``vllm_gen`` prefix so production + benchmark agree on observability key namespacing. """ payload = metrics.model_dump(exclude={"observability", "finish_reason_distribution"}) diff --git a/tests/generation/test_vllm_benchmark.py b/tests/generation/test_vllm_benchmark.py index e838c5e2c..6fd73dafb 100644 --- a/tests/generation/test_vllm_benchmark.py +++ b/tests/generation/test_vllm_benchmark.py @@ -39,7 +39,7 @@ bracketed_ab_spec_ngram, default_matrix, ) -from nemo_safe_synthesizer.generation.vllm_observability import CellObservability +from nemo_safe_synthesizer.generation.vllm_observability import GenerationObservability # --------------------------------------------------------------------------- # Shared fixtures @@ -68,7 +68,7 @@ def empty_base() -> BenchmarkEngineConfig: class TestSchemaContracts: def test_candidate_metrics_composes_observability(self) -> None: - """``CandidateMetrics`` embeds ``CellObservability`` rather than re-declaring its fields.""" + """``CandidateMetrics`` embeds ``GenerationObservability`` rather than re-declaring its fields.""" m = CandidateMetrics( name="t", raw_tok_s=1.0, @@ -80,7 +80,7 @@ def test_candidate_metrics_composes_observability(self) -> None: prompts_accepted=9, total_output_tokens=100, total_wall_seconds=1.0, - observability=CellObservability(peak_vram_gb=64.5), + observability=GenerationObservability(peak_vram_gb=64.5), ) assert m.observability.peak_vram_gb == 64.5 @@ -101,7 +101,7 @@ def test_benchmark_output_roundtrips_through_json(self) -> None: prompts_accepted=9, total_output_tokens=100, total_wall_seconds=1.0, - observability=CellObservability(peak_vram_gb=64.5, loadavg_pre=(1.0, 2.0, 3.0)), + observability=GenerationObservability(peak_vram_gb=64.5, loadavg_pre=(1.0, 2.0, 3.0)), ), ], ) From 617fbc1fafc4a66dc5a2ba0e7a284820ba06aa51 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 15 Jun 2026 21:03:26 +0000 Subject: [PATCH 13/17] Fix vLLM benchmark PR review issues Signed-off-by: Aaron Gonzales --- docs/developer-guide/vllm-benchmark.md | 103 +++++++++ mkdocs.yml | 1 + .../generation/vllm_benchmark.py | 133 ++++++----- .../generation/vllm_benchmark_analysis.py | 62 ++--- .../generation/vllm_benchmark_presets.py | 32 +-- .../generation/vllm_benchmark_wandb.py | 28 +-- tests/generation/test_vllm_benchmark.py | 211 +++++++++++++++++- .../test_vllm_benchmark_analysis.py | 14 +- tools/vllm_benchmark.py | 13 +- 9 files changed, 462 insertions(+), 135 deletions(-) create mode 100644 docs/developer-guide/vllm-benchmark.md diff --git a/docs/developer-guide/vllm-benchmark.md b/docs/developer-guide/vllm-benchmark.md new file mode 100644 index 000000000..f42a528f3 --- /dev/null +++ b/docs/developer-guide/vllm-benchmark.md @@ -0,0 +1,103 @@ + + + +# vLLM Benchmark Harness + +The vLLM benchmark harness replays a captured generation trace against one or +more benchmark candidates. It is a developer tool for comparing engine and +sampling configurations. It is not part of the main Safe Synthesizer CLI +workflow. + +Run it through `uv` with the full engine environment: + +```bash +uv run --frozen --extra cu129 --extra engine --group dev \ + python tools/vllm_benchmark.py list +``` + +## Corpus Format + +The input corpus is a JSONL file with one header record followed by prompt +records: + +```json +{"kind": "header", "run_id": "run-1", "pretrained_model": "model-ref", "dataset_schema": {}, "engine_parameters": {}} +{"kind": "record", "row_index": 0, "prompt": "...", "sampling_params": {"temperature": 0.0}} +``` + +The header supplies the model, optional LoRA path, dataset schema, and captured +engine parameters. Each record supplies the exact prompt and sampling parameters +to replay. + +## Run a Matrix + +Use a preset matrix: + +```bash +uv run --frozen --extra cu129 --extra engine --group dev \ + python tools/vllm_benchmark.py run \ + /path/to/trace.jsonl \ + --output /path/to/benchmark.json \ + --candidates default_matrix +``` + +Use `list` to see available presets. The `bracketed_ab_*` presets emit repeated +baseline and candidate runs for noisier comparisons. + +Use a custom candidate file when a preset is too broad: + +```json +{ + "candidates": [ + { + "name": "baseline", + "engine_config": {}, + "sampling_overrides": {"seed": 42} + } + ] +} +``` + +Then run: + +```bash +uv run --frozen --extra cu129 --extra engine --group dev \ + python tools/vllm_benchmark.py run \ + /path/to/trace.jsonl \ + --output /path/to/benchmark.json \ + --candidates-file candidates.json +``` + +## Compare And Analyze + +Render one benchmark JSON: + +```bash +uv run --frozen --extra cu129 --extra engine --group dev \ + python tools/vllm_benchmark.py compare /path/to/benchmark.json +``` + +Analyze every `*.json` result in a directory: + +```bash +uv run --frozen --extra cu129 --extra engine --group dev \ + python tools/vllm_benchmark.py analyze /path/to/results-dir \ + --cluster-signal auto \ + --json-out /path/to/analysis.json +``` + +The analyzer reports candidate-run aggregates by condition. It keeps the JSON +field name `n_cells` for compatibility; in this context a "cell" means one +candidate run in the benchmark matrix. + +Use `--min-runs-per-condition` to raise or lower the refusal threshold. The +older `--min-cells-per-condition` spelling remains accepted for compatibility. + +## WandB Sink + +The harness can use WandB as a metrics sink. Each benchmark candidate run becomes +one WandB run in a shared group. WandB failures do not fail the benchmark; the +benchmark JSON is still written. + +WandB mode defaults to disabled. Use `WANDB_MODE`, `NSS_WANDB_PROJECT`, or +`WANDB_PROJECT` consistently with the rest of Safe Synthesizer. diff --git a/mkdocs.yml b/mkdocs.yml index 8b7864102..47db786e5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -194,6 +194,7 @@ nav: - Example Generation: developer-guide/example-generation.md - Observability: developer-guide/observability.md - Preflight Plugins: developer-guide/preflight-plugins.md + - vLLM Benchmark Harness: developer-guide/vllm-benchmark.md - API Reference: reference/ - Dev Notes: - dev-notes/index.md diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index b060210fc..92f079bb0 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -11,19 +11,19 @@ Architecture: -- :class:`BenchmarkCorpus` — the replayable input (one corpus per +- :class:`BenchmarkCorpus` - the replayable input (one corpus per dataset, captured once via the production ``VllmBackend`` trace surface). Header carries the model reference + LoRA path + the engine kwargs at capture time. Prompt records carry the original sampling params so the harness can replay them faithfully. -- :class:`BenchmarkCandidate` — one configuration to benchmark - (engine kwargs overlay + sparse sampling overrides + per-cell +- :class:`BenchmarkCandidate` - one configuration to benchmark + (engine kwargs overlay + sparse sampling overrides + per-candidate-run identity for sweep grouping). -- :class:`CandidateMetrics` — per-cell measured outputs: throughput, +- :class:`CandidateMetrics` - per-candidate-run measured outputs: throughput, acceptance, TTFT, etc. Composes :class:`GenerationObservability` from PR-A's ``vllm_observability`` module so the benchmark schema doesn't - re-define observability primitives — it consumes them. -- :class:`BenchmarkOutput` — JSON-serialised result of one matrix + re-define observability primitives - it consumes them. +- :class:`BenchmarkOutput` - JSON-serialised result of one matrix invocation, with skip records for candidates that failed. The runner (next commit) is in ``vllm_benchmark.py`` alongside these @@ -72,7 +72,7 @@ """Maximum bytes of captured stderr to record on a subprocess failure.""" PromptAssemblyMode = Literal["multi_record", "per_record"] -"""Prompt-assembly regime — controls how max_tokens partitions the budget.""" +"""Prompt-assembly regime - controls how max_tokens partitions the budget.""" BatchDispatchMode = Literal["replicate", "n_fanout"] """How corpus prompts get submitted to vLLM. @@ -201,7 +201,7 @@ class BenchmarkEngineConfig(BaseModel): against that; for now it stands alone. ``extra='ignore'`` so the CLI can validate a corpus header's raw - ``engine_parameters`` dict into this model — header dicts may carry + ``engine_parameters`` dict into this model - header dicts may carry capture-time kwargs we don't expose as typed fields (those flow through ``_build_vllm_kwargs`` as the base layer regardless). """ @@ -214,7 +214,10 @@ class BenchmarkEngineConfig(BaseModel): ) structured_generation_backend: str = Field( default="xgrammar", - description="Structured-outputs backend. ``'xgrammar'`` is vLLM's current default; ``'outlines'`` and ``'guidance'`` are the alternatives.", + description=( + "Structured-outputs backend used by benchmark sweeps. The preset " + "matrix covers ``'xgrammar'``, ``'outlines'``, and ``'guidance'``." + ), ) max_model_len: int | None = Field( default=None, @@ -225,7 +228,7 @@ class BenchmarkEngineConfig(BaseModel): description=( "Forwarded to ``vllm.LLM(enable_prefix_caching=...)``. " "On for shared-schema tabular workloads (the prefix amortises across the batch); " - "off when measuring per-cell cold-start behaviour. ``None`` keeps vLLM's default." + "off when measuring per-candidate-run cold-start behaviour. ``None`` keeps vLLM's default." ), ) max_num_seqs: int | None = Field(default=None, description="vLLM scheduler ``max_num_seqs`` cap.") @@ -276,9 +279,9 @@ class BenchmarkEngineConfig(BaseModel): class BenchmarkCandidate(BaseModel): - """One configuration to benchmark — engine kwargs + sampling overrides + identity. + """One configuration to benchmark - engine kwargs + sampling overrides + identity. - ``sampling_overrides`` is sparse — only fields that differ from the + ``sampling_overrides`` is sparse - only fields that differ from the corpus default. The runner merges it on top of each prompt's ``original_sampling_params``. Pass ``{"seed": int}`` to pin ``SamplingParams.seed`` for reproducible acceptance-rate @@ -318,30 +321,30 @@ class BenchmarkCandidate(BaseModel): "Sweep-level condition this candidate measures, e.g. " "``'baseline'``, ``'n_fanout'``, ``'spec_ngram'``, ``'fp8'``. " "Set by the ``bracketed_ab`` preset family. Used by the " - "cluster-conditioned analyzer to group cells by condition " - "regardless of per-cell name suffixes." + "cluster-conditioned analyzer to group candidate runs by condition " + "regardless of candidate-run name suffixes." ), ) bracket_position: int = Field( default=0, ge=0, description=( - "Sequence index within a ``bracketed_ab`` cell stream " + "Sequence index within a ``bracketed_ab`` candidate-run stream " "(baseline_0=0, candidate_0=1, baseline_1=2, candidate_1=3, " - "etc.). Used by the analyzer to align candidate cells with " + "etc.). Used by the analyzer to align candidate runs with " "their bracketing baselines for drift detection." ), ) class CandidateMetrics(BaseModel): - """Measured outputs for one benchmark cell. + """Measured outputs for one benchmark candidate run. - Carries cell-specific bench measurements (throughput, acceptance, + Carries candidate-run bench measurements (throughput, acceptance, TTFT, etc.) directly; observability primitives (peak VRAM, KV cache usage, loadavg, engine_runtime_config) are composed from PR-A's :class:`GenerationObservability` schema in the ``observability`` - field. This composition keeps the schema DRY — adding a new + field. This composition keeps the schema DRY - adding a new observability primitive in PR-A automatically flows through to benchmark output via ``model_dump()``. """ @@ -350,14 +353,14 @@ class CandidateMetrics(BaseModel): name: str = Field(description="Matches :attr:`BenchmarkCandidate.name`.") - # Cell-level throughput + acceptance. + # Candidate-run throughput + acceptance. raw_tok_s: float = Field(description="Output tokens / wall seconds, ignoring validity.") acceptance_rate: float = Field(description="Fraction of generated records that passed validation (0..1).") effective_tok_s: float = Field(description="``raw_tok_s * acceptance_rate``; the operator-relevant headline.") # Per-request latency stats. TTFT is queue-inclusive under batched # submission (vLLM's ``first_token_latency`` = ``first_token_ts - - # arrival_time`` — so prompts later in the batch contribute their + # arrival_time`` - so prompts later in the batch contribute their # queue wait). Useful for spotting tail-of-batch wait time but not # for per-candidate comparisons under varying batch shapes. ttft_p50_ms: float = Field(description="Median time-to-first-token in milliseconds (queue-inclusive).") @@ -393,7 +396,7 @@ class CandidateMetrics(BaseModel): description="Counts of vLLM finish reasons (``stop``, ``length``, etc.) across replays.", ) - # Sweep grouping — copied from BenchmarkCandidate so the analyzer + # Sweep grouping - copied from BenchmarkCandidate so the analyzer # can read them off CandidateMetrics without re-joining to the # BenchmarkCandidate by name. condition_label: str = Field( @@ -438,7 +441,7 @@ class SkipRecord(BaseModel): class BenchmarkOutput(BaseModel): - """Full result of one matrix invocation — JSON-serialised for diffing.""" + """Full result of one matrix invocation - JSON-serialised for diffing.""" model_config = ConfigDict(extra="forbid") @@ -534,14 +537,14 @@ def _build_vllm_kwargs(header: TraceHeader, engine_config: BenchmarkEngineConfig # Required-positional kwargs that aren't in BenchmarkEngineConfig: base["model"] = header.pretrained_model base.setdefault("enable_lora", header.lora_path is not None) - # ``attention_backend`` → ``attention_config`` translation. vLLM's + # ``attention_backend`` -> ``attention_config`` translation. vLLM's # public API takes a config dict rather than a bare string. attention_backend = base.pop("attention_backend", None) if attention_backend not in (None, "auto"): base["attention_config"] = {"backend": attention_backend} - # ``structured_generation_backend`` → ``structured_outputs_config``. + # ``structured_generation_backend`` -> ``structured_outputs_config``. from vllm.config import ( - StructuredOutputsConfig, # noqa: PLC0415 — lazy, vLLM is heavy + StructuredOutputsConfig, # noqa: PLC0415 - lazy, vLLM is heavy ) sg_backend = base.pop("structured_generation_backend", None) @@ -554,13 +557,15 @@ def _build_vllm_kwargs(header: TraceHeader, engine_config: BenchmarkEngineConfig class _EngineInitResult: """Typed channel for the async engine-init thread's outcome. - Exactly one field is populated once the ``ready`` event fires: - ``llm`` on success, ``exception`` on failure. Mutated by the worker - thread; read by the runner after the join. + Once the ``ready`` event fires, ``llm`` is populated on success or + ``exception`` on failure; ``init_seconds`` records the worker-side + engine-construction duration. Mutated by the worker thread; read by + the runner after the join. """ llm: LLM | None = None exception: BaseException | None = None + init_seconds: float = 0.0 def _build_engine_async( @@ -575,18 +580,20 @@ def _build_engine_async( before joining via ``ready_event.wait()``, measuring how much of the engine-init cost can be hidden behind concurrent training. """ - from vllm import LLM as vLLM # noqa: PLC0415 — lazy + from vllm import LLM as vLLM # noqa: PLC0415 - lazy kwargs = _build_vllm_kwargs(header, engine_config) ready = threading.Event() result = _EngineInitResult() def worker() -> None: + init_start = time.monotonic() try: result.llm = vLLM(**kwargs) - except BaseException as exc: # noqa: BLE001 — surface via the result object + except BaseException as exc: # noqa: BLE001 - surface via the result object result.exception = exc finally: + result.init_seconds = max(0.0, time.monotonic() - init_start) ready.set() thread = threading.Thread(target=worker, name="vllm-benchmark-engine-init", daemon=True) @@ -624,7 +631,7 @@ def run_benchmark( when the engine's effective runtime config disagrees with the candidate's intended ``engine_config`` on any checked field. """ - # Lazy imports — keep this module CPU-importable. + # Lazy imports - keep this module CPU-importable. from vllm.lora.request import LoRARequest # noqa: PLC0415 from vllm.sampling_params import SamplingParams # noqa: PLC0415 @@ -651,7 +658,10 @@ def run_benchmark( raise InternalError("engine init thread finished without an LLM or an exception") llm = init_result.llm - startup_overlap_savings_seconds = min(overlap, startup_seconds + overlap) if overlap > 0.0 else 0.0 + engine_init_seconds = max(0.0, getattr(init_result, "init_seconds", 0.0)) + startup_overlap_savings_seconds = ( + min(overlap, max(0.0, engine_init_seconds - startup_seconds)) if overlap > 0.0 else 0.0 + ) # Probe the engine's effective runtime config + check for # candidate-intent / engine-actual disagreements. @@ -674,27 +684,32 @@ def run_benchmark( ) # Build SamplingParams from corpus default + candidate overrides. + base_sampling: dict[str, Any] = {} if corpus.prompts: base_sampling = corpus.prompts[0].original_sampling_params - else: - base_sampling = {} sampling_kwargs = _merge_sampling_kwargs(base_sampling, candidate.sampling_overrides) - # n=1 unless the caller's override sets it (n_fanout sets it explicitly). - sampling_kwargs.setdefault("n", 1) + prompts: list[str] = [p.prompt for p in corpus.prompts] + if candidate.batch_dispatch_mode == "n_fanout" and len(set(prompts)) > 1: + raise ValueError("batch_dispatch_mode='n_fanout' requires every corpus prompt to be identical") + if candidate.batch_dispatch_mode == "n_fanout": + sampling_kwargs["n"] = max(1, len(prompts)) + dispatch_prompts = prompts[:1] + else: + sampling_kwargs.setdefault("n", 1) + dispatch_prompts = prompts sampling_params = SamplingParams(**sampling_kwargs) # Dispatch. - prompts: list[str] = [p.prompt for p in corpus.prompts] gen_start = time.perf_counter() outputs: list[Any] = ( list( llm.generate( - prompts=prompts, + prompts=dispatch_prompts, sampling_params=sampling_params, lora_request=lora_request, ) ) - if prompts + if dispatch_prompts else [] ) total_wall = max(time.perf_counter() - gen_start, 0.0) @@ -706,27 +721,31 @@ def run_benchmark( total_valid = 0 total_invalid = 0 prompts_accepted = 0 + completion_idx = 0 for prompt_idx, output in enumerate(outputs): ttft = _extract_ttft_ms(output) if ttft is not None: ttft_ms_samples.append(ttft) - best = output.outputs[0] if getattr(output, "outputs", None) else None - if best is None: - continue - total_output_tokens += len(getattr(best, "token_ids", []) or []) - finish_reason = str(getattr(best, "finish_reason", None) or "unknown") - finish_reasons[finish_reason] = finish_reasons.get(finish_reason, 0) + 1 - text = getattr(best, "text", "") or "" - # ``Processor.__call__(prompt_number, text) -> ParsedResponse`` is - # the actual interface — see ``processors.py``. ``valid_records`` - # and ``invalid_records`` are properties on ``ParsedResponse``. - parsed = processor(prompt_idx, text) - valid_records = parsed.valid_records - invalid_records = parsed.invalid_records - total_valid += len(valid_records) - total_invalid += len(invalid_records) - if valid_records: - prompts_accepted += 1 + completions = list(getattr(output, "outputs", []) or []) + if candidate.batch_dispatch_mode != "n_fanout": + completions = completions[:1] + for completion in completions: + record_idx = completion_idx if candidate.batch_dispatch_mode == "n_fanout" else prompt_idx + completion_idx += 1 + total_output_tokens += len(getattr(completion, "token_ids", []) or []) + finish_reason = str(getattr(completion, "finish_reason", None) or "unknown") + finish_reasons[finish_reason] = finish_reasons.get(finish_reason, 0) + 1 + text = getattr(completion, "text", "") or "" + # ``Processor.__call__(prompt_number, text) -> ParsedResponse`` is + # the actual interface - see ``processors.py``. ``valid_records`` + # and ``invalid_records`` are properties on ``ParsedResponse``. + parsed = processor(record_idx, text) + valid_records = parsed.valid_records + invalid_records = parsed.invalid_records + total_valid += len(valid_records) + total_invalid += len(invalid_records) + if valid_records: + prompts_accepted += 1 total_records = total_valid + total_invalid acceptance = (total_valid / total_records) if total_records > 0 else 0.0 diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py index 674e0a1be..39f00bdf1 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py @@ -3,32 +3,32 @@ """Cluster-conditioned analysis for benchmark output. -This stack does not produce trustworthy single-cell measurements. -Pooled cross-cell CoV runs ~5-10%; in-cluster CoV is ~3%. This module -partitions cells by a per-dataset cluster signal, then reports +This stack does not produce trustworthy single-run measurements. +Pooled cross-run CoV runs ~5-10%; in-cluster CoV is ~3%. This module +partitions candidate runs by a per-dataset cluster signal, then reports per-condition aggregates both pooled and within-cluster, so the operator can see which clusters contained the candidate's samples -and trust the in-cluster Δ rather than the noisier pooled Δ. +and trust the in-cluster delta rather than the noisier pooled delta. Cluster signal per workload shape: -- ``wall_seconds`` for short-context workloads (bike_sales-shape) — +- ``wall_seconds`` for short-context workloads (bike_sales-shape) - bimodality is load-driven; partitioning on wall_seconds separates the fast vs normal-load clusters. - ``acceptance_rate`` for long-output workloads (call_transcripts- - shape) — bimodality is RNG/scheduler driven; partitioning on + shape) - bimodality is RNG/scheduler driven; partitioning on acceptance_rate separates the high vs low cluster the seed-pin validation found persists even with seed=42. -- ``auto`` — picks whichever signal has higher pooled CoV. +- ``auto`` - picks whichever signal has higher pooled CoV. -Cluster count is selected via silhouette score (post-hoc, k ∈ [2, 4]). -Refuses to compute Δ-style aggregates when a condition has fewer than -:data:`MIN_CELLS_PER_CONDITION` cells — single-cell measurements +Cluster count is selected via silhouette score (post-hoc, k in [2, 4]). +Refuses to compute delta-style aggregates when a condition has fewer than +:data:`MIN_CELLS_PER_CONDITION` candidate runs - single-run measurements should never drive promote/reject decisions on this stack. -Effect-size + 95% CI reporting lives in this module too — see +Effect-size + 95% CI reporting lives in this module too - see :class:`EffectSize`, computed via Welch's t-test on the difference of -means + Welch–Satterthwaite degrees of freedom. +means + Welch-Satterthwaite degrees of freedom. """ from __future__ import annotations @@ -46,7 +46,7 @@ ClusterSignal = Literal["wall_seconds", "acceptance_rate", "auto"] MIN_CELLS_PER_CONDITION: int = 6 -"""Minimum cells per condition before Δ-style aggregates are computed. +"""Minimum candidate runs per condition before delta-style aggregates are computed. Matches :data:`DEFAULT_BRACKETED_AB_N`. Below this threshold the analyzer records a refusal in the report rather than emitting @@ -63,7 +63,7 @@ class ClusterAssignment(BaseModel): - """One cell's assignment to a cluster + its raw signal value.""" + """One candidate run's assignment to a cluster + its raw signal value.""" model_config = ConfigDict(extra="forbid") @@ -87,7 +87,7 @@ class ClusterStats(BaseModel): class EffectSize(BaseModel): - """Welch's-t Δ ± 95% CI for one condition vs a baseline reference. + """Welch's-t delta and 95% CI for one condition vs a baseline reference. ``cluster_id=None`` is a pooled effect across all clusters; integer values mean within-cluster (the brief-mandated headline @@ -112,7 +112,7 @@ class EffectSize(BaseModel): class ConditionClusterAggregate(BaseModel): - """Per-(condition × cluster) aggregate: in-cluster condition stats.""" + """Per-(condition x cluster) aggregate: in-cluster condition stats.""" model_config = ConfigDict(extra="forbid") @@ -160,13 +160,13 @@ class AnalysisReport(BaseModel): def to_markdown_summary(self) -> str: """Render a human-readable summary.""" lines: list[str] = [ - f"# Cluster-conditioned analysis ({self.n_cells} cells, k={self.n_clusters})", + f"# Cluster-conditioned analysis ({self.n_cells} candidate runs, k={self.n_clusters})", "", - f"**Cluster signal**: `{self.cluster_signal}`", + f"Cluster signal: `{self.cluster_signal}`", "", "## Clusters", "", - "| cluster | n_cells | signal_mean | signal_stddev | signal_cov |", + "| cluster | candidate_runs | signal_mean | signal_stddev | signal_cov |", "|--------:|--------:|------------:|--------------:|-----------:|", ] for cs in self.cluster_stats: @@ -179,8 +179,8 @@ def to_markdown_summary(self) -> str: lines.append(f"### `{agg.condition_label}` (n={agg.n_cells})") lines.append("") lines.append( - f"- **Pooled**: eff_tok_s={agg.pooled_mean_effective_tok_s:.1f} " - f"± {agg.pooled_stddev_effective_tok_s:.1f} " + f"- Pooled: eff_tok_s={agg.pooled_mean_effective_tok_s:.1f} " + f"+/- {agg.pooled_stddev_effective_tok_s:.1f} " f"(CoV {agg.pooled_cov_effective_tok_s * 100:.2f}%); " f"accept={agg.pooled_mean_acceptance_rate:.4f} " f"(CoV {agg.pooled_cov_acceptance_rate * 100:.2f}%)" @@ -188,24 +188,24 @@ def to_markdown_summary(self) -> str: if agg.pooled_effect_size_vs_baseline is not None: es = agg.pooled_effect_size_vs_baseline lines.append( - f" - Δ vs baseline (pooled): {es.delta_absolute:+.1f} tok/s " + f" - Delta vs baseline (pooled): {es.delta_absolute:+.1f} tok/s " f"({es.delta_pct:+.2f}%) [95% CI: {es.ci95_low:+.1f}, {es.ci95_high:+.1f}; " f"n={es.n_candidate}+{es.n_baseline}, Welch df={es.welch_df:.1f}]" ) if agg.in_cluster: - lines.append("- **In-cluster**:") + lines.append("- In-cluster:") for ic in agg.in_cluster: lines.append( f" - cluster {ic.cluster_id}: n={ic.n_cells}, " f"eff_tok_s={ic.mean_effective_tok_s:.1f} " - f"± {ic.stddev_effective_tok_s:.1f} " + f"+/- {ic.stddev_effective_tok_s:.1f} " f"(CoV {ic.cov_effective_tok_s * 100:.2f}%); " f"accept={ic.mean_acceptance_rate:.4f}" ) if ic.effect_size_vs_baseline is not None: es = ic.effect_size_vs_baseline lines.append( - f" - Δ vs baseline cluster {es.cluster_id}: " + f" - Delta vs baseline cluster {es.cluster_id}: " f"{es.delta_absolute:+.1f} tok/s ({es.delta_pct:+.2f}%) " f"[95% CI: {es.ci95_low:+.1f}, {es.ci95_high:+.1f}; " f"n={es.n_candidate}+{es.n_baseline}, Welch df={es.welch_df:.1f}]" @@ -244,7 +244,7 @@ def _pooled_cov(values: list[float]) -> tuple[float, float, float]: def _auto_select_signal(cells: list[CandidateMetrics]) -> str: - """Pick whichever signal has higher pooled CoV across the cells. + """Pick whichever signal has higher pooled CoV across the candidate runs. Defaults to ``wall_seconds`` on ties (short-context bimodality is the more common workload shape). @@ -259,7 +259,7 @@ def _auto_select_signal(cells: list[CandidateMetrics]) -> str: def _select_n_clusters(values: np.ndarray, k_max: int = _DEFAULT_K_MAX) -> int: """Silhouette-score-best ``k`` in range ``[2, min(k_max, n-1)]``. - Returns 1 when there are too few cells (<4) to cluster meaningfully. + Returns 1 when there are too few candidate runs (<4) to cluster meaningfully. """ from sklearn.cluster import KMeans # noqa: PLC0415 from sklearn.metrics import silhouette_score # noqa: PLC0415 @@ -373,7 +373,7 @@ def _effect_size( def load_cells(output_dir: Path) -> list[CandidateMetrics]: - """Read every ``BenchmarkOutput`` JSON in ``output_dir``, flatten candidates. + """Read every ``BenchmarkOutput`` JSON in ``output_dir``, flatten candidate runs. Subdirectories are NOT recursed. Callers wanting cross-dataset analysis should invoke once per dataset dir. @@ -393,7 +393,7 @@ def analyze( cluster_signal: ClusterSignal = "auto", min_cells_per_condition: int = MIN_CELLS_PER_CONDITION, ) -> AnalysisReport: - """Full pipeline: load → cluster → per-condition aggregate → effect-size → report.""" + """Full pipeline: load -> cluster -> per-condition aggregate -> effect-size -> report.""" cells = load_cells(output_dir) if not cells: raise ValueError(f"No BenchmarkOutput JSONs found under {output_dir}") @@ -444,8 +444,8 @@ def analyze( labeled = by_condition[condition] if len(labeled) < min_cells_per_condition: refusals.append( - f"condition {condition!r} has only {len(labeled)} cells; " - f"need ≥{min_cells_per_condition} — refusing aggregate" + f"condition {condition!r} has only {len(labeled)} candidate runs; " + f"need >={min_cells_per_condition} - refusing aggregate" ) continue pooled_eff = [c.effective_tok_s for _, c in labeled] diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py index fd93d6fc7..91f7c757e 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py @@ -25,7 +25,7 @@ from .vllm_benchmark import BenchmarkCandidate, BenchmarkEngineConfig PresetFn = Callable[[BenchmarkEngineConfig], list[BenchmarkCandidate]] -"""A preset resolves the corpus-default engine config into candidate cells.""" +"""A preset resolves the corpus-default engine config into candidate runs.""" DEFAULT_BENCHMARK_SEED: int = 42 """Default ``SamplingParams.seed`` value for preset-built candidates. @@ -41,7 +41,7 @@ """Fallback ``max_model_len`` for presets when the corpus header doesn't have a resolver hint. vLLM 0.20 defaults this to whatever the model's tokenizer reports, which -for Mistral-7B is 32768 — over-provisions the KV cache budget for the +for Mistral-7B is 32768 - over-provisions the KV cache budget for the tabular workloads we benchmark. """ @@ -107,7 +107,7 @@ def _named_copy(base: BenchmarkEngineConfig, name: str, **updates: Any) -> Bench def baseline(base: BenchmarkEngineConfig) -> BenchmarkCandidate: - """Pass-through candidate — runs the corpus's default config unchanged. + """Pass-through candidate - runs the corpus's default config unchanged. Still seed-pinned for reproducibility. See :func:`_named_copy`. """ @@ -165,7 +165,7 @@ def max_model_len_sweep(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate] def default_matrix(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: """Concatenation of every sweep, deduplicated by (engine config, overrides). - Dedup is conservative — two candidates with identical engine config + Dedup is conservative - two candidates with identical engine config but different ``name`` are still considered duplicates because the resulting engine + sampling combination is what the runner measures. """ @@ -191,12 +191,12 @@ def default_matrix(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: DEFAULT_BRACKETED_AB_N: int = 6 -"""Cells per condition in :func:`bracketed_ab`. +"""Candidate runs per condition in :func:`bracketed_ab`. N=6 is the methodology-critic-recommended compromise between N=4 (~20% statistical power for detecting +5% effects under the stack's observed -~3% in-cluster CoV) and N=8 (busts the 24h × $350 envelope on a -4-dataset × 4-condition matrix). +~3% in-cluster CoV) and N=8 (busts the 24h x $350 envelope on a +4-dataset x 4-condition matrix). """ @@ -209,18 +209,18 @@ def bracketed_ab( condition_label: str, n_samples_per_condition: int = DEFAULT_BRACKETED_AB_N, ) -> list[BenchmarkCandidate]: - """Emit an interleaved baseline-candidate cell sequence for bracketed A/B. + """Emit an interleaved baseline-candidate run sequence for bracketed A/B. - Returns ``2 * n_samples_per_condition`` cells: ``n`` baselines - interleaved with ``n`` candidate cells. ``bracket_position`` is set - on each cell so the cluster-conditioned analyzer can align - candidate cells with their bracketing baselines for drift detection. + Returns ``2 * n_samples_per_condition`` runs: ``n`` baselines + interleaved with ``n`` candidate runs. ``bracket_position`` is set + on each run so the cluster-conditioned analyzer can align + candidate runs with their bracketing baselines for drift detection. Both baselines and candidates pin ``SamplingParams.seed=DEFAULT_BENCHMARK_SEED``. The seed-pin verification (2026-05-26) showed this does NOT collapse acceptance variance to <0.5% CoV on long-output workloads but DOES eliminate - the per-request RNG portion — residual ~5% pooled CoV is structural + the per-request RNG portion - residual ~5% pooled CoV is structural non-RNG variance that the cluster-conditioned analyzer partitions out post-hoc. """ @@ -250,12 +250,12 @@ def bracketed_ab( return cells -# Phase B matrix-condition wrappers. Each yields a 2N-cell sequence +# Phase B matrix-condition wrappers. Each yields a 2N-run sequence # (N baselines + N condition-specific candidates, interleaved). def bracketed_ab_baseline_pool(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: - """N baseline cells with bracket_position labels — the shared baseline pool.""" + """N baseline runs with bracket_position labels - the shared baseline pool.""" return [ BenchmarkCandidate( name=f"bracket_baseline_pool_{i}", @@ -315,7 +315,7 @@ def bracketed_ab_fp8(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: ) -# Map of preset name → callable used by the CLI to resolve ``--candidates``. +# Map of preset name -> callable used by the CLI to resolve ``--candidates``. PRESETS: dict[str, PresetFn] = { "baseline": lambda base: [baseline(base)], "attention_backend_sweep": attention_backend_sweep, diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py index 09cb2959f..7610bef4c 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py @@ -1,13 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Benchmark-side wandb integration — per-cell new-run mode. +"""WandB metrics sink for vLLM benchmark runs. -Each benchmark cell becomes one wandb run (grouped by sweep ID, tagged +Each benchmark candidate run becomes one WandB run (grouped by sweep ID, tagged by condition_label + dataset). Distinct from PR-A's production ``log_observability_event`` pattern, which logs to the *currently -active* wandb run — production logs generations as a time-series within one -run, benchmark logs each cell as its own run for per-condition +active* WandB run: production logs generations as a time-series within one +run, while benchmark mode logs each candidate run as its own run for per-condition isolation in the wandb UI. Reuses :class:`WandbSettings` from ``nemo_safe_synthesizer.cli.wandb_setup`` @@ -17,7 +17,7 @@ Soft dependency: wandb is observability, not a hard requirement. Any failure (missing package, missing netrc, ``wandb.init`` exception) -logs a warning and the harness continues — the benchmark JSON output +logs a warning and the harness continues; the benchmark JSON output is still written. """ @@ -35,7 +35,7 @@ logger = get_logger(__name__) -# A benchmark cell is structurally a GENERATE invocation that's +# A benchmark candidate run is structurally a GENERATE invocation that's # measured rather than consumed. Phase aligns with production GENERATE # runs; job_type distinguishes benchmark from production at the # wandb-UI level. @@ -46,8 +46,8 @@ def resolve_sweep_id() -> str: """Resolve the wandb group identifier for the current sweep. Reads ``WANDB_RUN_GROUP`` when set (the orchestrator sets this once - per sweep to group all cells). Falls back to an auto-generated - timestamp so single-cell invocations don't all collapse into one + per sweep to group all candidate runs). Falls back to an auto-generated + timestamp so single-run invocations don't all collapse into one bucket. """ return os.environ.get("WANDB_RUN_GROUP") or f"sweep-{datetime.now(timezone.utc):%Y%m%d-%H%M%S}" @@ -64,7 +64,7 @@ def init_cell_run( candidate_condition_label: str = "", candidate_bracket_position: int = 0, ) -> Any: - """Open one wandb run for a single benchmark cell. + """Open one WandB run for a single benchmark candidate run. Returns the run object on success, ``None`` when wandb is disabled, misconfigured, or the package is unavailable. Callers must treat @@ -79,16 +79,16 @@ def init_cell_run( path is set by the ``bracketed_ab`` preset family. Auth via ``~/.netrc`` (set up by ``wandb login``). - ``WANDB_API_KEY`` is NOT read — passing the key via env leaks it + ``WANDB_API_KEY`` is NOT read; passing the key via env leaks it via ``ps auxe``. """ settings = WandbSettings() if settings.wandb_mode == WandbMode.DISABLED: return None try: - import wandb # noqa: PLC0415 — soft dependency + import wandb # noqa: PLC0415 - soft dependency except ImportError: - logger.warning("wandb not installed; skipping wandb integration for this cell") + logger.warning("wandb not installed; skipping wandb metrics sink for this candidate run") return None condition_label = candidate_condition_label or os.environ.get("BENCHMARK_CONDITION_LABEL") or candidate_name @@ -121,7 +121,7 @@ def init_cell_run( "phase": WandbPhase.GENERATE.value, }, ) - except Exception as exc: # noqa: BLE001 — degraded mode by design + except Exception as exc: # noqa: BLE001 - degraded mode by design logger.warning("wandb.init failed; continuing without wandb", exc_info=exc) return None @@ -158,5 +158,5 @@ def log_and_finish(run: Any, metrics: CandidateMetrics | None, exit_code: int = wandb.log(_flatten_metrics(metrics)) run.finish(exit_code=exit_code) - except Exception as exc: # noqa: BLE001 — degraded mode + except Exception as exc: # noqa: BLE001 - degraded mode logger.warning("wandb finish failed; continuing", exc_info=exc) diff --git a/tests/generation/test_vllm_benchmark.py b/tests/generation/test_vllm_benchmark.py index 6fd73dafb..14755fc7f 100644 --- a/tests/generation/test_vllm_benchmark.py +++ b/tests/generation/test_vllm_benchmark.py @@ -1,25 +1,32 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for benchmark harness data models, helpers, and presets. +"""Tests for benchmark harness data models, helpers, presets, and runner dispatch. -Scope: consumer-facing contracts. Skip the runner itself (requires -spinning up vLLM); covered by the actual production cell smokes. +Scope: consumer-facing contracts. Runner tests patch vLLM imports and +engine construction so they exercise dispatch/metric semantics without +spinning up a real vLLM engine. """ from __future__ import annotations import json +import sys +import types +from types import SimpleNamespace from typing import Any import pytest from pydantic import ValidationError +import nemo_safe_synthesizer.generation.processors as processors_mod +import nemo_safe_synthesizer.generation.vllm_benchmark as benchmark_mod from nemo_safe_synthesizer.generation.vllm_benchmark import ( BenchmarkCandidate, BenchmarkCorpus, BenchmarkEngineConfig, BenchmarkOutput, + BenchmarkPrompt, CandidateMetrics, SubprocessRunResult, TraceHeader, @@ -29,6 +36,7 @@ _parse_error_class, _percentile, _truncate_stderr, + run_benchmark, ) from nemo_safe_synthesizer.generation.vllm_benchmark_presets import ( DEFAULT_BENCHMARK_SEED, @@ -61,6 +69,139 @@ def empty_base() -> BenchmarkEngineConfig: return BenchmarkEngineConfig() +class _FakeSamplingParams: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + +class _FakeLoRARequest: + def __init__(self, *args: Any) -> None: + self.args = args + + +class _NoopPeakSampler: + peak_gb: float | None = None + + def __enter__(self) -> _NoopPeakSampler: + return self + + def __exit__(self, *args: object) -> None: + return None + + +class _ReadyEvent: + def wait(self) -> None: + return None + + +class _FakeCompletion: + def __init__(self, text: str, token_ids: list[int], finish_reason: str = "stop") -> None: + self.text = text + self.token_ids = token_ids + self.finish_reason = finish_reason + + +class _FakeRequestOutput: + def __init__(self, outputs: list[_FakeCompletion], ttft_s: float = 0.01) -> None: + self.outputs = outputs + self.metrics = SimpleNamespace(first_token_latency=ttft_s) + + +class _FakeLLM: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def generate(self, *, prompts: list[str], sampling_params: _FakeSamplingParams, lora_request: Any) -> list[Any]: + self.calls.append( + { + "prompts": list(prompts), + "sampling_params": sampling_params, + "lora_request": lora_request, + }, + ) + if len(prompts) == 1 and sampling_params.kwargs.get("n", 1) > 1: + return [ + _FakeRequestOutput( + [ + _FakeCompletion(text=f'{{"col": "v{i}"}}', token_ids=[i, i + 100]) + for i in range(sampling_params.kwargs["n"]) + ], + ), + ] + return [ + _FakeRequestOutput([_FakeCompletion(text=f'{{"col": "v{i}"}}', token_ids=[i])]) + for i, _prompt in enumerate(prompts) + ] + + def get_metrics(self) -> list[Any]: + return [] + + +class _FakeProcessor: + def __init__(self, calls: list[tuple[int, str]]) -> None: + self._calls = calls + + def __call__(self, prompt_number: int, text: str) -> Any: + self._calls.append((prompt_number, text)) + return SimpleNamespace(valid_records=[{"text": text}], invalid_records=[]) + + +def _install_fake_vllm_modules(monkeypatch: pytest.MonkeyPatch) -> None: + vllm_mod = types.ModuleType("vllm") + vllm_mod.__path__ = [] + lora_mod = types.ModuleType("vllm.lora") + lora_mod.__path__ = [] + request_mod = types.ModuleType("vllm.lora.request") + setattr(request_mod, "LoRARequest", _FakeLoRARequest) + sampling_mod = types.ModuleType("vllm.sampling_params") + setattr(sampling_mod, "SamplingParams", _FakeSamplingParams) + monkeypatch.setitem(sys.modules, "vllm", vllm_mod) + monkeypatch.setitem(sys.modules, "vllm.lora", lora_mod) + monkeypatch.setitem(sys.modules, "vllm.lora.request", request_mod) + monkeypatch.setitem(sys.modules, "vllm.sampling_params", sampling_mod) + + +def _install_runner_fakes( + monkeypatch: pytest.MonkeyPatch, + llm: _FakeLLM, + *, + init_seconds: float = 0.0, + processor_calls: list[tuple[int, str]] | None = None, +) -> None: + _install_fake_vllm_modules(monkeypatch) + calls = processor_calls if processor_calls is not None else [] + monkeypatch.setattr(processors_mod, "TabularDataProcessor", lambda *args, **kwargs: _FakeProcessor(calls)) + monkeypatch.setattr(benchmark_mod, "NvmlPeakSampler", _NoopPeakSampler) + monkeypatch.setattr(benchmark_mod, "read_loadavg", lambda: (0.0, 0.0, 0.0)) + monkeypatch.setattr(benchmark_mod, "probe_engine_runtime_config", lambda llm: {}) + monkeypatch.setattr(benchmark_mod, "flag_engagement_mismatches", lambda intended, actual: []) + monkeypatch.setattr( + benchmark_mod, + "read_vllm_runtime_metrics", + lambda llm: {"kv_cache_usage_perc": None, "prefix_cache_hit_rate": None, "spec_accept_rate": None}, + ) + + def fake_build_engine_async(header: TraceHeader, engine_config: BenchmarkEngineConfig) -> tuple[Any, Any, Any]: + result = SimpleNamespace(llm=llm, exception=None, init_seconds=init_seconds) + return SimpleNamespace(), _ReadyEvent(), result + + monkeypatch.setattr(benchmark_mod, "_build_engine_async", fake_build_engine_async) + + +def _benchmark_corpus(prompts: list[str]) -> BenchmarkCorpus: + return BenchmarkCorpus( + header=TraceHeader(run_id="r", pretrained_model="m", dataset_schema={"col": "string"}), + prompts=[ + BenchmarkPrompt( + row_index=i, + prompt=prompt, + original_sampling_params={"temperature": 0.0, "max_tokens": 8}, + ) + for i, prompt in enumerate(prompts) + ], + ) + + # --------------------------------------------------------------------------- # Data model contracts # --------------------------------------------------------------------------- @@ -117,7 +258,7 @@ def test_engine_config_tolerates_unknown_kwargs(self) -> None: assert cfg.attention_backend == "FLASHINFER" def test_candidate_extra_fields_forbidden(self) -> None: - """``BenchmarkCandidate.extra='forbid'`` — adding a field must update the schema.""" + """``BenchmarkCandidate.extra='forbid'`` means adding a field must update the schema.""" with pytest.raises(ValidationError): BenchmarkCandidate.model_validate({"name": "t", "unknown_field": 42}) @@ -228,12 +369,68 @@ def test_drops_none_valued_overrides(self, header: TraceHeader, empty_base: Benc def test_auto_attention_backend_is_treated_as_unset( self, header: TraceHeader, empty_base: BenchmarkEngineConfig ) -> None: - """``attention_backend='auto'`` means "let vLLM pick" — no attention_config kwarg should appear.""" + """``attention_backend='auto'`` means no attention_config kwarg should appear.""" cfg = empty_base.model_copy(update={"attention_backend": "auto"}) kwargs = _build_vllm_kwargs(header, cfg) assert "attention_config" not in kwargs +# --------------------------------------------------------------------------- +# Runner contracts +# --------------------------------------------------------------------------- + + +class TestRunBenchmark: + def test_n_fanout_dispatches_one_prompt_with_one_completion_per_corpus_prompt( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + llm = _FakeLLM() + processor_calls: list[tuple[int, str]] = [] + _install_runner_fakes(monkeypatch, llm, processor_calls=processor_calls) + + metrics = run_benchmark( + BenchmarkCandidate(name="fanout", batch_dispatch_mode="n_fanout"), + _benchmark_corpus(["same", "same", "same"]), + ) + + assert llm.calls[0]["prompts"] == ["same"] + assert llm.calls[0]["sampling_params"].kwargs["n"] == 3 + assert processor_calls == [ + (0, '{"col": "v0"}'), + (1, '{"col": "v1"}'), + (2, '{"col": "v2"}'), + ] + assert metrics.prompts_attempted == 3 + assert metrics.prompts_accepted == 3 + assert metrics.total_output_tokens == 6 + assert metrics.finish_reason_distribution == {"stop": 3} + + def test_n_fanout_rejects_non_identical_prompts(self, monkeypatch: pytest.MonkeyPatch) -> None: + llm = _FakeLLM() + _install_runner_fakes(monkeypatch, llm) + + with pytest.raises(ValueError, match="requires every corpus prompt to be identical"): + run_benchmark( + BenchmarkCandidate(name="fanout", batch_dispatch_mode="n_fanout"), + _benchmark_corpus(["first", "second"]), + ) + + def test_overlap_savings_are_capped_by_actual_engine_init_time(self, monkeypatch: pytest.MonkeyPatch) -> None: + llm = _FakeLLM() + _install_runner_fakes(monkeypatch, llm, init_seconds=1.25) + monkeypatch.setattr(benchmark_mod.time, "sleep", lambda seconds: None) + monkeypatch.setattr(benchmark_mod.time, "monotonic", lambda: 100.0) + + metrics = run_benchmark( + BenchmarkCandidate(name="baseline"), + _benchmark_corpus(["p0"]), + simulate_training_overlap_seconds=10.0, + ) + + assert metrics.startup_seconds == 0.0 + assert metrics.startup_overlap_savings_seconds == 1.25 + + # --------------------------------------------------------------------------- # Preset contracts # --------------------------------------------------------------------------- @@ -258,7 +455,7 @@ def test_attention_backend_sweep_covers_known_backends(self, empty_base: Benchma class TestBracketedAb: def test_emits_2n_cells_interleaved(self, empty_base: BenchmarkEngineConfig) -> None: - """N baselines + N candidate cells, interleaved by bracket_position.""" + """N baselines + N candidate runs, interleaved by bracket_position.""" cells = bracketed_ab( empty_base, candidate_engine_overrides={"enable_prefix_caching": True}, @@ -274,7 +471,7 @@ def test_emits_2n_cells_interleaved(self, empty_base: BenchmarkEngineConfig) -> def test_spec_ngram_wrapper_applies_speculative_config(self, empty_base: BenchmarkEngineConfig) -> None: cells = bracketed_ab_spec_ngram(empty_base) - # Candidate cells have speculative_config; baselines don't. + # Candidate runs have speculative_config; baselines don't. candidates = [c for c in cells if c.condition_label == "spec_ngram"] baselines = [c for c in cells if c.condition_label == "baseline"] assert len(candidates) == DEFAULT_BRACKETED_AB_N diff --git a/tests/generation/test_vllm_benchmark_analysis.py b/tests/generation/test_vllm_benchmark_analysis.py index 66e25be8b..851330fd1 100644 --- a/tests/generation/test_vllm_benchmark_analysis.py +++ b/tests/generation/test_vllm_benchmark_analysis.py @@ -3,7 +3,7 @@ """Tests for the cluster-conditioned analyzer. -Scope: contracts consumers depend on — cluster partitioning produces +Scope: contracts consumers depend on: cluster partitioning produces the expected shape, refusals fire when sample size is too small, effect-size CI brackets behave correctly, JSON output round-trips. """ @@ -26,7 +26,7 @@ def _metric(name: str, condition: str, *, eff: float, accept: float, wall: float, bracket: int = 0) -> CandidateMetrics: - """Build a CandidateMetrics with minimal scaffolding — used to seed the analyzer.""" + """Build a CandidateMetrics with minimal scaffolding to seed the analyzer.""" return CandidateMetrics( name=name, raw_tok_s=eff / max(accept, 0.01), @@ -52,7 +52,7 @@ def _write_output_dir(tmp_path: Path, cells: list[CandidateMetrics]) -> Path: @pytest.fixture def synthetic_sweep_dir(tmp_path: Path) -> Path: - """6 baselines + 6 spec_ngram cells, both with realistic-noise spread.""" + """6 baselines + 6 spec_ngram runs, both with realistic-noise spread.""" cells = [ # Baselines: ~1500 eff_tok_s, ~0.99 acceptance. *( @@ -75,7 +75,7 @@ def synthetic_sweep_dir(tmp_path: Path) -> Path: class TestWelchTtestCi: def test_clear_difference_excludes_zero(self) -> None: - """Genuine difference in means → CI doesn't bracket 0.""" + """Genuine difference in means means CI doesn't bracket 0.""" res = _welch_ttest_ci([1700, 1720, 1690, 1710, 1705, 1715], [1500, 1510, 1495, 1505, 1502, 1498]) assert res is not None mean_diff, ci_low, ci_high, _df = res @@ -83,7 +83,7 @@ def test_clear_difference_excludes_zero(self) -> None: assert ci_low > 0 # CI excludes 0 def test_identical_means_brackets_zero(self) -> None: - """Same means + nonzero variance → CI brackets 0.""" + """Same means + nonzero variance means CI brackets 0.""" res = _welch_ttest_ci([1500, 1510, 1495, 1505, 1502, 1498], [1500, 1510, 1495, 1505, 1502, 1498]) assert res is not None mean_diff, ci_low, ci_high, _df = res @@ -127,14 +127,14 @@ def test_emits_effect_size_for_non_baseline_conditions(self, synthetic_sweep_dir baseline_agg = next(agg for agg in report.condition_aggregates if agg.condition_label == "baseline") assert spec_agg.pooled_effect_size_vs_baseline is not None assert baseline_agg.pooled_effect_size_vs_baseline is None - # The Δ should be roughly +200 tok/s, CI excludes 0. + # The delta should be roughly +200 tok/s, CI excludes 0. es = spec_agg.pooled_effect_size_vs_baseline assert es.delta_absolute > 0 assert es.ci95_low > 0 def test_refuses_aggregates_below_min_cells(self, tmp_path: Path) -> None: """Conditions with N<6 land in refusals, not aggregates.""" - # 6 baselines + only 3 candidate cells → spec_ngram should be refused. + # 6 baselines + only 3 candidate runs means spec_ngram should be refused. cells = [ *(_metric(f"baseline_{i}", "baseline", eff=1500.0, accept=0.99, wall=130.0) for i in range(6)), *(_metric(f"spec_{i}", "spec_ngram", eff=1700.0, accept=0.99, wall=115.0) for i in range(3)), diff --git a/tools/vllm_benchmark.py b/tools/vllm_benchmark.py index f9be5fe82..a12acca18 100644 --- a/tools/vllm_benchmark.py +++ b/tools/vllm_benchmark.py @@ -4,6 +4,8 @@ r"""vllm-benchmark: drive the vLLM benchmark harness from the command line. +Full usage notes live in docs/developer-guide/vllm-benchmark.md. + Subcommands: list Print the available preset matrices. run CORPUS --output PATH ... Replay CORPUS against the chosen candidates @@ -213,7 +215,7 @@ def compare_cmd(output_path: Path) -> None: f"{m.acceptance_rate:.3f}", f"{m.ttft_p50_ms:.1f}", f"{m.ttft_p99_ms:.1f}", - f"{m.observability.peak_vram_gb:.2f}" if m.observability.peak_vram_gb is not None else "—", + f"{m.observability.peak_vram_gb:.2f}" if m.observability.peak_vram_gb is not None else "--", f"{m.startup_seconds:.1f}", f"{m.prompts_accepted}/{m.prompts_attempted}", ) @@ -235,14 +237,19 @@ def compare_cmd(output_path: Path) -> None: type=click.Choice(["wall_seconds", "acceptance_rate", "auto"]), default="auto", show_default=True, - help="Which per-cell metric to partition cells on. Use 'wall_seconds' for short-context (load-driven bimodality), 'acceptance_rate' for long-output (RNG/scheduler driven), 'auto' to pick whichever has higher pooled CoV.", + help=( + "Which per-candidate-run metric to partition on. Use 'wall_seconds' for " + "short-context workloads, 'acceptance_rate' for long-output workloads, " + "or 'auto' to pick whichever has higher pooled CoV." + ), ) @click.option( + "--min-runs-per-condition", "--min-cells-per-condition", type=int, default=None, show_default="MIN_CELLS_PER_CONDITION (6)", - help="Refuse aggregates for conditions below this N. Brief mandates N≥6.", + help="Refuse aggregates for conditions below this N. Brief mandates N>=6 candidate runs.", ) @click.option( "--json-out", From 0ab073b2234a1f4d6e6bbd416574c8725afe64cb Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 15 Jun 2026 21:56:09 +0000 Subject: [PATCH 14/17] Refine vLLM benchmark harness Signed-off-by: Aaron Gonzales --- docs/developer-guide/vllm-benchmark.md | 7 +- .../generation/vllm_benchmark.py | 124 +++++++++++------- .../generation/vllm_benchmark_analysis.py | 92 ++++++------- .../generation/vllm_benchmark_presets.py | 12 +- .../generation/vllm_benchmark_single_run.py | 2 +- .../generation/vllm_benchmark_wandb.py | 2 +- .../generation/vllm_observability.py | 4 +- tests/generation/test_vllm_benchmark.py | 85 +++++++++--- .../test_vllm_benchmark_analysis.py | 18 +-- tests/generation/test_vllm_benchmark_cli.py | 73 +++++++++++ tools/vllm_benchmark.py | 26 ++-- 11 files changed, 301 insertions(+), 144 deletions(-) create mode 100644 tests/generation/test_vllm_benchmark_cli.py diff --git a/docs/developer-guide/vllm-benchmark.md b/docs/developer-guide/vllm-benchmark.md index f42a528f3..172032116 100644 --- a/docs/developer-guide/vllm-benchmark.md +++ b/docs/developer-guide/vllm-benchmark.md @@ -86,12 +86,11 @@ uv run --frozen --extra cu129 --extra engine --group dev \ --json-out /path/to/analysis.json ``` -The analyzer reports candidate-run aggregates by condition. It keeps the JSON -field name `n_cells` for compatibility; in this context a "cell" means one -candidate run in the benchmark matrix. +The analyzer reports candidate-run aggregates by condition. The report JSON +uses `n_candidate_runs` for aggregate sample counts. Use `--min-runs-per-condition` to raise or lower the refusal threshold. The -older `--min-cells-per-condition` spelling remains accepted for compatibility. +default threshold is 6 candidate runs per condition. ## WandB Sink diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index 92f079bb0..80e44661e 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -39,10 +39,11 @@ import tempfile import threading import time +from collections.abc import Sequence from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, Self from pydantic import BaseModel, ConfigDict, Field @@ -137,6 +138,28 @@ class BenchmarkPrompt(BaseModel): ) +class TracePromptRecord(BaseModel): + """Raw ``kind='record'`` line from a captured trace JSONL.""" + + model_config = ConfigDict(extra="ignore") + + row_index: int + prompt: str + sampling_params: dict[str, Any] = Field(default_factory=dict) + finish_reason: str | None = None + output_text: str = "" + + def to_benchmark_prompt(self) -> BenchmarkPrompt: + """Convert capture-time field names to the benchmark corpus schema.""" + return BenchmarkPrompt( + row_index=self.row_index, + prompt=self.prompt, + original_sampling_params=dict(self.sampling_params), + expected_finish_reason=self.finish_reason, + original_output_text=self.output_text, + ) + + class BenchmarkCorpus(BaseModel): """One captured workload corpus, loaded from a trace JSONL. @@ -170,15 +193,7 @@ def from_trace_jsonl(cls, path: str | Path) -> BenchmarkCorpus: elif kind == "record": if header is None: raise ValueError(f"{path}: record on line {line_no} before any header") - prompts.append( - BenchmarkPrompt( - row_index=int(payload["row_index"]), - prompt=str(payload["prompt"]), - original_sampling_params=dict(payload.get("sampling_params") or {}), - expected_finish_reason=payload.get("finish_reason"), - original_output_text=str(payload.get("output_text", "")), - ), - ) + prompts.append(TracePromptRecord.model_validate(payload).to_benchmark_prompt()) else: raise ValueError(f"{path}: unknown kind={kind!r} on line {line_no}") if header is None: @@ -336,6 +351,52 @@ class BenchmarkCandidate(BaseModel): ), ) + def sampling_kwargs(self, base: dict[str, Any]) -> dict[str, Any]: + """Compose vLLM ``SamplingParams`` kwargs for this candidate. + + The corpus's captured sampling params can carry metadata fields + that vLLM does not accept. Candidate overrides win, and an + override can intentionally reintroduce a normally stripped key. + """ + merged: dict[str, Any] = {**base, **self.sampling_overrides} + for field in _NON_SAMPLING_FIELDS: + if field not in self.sampling_overrides: + merged.pop(field, None) + return merged + + def dispatch_plan(self, prompts: Sequence[str], base_sampling: dict[str, Any]) -> BenchmarkDispatchPlan: + """Build prompt and sampling kwargs for the candidate's dispatch mode.""" + prompt_list = list(prompts) + sampling_kwargs = self.sampling_kwargs(base_sampling) + if self.batch_dispatch_mode == "n_fanout" and len(set(prompt_list)) > 1: + raise ValueError("batch_dispatch_mode='n_fanout' requires every corpus prompt to be identical") + if self.batch_dispatch_mode == "n_fanout": + sampling_kwargs["n"] = max(1, len(prompt_list)) + return BenchmarkDispatchPlan(prompts=prompt_list[:1], sampling_kwargs=sampling_kwargs) + sampling_kwargs.setdefault("n", 1) + return BenchmarkDispatchPlan(prompts=prompt_list, sampling_kwargs=sampling_kwargs) + + +class BenchmarkCandidateDocument(BaseModel): + """JSON document shape consumed by ``tools/vllm_benchmark.py --candidates-file``.""" + + model_config = ConfigDict(extra="forbid") + + candidates: list[BenchmarkCandidate] = Field(description="Candidate configurations to run.") + + @classmethod + def from_json_file(cls, path: str | Path) -> Self: + """Load and validate a candidate-file JSON document.""" + return cls.model_validate_json(Path(path).read_text(encoding="utf-8")) + + +@dataclass(frozen=True) +class BenchmarkDispatchPlan: + """Prompt payload plus sampling kwargs for one candidate dispatch.""" + + prompts: list[str] + sampling_kwargs: dict[str, Any] + class CandidateMetrics(BaseModel): """Measured outputs for one benchmark candidate run. @@ -417,7 +478,7 @@ class CandidateMetrics(BaseModel): observability: GenerationObservability = Field( default_factory=GenerationObservability, description=( - "Cell-level observability snapshot. Composed from PR-A's schema " + "Candidate-run observability snapshot. Composed from PR-A's schema " "so benchmark consumers can read e.g. ``metrics.observability." "peak_vram_gb`` and ``metrics.observability.kv_cache_usage_perc`` " "without the benchmark schema re-defining those fields." @@ -459,25 +520,10 @@ class BenchmarkOutput(BaseModel): # --------------------------------------------------------------------------- -# Sampling-params + percentile helpers +# Percentile + metric helpers # --------------------------------------------------------------------------- -def _merge_sampling_kwargs(base: dict[str, Any], overrides: dict[str, Any]) -> dict[str, Any]: - """Compose ``SamplingParams`` kwargs from corpus default + candidate overrides. - - The corpus's captured ``original_sampling_params`` may carry fields - that ``vllm.SamplingParams`` doesn't accept (e.g. structured-output - presence summaries). Strip those unless the override explicitly - provides them. Overrides win on conflict. - """ - merged: dict[str, Any] = {**base, **overrides} - for field in _NON_SAMPLING_FIELDS: - if field not in overrides: - merged.pop(field, None) - return merged - - def _percentile(values: list[float], pct: float) -> float: """Linear-interpolation percentile (``pct`` in [0, 100]); ``0.0`` on empty input.""" if not values: @@ -687,29 +733,20 @@ def run_benchmark( base_sampling: dict[str, Any] = {} if corpus.prompts: base_sampling = corpus.prompts[0].original_sampling_params - sampling_kwargs = _merge_sampling_kwargs(base_sampling, candidate.sampling_overrides) - prompts: list[str] = [p.prompt for p in corpus.prompts] - if candidate.batch_dispatch_mode == "n_fanout" and len(set(prompts)) > 1: - raise ValueError("batch_dispatch_mode='n_fanout' requires every corpus prompt to be identical") - if candidate.batch_dispatch_mode == "n_fanout": - sampling_kwargs["n"] = max(1, len(prompts)) - dispatch_prompts = prompts[:1] - else: - sampling_kwargs.setdefault("n", 1) - dispatch_prompts = prompts - sampling_params = SamplingParams(**sampling_kwargs) + dispatch = candidate.dispatch_plan([p.prompt for p in corpus.prompts], base_sampling) + sampling_params = SamplingParams(**dispatch.sampling_kwargs) # Dispatch. gen_start = time.perf_counter() outputs: list[Any] = ( list( llm.generate( - prompts=dispatch_prompts, + prompts=dispatch.prompts, sampling_params=sampling_params, lora_request=lora_request, ) ) - if dispatch_prompts + if dispatch.prompts else [] ) total_wall = max(time.perf_counter() - gen_start, 0.0) @@ -846,9 +883,8 @@ def run_benchmark_in_subprocess( interpreter; the OS reclaims everything on child exit. """ candidate_json = candidate.model_dump_json() - with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False, encoding="utf-8") as result_fh: - result_path = Path(result_fh.name) - try: + with tempfile.TemporaryDirectory(prefix="nss-vllm-benchmark-") as result_dir: + result_path = Path(result_dir) / "result.json" completed = subprocess.run( [ sys.executable, @@ -880,5 +916,3 @@ def run_benchmark_in_subprocess( ) metrics = CandidateMetrics.model_validate_json(result_path.read_text(encoding="utf-8")) return SubprocessRunResult(metrics=metrics) - finally: - result_path.unlink(missing_ok=True) diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py index 39f00bdf1..a2d0153d5 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py @@ -23,7 +23,7 @@ Cluster count is selected via silhouette score (post-hoc, k in [2, 4]). Refuses to compute delta-style aggregates when a condition has fewer than -:data:`MIN_CELLS_PER_CONDITION` candidate runs - single-run measurements +:data:`MIN_CANDIDATE_RUNS_PER_CONDITION` candidate runs - single-run measurements should never drive promote/reject decisions on this stack. Effect-size + 95% CI reporting lives in this module too - see @@ -45,7 +45,7 @@ ClusterSignal = Literal["wall_seconds", "acceptance_rate", "auto"] -MIN_CELLS_PER_CONDITION: int = 6 +MIN_CANDIDATE_RUNS_PER_CONDITION: int = 6 """Minimum candidate runs per condition before delta-style aggregates are computed. Matches :data:`DEFAULT_BRACKETED_AB_N`. Below this threshold the @@ -80,7 +80,7 @@ class ClusterStats(BaseModel): model_config = ConfigDict(extra="forbid") cluster_id: int - n_cells: int + n_candidate_runs: int signal_mean: float signal_stddev: float signal_cov: float @@ -118,7 +118,7 @@ class ConditionClusterAggregate(BaseModel): condition_label: str cluster_id: int - n_cells: int + n_candidate_runs: int mean_effective_tok_s: float stddev_effective_tok_s: float cov_effective_tok_s: float @@ -133,7 +133,7 @@ class ConditionAggregate(BaseModel): model_config = ConfigDict(extra="forbid") condition_label: str - n_cells: int + n_candidate_runs: int pooled_mean_effective_tok_s: float pooled_stddev_effective_tok_s: float pooled_cov_effective_tok_s: float @@ -151,7 +151,7 @@ class AnalysisReport(BaseModel): cluster_signal: str n_clusters: int - n_cells: int + n_candidate_runs: int cluster_assignments: list[ClusterAssignment] cluster_stats: list[ClusterStats] condition_aggregates: list[ConditionAggregate] @@ -160,7 +160,7 @@ class AnalysisReport(BaseModel): def to_markdown_summary(self) -> str: """Render a human-readable summary.""" lines: list[str] = [ - f"# Cluster-conditioned analysis ({self.n_cells} candidate runs, k={self.n_clusters})", + f"# Cluster-conditioned analysis ({self.n_candidate_runs} candidate runs, k={self.n_clusters})", "", f"Cluster signal: `{self.cluster_signal}`", "", @@ -171,12 +171,12 @@ def to_markdown_summary(self) -> str: ] for cs in self.cluster_stats: lines.append( - f"| {cs.cluster_id} | {cs.n_cells} | {cs.signal_mean:.4f} | " + f"| {cs.cluster_id} | {cs.n_candidate_runs} | {cs.signal_mean:.4f} | " f"{cs.signal_stddev:.4f} | {cs.signal_cov * 100:.2f}% |" ) lines.extend(("", "## Per-condition aggregates", "")) for agg in self.condition_aggregates: - lines.append(f"### `{agg.condition_label}` (n={agg.n_cells})") + lines.append(f"### `{agg.condition_label}` (n={agg.n_candidate_runs})") lines.append("") lines.append( f"- Pooled: eff_tok_s={agg.pooled_mean_effective_tok_s:.1f} " @@ -196,7 +196,7 @@ def to_markdown_summary(self) -> str: lines.append("- In-cluster:") for ic in agg.in_cluster: lines.append( - f" - cluster {ic.cluster_id}: n={ic.n_cells}, " + f" - cluster {ic.cluster_id}: n={ic.n_candidate_runs}, " f"eff_tok_s={ic.mean_effective_tok_s:.1f} " f"+/- {ic.stddev_effective_tok_s:.1f} " f"(CoV {ic.cov_effective_tok_s * 100:.2f}%); " @@ -224,12 +224,12 @@ def to_markdown_summary(self) -> str: # --------------------------------------------------------------------------- -def _signal_value(cell: CandidateMetrics, signal: str) -> float: +def _signal_value(candidate_run: CandidateMetrics, signal: str) -> float: """Extract the signal value for clustering; raises on unknown signal.""" if signal == "wall_seconds": - return cell.total_wall_seconds + return candidate_run.total_wall_seconds if signal == "acceptance_rate": - return cell.acceptance_rate + return candidate_run.acceptance_rate raise ValueError(f"unknown cluster signal: {signal!r}") @@ -243,16 +243,16 @@ def _pooled_cov(values: list[float]) -> tuple[float, float, float]: return mean, stddev, cov -def _auto_select_signal(cells: list[CandidateMetrics]) -> str: +def _auto_select_signal(candidate_runs: list[CandidateMetrics]) -> str: """Pick whichever signal has higher pooled CoV across the candidate runs. Defaults to ``wall_seconds`` on ties (short-context bimodality is the more common workload shape). """ - if not cells: + if not candidate_runs: return "wall_seconds" - _, _, cov_wall = _pooled_cov([c.total_wall_seconds for c in cells]) - _, _, cov_acc = _pooled_cov([c.acceptance_rate for c in cells]) + _, _, cov_wall = _pooled_cov([c.total_wall_seconds for c in candidate_runs]) + _, _, cov_acc = _pooled_cov([c.acceptance_rate for c in candidate_runs]) return "acceptance_rate" if cov_acc > cov_wall else "wall_seconds" @@ -372,46 +372,46 @@ def _effect_size( # --------------------------------------------------------------------------- -def load_cells(output_dir: Path) -> list[CandidateMetrics]: +def load_candidate_runs(output_dir: Path) -> list[CandidateMetrics]: """Read every ``BenchmarkOutput`` JSON in ``output_dir``, flatten candidate runs. Subdirectories are NOT recursed. Callers wanting cross-dataset analysis should invoke once per dataset dir. """ - cells: list[CandidateMetrics] = [] + candidate_runs: list[CandidateMetrics] = [] for path in sorted(output_dir.glob("*.json")): try: doc = BenchmarkOutput.model_validate_json(path.read_text(encoding="utf-8")) except (json.JSONDecodeError, ValueError): continue - cells.extend(doc.candidates) - return cells + candidate_runs.extend(doc.candidates) + return candidate_runs def analyze( output_dir: Path, cluster_signal: ClusterSignal = "auto", - min_cells_per_condition: int = MIN_CELLS_PER_CONDITION, + min_candidate_runs_per_condition: int = MIN_CANDIDATE_RUNS_PER_CONDITION, ) -> AnalysisReport: """Full pipeline: load -> cluster -> per-condition aggregate -> effect-size -> report.""" - cells = load_cells(output_dir) - if not cells: + candidate_runs = load_candidate_runs(output_dir) + if not candidate_runs: raise ValueError(f"No BenchmarkOutput JSONs found under {output_dir}") - resolved_signal = _auto_select_signal(cells) if cluster_signal == "auto" else cluster_signal - values = np.array([_signal_value(c, resolved_signal) for c in cells], dtype=float) + resolved_signal = _auto_select_signal(candidate_runs) if cluster_signal == "auto" else cluster_signal + values = np.array([_signal_value(candidate_run, resolved_signal) for candidate_run in candidate_runs], dtype=float) n_clusters = _select_n_clusters(values) labels = _assign_clusters(values, n_clusters) assignments = [ ClusterAssignment( - candidate_name=cell.name, - condition_label=cell.condition_label, - bracket_position=cell.bracket_position, + candidate_name=candidate_run.name, + condition_label=candidate_run.condition_label, + bracket_position=candidate_run.bracket_position, cluster_id=int(lbl), signal_value=float(val), ) - for cell, lbl, val in zip(cells, labels, values, strict=True) + for candidate_run, lbl, val in zip(candidate_runs, labels, values, strict=True) ] cluster_stats: list[ClusterStats] = [] @@ -421,7 +421,7 @@ def analyze( cluster_stats.append( ClusterStats( cluster_id=cid, - n_cells=len(cluster_values), + n_candidate_runs=len(cluster_values), signal_mean=mean, signal_stddev=stddev, signal_cov=cov, @@ -429,23 +429,23 @@ def analyze( ) by_condition: dict[str, list[tuple[int, CandidateMetrics]]] = {} - for cell, lbl in zip(cells, labels, strict=True): - by_condition.setdefault(cell.condition_label, []).append((int(lbl), cell)) + for candidate_run, lbl in zip(candidate_runs, labels, strict=True): + by_condition.setdefault(candidate_run.condition_label, []).append((int(lbl), candidate_run)) baseline_pooled_eff: list[float] = [] baseline_per_cluster_eff: dict[int, list[float]] = {} - for lbl, cell in by_condition.get("baseline", []): - baseline_pooled_eff.append(cell.effective_tok_s) - baseline_per_cluster_eff.setdefault(lbl, []).append(cell.effective_tok_s) + for lbl, candidate_run in by_condition.get("baseline", []): + baseline_pooled_eff.append(candidate_run.effective_tok_s) + baseline_per_cluster_eff.setdefault(lbl, []).append(candidate_run.effective_tok_s) aggregates: list[ConditionAggregate] = [] refusals: list[str] = [] for condition in sorted(by_condition): labeled = by_condition[condition] - if len(labeled) < min_cells_per_condition: + if len(labeled) < min_candidate_runs_per_condition: refusals.append( f"condition {condition!r} has only {len(labeled)} candidate runs; " - f"need >={min_cells_per_condition} - refusing aggregate" + f"need >={min_candidate_runs_per_condition} - refusing aggregate" ) continue pooled_eff = [c.effective_tok_s for _, c in labeled] @@ -455,12 +455,12 @@ def analyze( in_cluster: list[ConditionClusterAggregate] = [] for cid in range(n_clusters): - cluster_cells = [c for lbl, c in labeled if lbl == cid] - if not cluster_cells: + cluster_candidate_runs = [c for lbl, c in labeled if lbl == cid] + if not cluster_candidate_runs: continue - ic_eff = [c.effective_tok_s for c in cluster_cells] - ic_acc = [c.acceptance_rate for c in cluster_cells] - ic_raw = [c.raw_tok_s for c in cluster_cells] + ic_eff = [c.effective_tok_s for c in cluster_candidate_runs] + ic_acc = [c.acceptance_rate for c in cluster_candidate_runs] + ic_raw = [c.raw_tok_s for c in cluster_candidate_runs] mean_eff, stddev_eff, cov_eff = _pooled_cov(ic_eff) ic_effect: EffectSize | None = None if condition != "baseline" and cid in baseline_per_cluster_eff: @@ -476,7 +476,7 @@ def analyze( ConditionClusterAggregate( condition_label=condition, cluster_id=cid, - n_cells=len(cluster_cells), + n_candidate_runs=len(cluster_candidate_runs), mean_effective_tok_s=mean_eff, stddev_effective_tok_s=stddev_eff, cov_effective_tok_s=cov_eff, @@ -498,7 +498,7 @@ def analyze( aggregates.append( ConditionAggregate( condition_label=condition, - n_cells=len(labeled), + n_candidate_runs=len(labeled), pooled_mean_effective_tok_s=eff_mean, pooled_stddev_effective_tok_s=eff_stddev, pooled_cov_effective_tok_s=eff_cov, @@ -513,7 +513,7 @@ def analyze( return AnalysisReport( cluster_signal=resolved_signal, n_clusters=n_clusters, - n_cells=len(cells), + n_candidate_runs=len(candidate_runs), cluster_assignments=assignments, cluster_stats=cluster_stats, condition_aggregates=aggregates, diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py index 91f7c757e..b0ce05e02 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py @@ -22,7 +22,7 @@ from collections.abc import Callable from typing import Any -from .vllm_benchmark import BenchmarkCandidate, BenchmarkEngineConfig +from .vllm_benchmark import BatchDispatchMode, BenchmarkCandidate, BenchmarkEngineConfig PresetFn = Callable[[BenchmarkEngineConfig], list[BenchmarkCandidate]] """A preset resolves the corpus-default engine config into candidate runs.""" @@ -205,7 +205,7 @@ def bracketed_ab( *, candidate_engine_overrides: dict[str, Any] | None = None, candidate_sampling_overrides: dict[str, Any] | None = None, - candidate_batch_dispatch_mode: str | None = None, + candidate_batch_dispatch_mode: BatchDispatchMode | None = None, condition_label: str, n_samples_per_condition: int = DEFAULT_BRACKETED_AB_N, ) -> list[BenchmarkCandidate]: @@ -226,9 +226,9 @@ def bracketed_ab( """ engine_overrides = candidate_engine_overrides or {} sampling_extra = candidate_sampling_overrides or {} - cells: list[BenchmarkCandidate] = [] + candidate_runs: list[BenchmarkCandidate] = [] for i in range(n_samples_per_condition): - cells.append( + candidate_runs.append( BenchmarkCandidate( name=f"bracket_baseline_{i}", engine_config=_with_default_max_model_len(base), @@ -246,8 +246,8 @@ def bracketed_ab( } if candidate_batch_dispatch_mode is not None: cand_kwargs["batch_dispatch_mode"] = candidate_batch_dispatch_mode - cells.append(BenchmarkCandidate(**cand_kwargs)) - return cells + candidate_runs.append(BenchmarkCandidate(**cand_kwargs)) + return candidate_runs # Phase B matrix-condition wrappers. Each yields a 2N-run sequence diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py index bda272185..1b8e546c8 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Subprocess entry point for one isolated benchmark cell. +"""Subprocess entry point for one isolated benchmark candidate run. Invoked by :func:`vllm_benchmark.run_benchmark_in_subprocess` via ``python -m nemo_safe_synthesizer.generation.vllm_benchmark_single_run``. diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py index 7610bef4c..34c60137f 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_wandb.py @@ -53,7 +53,7 @@ def resolve_sweep_id() -> str: return os.environ.get("WANDB_RUN_GROUP") or f"sweep-{datetime.now(timezone.utc):%Y%m%d-%H%M%S}" -def init_cell_run( +def init_candidate_run( *, candidate_name: str, candidate_idx: int, diff --git a/src/nemo_safe_synthesizer/generation/vllm_observability.py b/src/nemo_safe_synthesizer/generation/vllm_observability.py index 8ea043920..0ef289528 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_observability.py +++ b/src/nemo_safe_synthesizer/generation/vllm_observability.py @@ -5,7 +5,7 @@ Schema-frozen generation-observability events emitted by ``VllmBackend.generate()`` and consumed by downstream observability surfaces (structured logs, wandb, -the benchmark harness's per-cell aggregator). +the benchmark harness's candidate-run aggregator). Four primitives, all degraded-mode by design: @@ -105,7 +105,7 @@ class GenerationObservability(BaseModel): ``logger.runtime.info(...)`` like the rest of PR-1's trace telemetry). - Wandb (when a run is active) — logged to the current wandb run. - - The benchmark harness's per-cell aggregator (composes this into + - The benchmark harness's candidate-run aggregator (composes this into its richer ``CandidateMetrics`` schema). Every measurement field is optional; producers should populate what diff --git a/tests/generation/test_vllm_benchmark.py b/tests/generation/test_vllm_benchmark.py index 14755fc7f..7a2657fdc 100644 --- a/tests/generation/test_vllm_benchmark.py +++ b/tests/generation/test_vllm_benchmark.py @@ -23,6 +23,7 @@ import nemo_safe_synthesizer.generation.vllm_benchmark as benchmark_mod from nemo_safe_synthesizer.generation.vllm_benchmark import ( BenchmarkCandidate, + BenchmarkCandidateDocument, BenchmarkCorpus, BenchmarkEngineConfig, BenchmarkOutput, @@ -30,9 +31,9 @@ CandidateMetrics, SubprocessRunResult, TraceHeader, + TracePromptRecord, _build_vllm_kwargs, _extract_ttft_ms, - _merge_sampling_kwargs, _parse_error_class, _percentile, _truncate_stderr, @@ -262,6 +263,19 @@ def test_candidate_extra_fields_forbidden(self) -> None: with pytest.raises(ValidationError): BenchmarkCandidate.model_validate({"name": "t", "unknown_field": 42}) + def test_candidate_document_loads_candidate_file(self, tmp_path: Any) -> None: + path = tmp_path / "candidates.json" + path.write_text( + json.dumps({"candidates": [{"name": "baseline", "sampling_overrides": {"seed": 42}}]}), + encoding="utf-8", + ) + + doc = BenchmarkCandidateDocument.from_json_file(path) + + assert len(doc.candidates) == 1 + assert doc.candidates[0].name == "baseline" + assert doc.candidates[0].sampling_overrides == {"seed": 42} + # --------------------------------------------------------------------------- # Corpus loader @@ -269,6 +283,24 @@ def test_candidate_extra_fields_forbidden(self) -> None: class TestBenchmarkCorpus: + def test_trace_prompt_record_converts_capture_keys(self) -> None: + record = TracePromptRecord.model_validate( + { + "row_index": "3", + "prompt": "p", + "sampling_params": {"temperature": 0.7}, + "finish_reason": "stop", + "output_text": '{"col": "x"}', + }, + ) + + prompt = record.to_benchmark_prompt() + + assert prompt.row_index == 3 + assert prompt.original_sampling_params == {"temperature": 0.7} + assert prompt.expected_finish_reason == "stop" + assert prompt.original_output_text == '{"col": "x"}' + def test_loads_jsonl_with_header_and_records(self, tmp_path: Any) -> None: path = tmp_path / "trace.jsonl" lines = [ @@ -300,18 +332,31 @@ def test_percentile_linear_interp(self) -> None: assert _percentile([1.0], 50) == 1.0 assert _percentile([], 50) == 0.0 # empty returns 0, not raises - def test_merge_sampling_strips_non_sampling_fields(self) -> None: + def test_candidate_sampling_kwargs_strips_non_sampling_fields(self) -> None: """``structured_outputs`` is capture-time metadata, not a SamplingParams kwarg.""" - merged = _merge_sampling_kwargs( - {"temperature": 0.7, "top_p": 0.9, "structured_outputs": "json"}, - {"seed": 42}, - ) + candidate = BenchmarkCandidate(name="baseline", sampling_overrides={"seed": 42}) + merged = candidate.sampling_kwargs({"temperature": 0.7, "top_p": 0.9, "structured_outputs": "json"}) assert merged == {"temperature": 0.7, "top_p": 0.9, "seed": 42} - def test_merge_overrides_win_on_conflict(self) -> None: - merged = _merge_sampling_kwargs({"temperature": 0.7}, {"temperature": 0.0}) + def test_candidate_sampling_overrides_win_on_conflict(self) -> None: + candidate = BenchmarkCandidate(name="baseline", sampling_overrides={"temperature": 0.0}) + merged = candidate.sampling_kwargs({"temperature": 0.7}) assert merged["temperature"] == 0.0 + def test_candidate_dispatch_plan_replicates_prompts(self) -> None: + candidate = BenchmarkCandidate(name="baseline") + plan = candidate.dispatch_plan(["a", "b"], {"temperature": 0.0}) + + assert plan.prompts == ["a", "b"] + assert plan.sampling_kwargs == {"temperature": 0.0, "n": 1} + + def test_candidate_dispatch_plan_fans_out_identical_prompts(self) -> None: + candidate = BenchmarkCandidate(name="fanout", batch_dispatch_mode="n_fanout") + plan = candidate.dispatch_plan(["same", "same", "same"], {"temperature": 0.0}) + + assert plan.prompts == ["same"] + assert plan.sampling_kwargs == {"temperature": 0.0, "n": 3} + @pytest.mark.parametrize( ("metrics_obj", "expected"), [ @@ -454,26 +499,26 @@ def test_attention_backend_sweep_covers_known_backends(self, empty_base: Benchma class TestBracketedAb: - def test_emits_2n_cells_interleaved(self, empty_base: BenchmarkEngineConfig) -> None: + def test_emits_2n_candidate_runs_interleaved(self, empty_base: BenchmarkEngineConfig) -> None: """N baselines + N candidate runs, interleaved by bracket_position.""" - cells = bracketed_ab( + candidate_runs = bracketed_ab( empty_base, candidate_engine_overrides={"enable_prefix_caching": True}, condition_label="prefix_on", n_samples_per_condition=3, ) - assert len(cells) == 6 + assert len(candidate_runs) == 6 # Even positions are baseline; odd positions are the candidate. - for i, cell in enumerate(cells): + for i, candidate_run in enumerate(candidate_runs): expected_label = "baseline" if i % 2 == 0 else "prefix_on" - assert cell.condition_label == expected_label - assert cell.bracket_position == i + assert candidate_run.condition_label == expected_label + assert candidate_run.bracket_position == i def test_spec_ngram_wrapper_applies_speculative_config(self, empty_base: BenchmarkEngineConfig) -> None: - cells = bracketed_ab_spec_ngram(empty_base) + candidate_runs = bracketed_ab_spec_ngram(empty_base) # Candidate runs have speculative_config; baselines don't. - candidates = [c for c in cells if c.condition_label == "spec_ngram"] - baselines = [c for c in cells if c.condition_label == "baseline"] + candidates = [c for c in candidate_runs if c.condition_label == "spec_ngram"] + baselines = [c for c in candidate_runs if c.condition_label == "baseline"] assert len(candidates) == DEFAULT_BRACKETED_AB_N assert len(baselines) == DEFAULT_BRACKETED_AB_N for c in candidates: @@ -482,9 +527,9 @@ def test_spec_ngram_wrapper_applies_speculative_config(self, empty_base: Benchma for b in baselines: assert b.engine_config.speculative_config is None - def test_all_cells_seed_pinned(self, empty_base: BenchmarkEngineConfig) -> None: - cells = bracketed_ab_spec_ngram(empty_base) - for c in cells: + def test_all_candidate_runs_seed_pinned(self, empty_base: BenchmarkEngineConfig) -> None: + candidate_runs = bracketed_ab_spec_ngram(empty_base) + for c in candidate_runs: assert c.sampling_overrides.get("seed") == DEFAULT_BENCHMARK_SEED diff --git a/tests/generation/test_vllm_benchmark_analysis.py b/tests/generation/test_vllm_benchmark_analysis.py index 851330fd1..b5793292b 100644 --- a/tests/generation/test_vllm_benchmark_analysis.py +++ b/tests/generation/test_vllm_benchmark_analysis.py @@ -43,9 +43,9 @@ def _metric(name: str, condition: str, *, eff: float, accept: float, wall: float ) -def _write_output_dir(tmp_path: Path, cells: list[CandidateMetrics]) -> Path: +def _write_output_dir(tmp_path: Path, candidate_runs: list[CandidateMetrics]) -> Path: """Write a single BenchmarkOutput JSON to ``tmp_path / out.json``.""" - out = BenchmarkOutput(corpus_run_id="r1", corpus_size=143, candidates=cells) + out = BenchmarkOutput(corpus_run_id="r1", corpus_size=143, candidates=candidate_runs) (tmp_path / "out.json").write_text(out.model_dump_json(), encoding="utf-8") return tmp_path @@ -53,7 +53,7 @@ def _write_output_dir(tmp_path: Path, cells: list[CandidateMetrics]) -> Path: @pytest.fixture def synthetic_sweep_dir(tmp_path: Path) -> Path: """6 baselines + 6 spec_ngram runs, both with realistic-noise spread.""" - cells = [ + candidate_runs = [ # Baselines: ~1500 eff_tok_s, ~0.99 acceptance. *( _metric(f"baseline_{i}", "baseline", eff=1500 + i * 5, accept=0.99, wall=130.0, bracket=2 * i) @@ -65,7 +65,7 @@ def synthetic_sweep_dir(tmp_path: Path) -> Path: for i in range(6) ), ] - return _write_output_dir(tmp_path, cells) + return _write_output_dir(tmp_path, candidate_runs) # --------------------------------------------------------------------------- @@ -111,13 +111,13 @@ def test_underdetermined_returns_none(self, cand: list[float], base: list[float] class TestAnalyze: def test_partitions_and_aggregates(self, synthetic_sweep_dir: Path) -> None: report = analyze(synthetic_sweep_dir, cluster_signal="wall_seconds") - assert report.n_cells == 12 + assert report.n_candidate_runs == 12 # Two conditions present. labels = {agg.condition_label for agg in report.condition_aggregates} assert labels == {"baseline", "spec_ngram"} # Each condition has the expected pooled aggregate. spec_agg = next(agg for agg in report.condition_aggregates if agg.condition_label == "spec_ngram") - assert spec_agg.n_cells == 6 + assert spec_agg.n_candidate_runs == 6 assert spec_agg.pooled_mean_effective_tok_s == pytest.approx(1712.5, abs=0.1) def test_emits_effect_size_for_non_baseline_conditions(self, synthetic_sweep_dir: Path) -> None: @@ -132,14 +132,14 @@ def test_emits_effect_size_for_non_baseline_conditions(self, synthetic_sweep_dir assert es.delta_absolute > 0 assert es.ci95_low > 0 - def test_refuses_aggregates_below_min_cells(self, tmp_path: Path) -> None: + def test_refuses_aggregates_below_min_candidate_runs(self, tmp_path: Path) -> None: """Conditions with N<6 land in refusals, not aggregates.""" # 6 baselines + only 3 candidate runs means spec_ngram should be refused. - cells = [ + candidate_runs = [ *(_metric(f"baseline_{i}", "baseline", eff=1500.0, accept=0.99, wall=130.0) for i in range(6)), *(_metric(f"spec_{i}", "spec_ngram", eff=1700.0, accept=0.99, wall=115.0) for i in range(3)), ] - out_dir = _write_output_dir(tmp_path, cells) + out_dir = _write_output_dir(tmp_path, candidate_runs) report = analyze(out_dir) labels = {agg.condition_label for agg in report.condition_aggregates} assert "baseline" in labels diff --git a/tests/generation/test_vllm_benchmark_cli.py b/tests/generation/test_vllm_benchmark_cli.py new file mode 100644 index 000000000..b6dc3ba14 --- /dev/null +++ b/tests/generation/test_vllm_benchmark_cli.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI regression tests for the vLLM benchmark tool.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +from click.testing import CliRunner + +from nemo_safe_synthesizer.generation.vllm_benchmark import BenchmarkOutput, CandidateMetrics + + +def _load_cli(): + tool_path = Path(__file__).resolve().parents[2] / "tools" / "vllm_benchmark.py" + spec = importlib.util.spec_from_file_location("vllm_benchmark_tool", tool_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load {tool_path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module.cli + + +cli = _load_cli() + + +def _metric(name: str, condition: str, *, eff: float, wall: float, bracket: int) -> CandidateMetrics: + return CandidateMetrics( + name=name, + raw_tok_s=eff / 0.99, + acceptance_rate=0.99, + effective_tok_s=eff, + ttft_p50_ms=0.0, + ttft_p99_ms=0.0, + prompts_attempted=4, + prompts_accepted=4, + total_output_tokens=1000, + total_wall_seconds=wall, + condition_label=condition, + bracket_position=bracket, + ) + + +def _write_output(path: Path) -> None: + candidate_runs = [ + *(_metric(f"baseline_{i}", "baseline", eff=1500.0 + i, wall=130.0 + i, bracket=2 * i) for i in range(6)), + *( + _metric(f"spec_ngram_{i}", "spec_ngram", eff=1700.0 + i, wall=115.0 + i, bracket=2 * i + 1) + for i in range(6) + ), + ] + path.write_text( + BenchmarkOutput(corpus_run_id="cli-test", corpus_size=4, candidates=candidate_runs).model_dump_json(), + encoding="utf-8", + ) + + +def test_analyze_accepts_min_runs_option(tmp_path: Path) -> None: + """The public min-run option should bind to the analyzer callback.""" + _write_output(tmp_path / "out.json") + + result = CliRunner().invoke( + cli, + ["analyze", str(tmp_path), "--cluster-signal", "wall_seconds", "--min-runs-per-condition", "6"], + color=False, + ) + + assert result.exit_code == 0, result.output + assert "Cluster-conditioned analysis" in result.output diff --git a/tools/vllm_benchmark.py b/tools/vllm_benchmark.py index a12acca18..9547f79c4 100644 --- a/tools/vllm_benchmark.py +++ b/tools/vllm_benchmark.py @@ -27,7 +27,6 @@ from __future__ import annotations -import json from datetime import datetime, timezone from pathlib import Path from typing import cast @@ -38,6 +37,7 @@ from nemo_safe_synthesizer.generation.vllm_benchmark import ( BenchmarkCandidate, + BenchmarkCandidateDocument, BenchmarkCorpus, BenchmarkEngineConfig, BenchmarkOutput, @@ -46,7 +46,7 @@ ) from nemo_safe_synthesizer.generation.vllm_benchmark_presets import PRESETS from nemo_safe_synthesizer.generation.vllm_benchmark_wandb import ( - init_cell_run, + init_candidate_run, log_and_finish, resolve_sweep_id, ) @@ -83,8 +83,10 @@ def _resolve_candidates( raise click.UsageError(f"Unknown preset {preset_name!r}; available: {sorted(PRESETS)}") return PRESETS[preset_name](base) if candidates_file: - doc = json.loads(Path(candidates_file).read_text(encoding="utf-8")) - return [BenchmarkCandidate.model_validate(c) for c in doc["candidates"]] + try: + return BenchmarkCandidateDocument.from_json_file(candidates_file).candidates + except ValueError as exc: + raise click.UsageError(f"Invalid candidates file {candidates_file}: {exc}") from exc raise click.UsageError("One of --candidates or --candidates-file is required.") @@ -142,7 +144,7 @@ def run_cmd( sweep_id = resolve_sweep_id() for idx, candidate in enumerate(candidates, start=1): console.print(f"[{idx}/{len(candidates)}] running candidate {candidate.name!r}") - wandb_run = init_cell_run( + wandb_run = init_candidate_run( candidate_name=candidate.name, candidate_idx=idx, total=len(candidates), @@ -245,10 +247,10 @@ def compare_cmd(output_path: Path) -> None: ) @click.option( "--min-runs-per-condition", - "--min-cells-per-condition", + "min_candidate_runs_per_condition", type=int, default=None, - show_default="MIN_CELLS_PER_CONDITION (6)", + show_default="MIN_CANDIDATE_RUNS_PER_CONDITION (6)", help="Refuse aggregates for conditions below this N. Brief mandates N>=6 candidate runs.", ) @click.option( @@ -260,12 +262,12 @@ def compare_cmd(output_path: Path) -> None: def analyze_cmd( output_dir: Path, cluster_signal: str, - min_cells_per_condition: int | None, + min_candidate_runs_per_condition: int | None, json_out: Path | None, ) -> None: """Cluster-conditioned analysis across every BenchmarkOutput JSON in OUTPUT_DIR.""" from nemo_safe_synthesizer.generation.vllm_benchmark_analysis import ( - MIN_CELLS_PER_CONDITION, + MIN_CANDIDATE_RUNS_PER_CONDITION, ClusterSignal, analyze, ) @@ -273,7 +275,11 @@ def analyze_cmd( report = analyze( output_dir, cluster_signal=cast(ClusterSignal, cluster_signal), - min_cells_per_condition=MIN_CELLS_PER_CONDITION if min_cells_per_condition is None else min_cells_per_condition, + min_candidate_runs_per_condition=( + MIN_CANDIDATE_RUNS_PER_CONDITION + if min_candidate_runs_per_condition is None + else min_candidate_runs_per_condition + ), ) console.print(report.to_markdown_summary()) if json_out is not None: From 241bc934ab0123e561078bc7ce09f8dd7ad83045 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Tue, 16 Jun 2026 18:56:15 +0000 Subject: [PATCH 15/17] Share benchmark preset definitions Signed-off-by: Aaron Gonzales --- src/nemo_safe_synthesizer/config/generate.py | 23 +++++++- .../generation/vllm_benchmark_analysis.py | 3 +- .../generation/vllm_benchmark_presets.py | 57 +++++++++++++------ tests/generation/test_vllm_benchmark.py | 23 +++++++- 4 files changed, 84 insertions(+), 22 deletions(-) diff --git a/src/nemo_safe_synthesizer/config/generate.py b/src/nemo_safe_synthesizer/config/generate.py index 47b1e05aa..22a7e8870 100644 --- a/src/nemo_safe_synthesizer/config/generate.py +++ b/src/nemo_safe_synthesizer/config/generate.py @@ -5,7 +5,7 @@ import warnings from collections.abc import Mapping -from typing import Annotated, Any, ClassVar, Literal, Self +from typing import Annotated, Any, ClassVar, Literal, Self, cast, get_args from pydantic import ( BaseModel, @@ -27,15 +27,32 @@ ResolvedStructuredGenerationSchemaMethod = Literal["regex", "json_schema", "structural_tag"] StructuredGenerationBackend = Literal["auto", "xgrammar", "guidance", "outlines", "lm-format-enforcer"] +SUPPORTED_STRUCTURED_GENERATION_BACKENDS = cast( + tuple[StructuredGenerationBackend, ...], + get_args(StructuredGenerationBackend), +) +"""Structured-output backend values accepted by generation config.""" + STRUCTURAL_TAG_COMPATIBLE_BACKENDS = frozenset({"auto", "xgrammar"}) +COMMON_ATTENTION_BACKENDS: tuple[str, ...] = ( + "FLASHINFER", + "FLASH_ATTN", + "TORCH_SDPA", + "TRITON_ATTN", + "FLEX_ATTENTION", +) +"""Common vLLM attention backend values accepted by ``generation.attention_backend``.""" + __all__ = [ + "COMMON_ATTENTION_BACKENDS", "GenerateParameters", "ResolvedStructuredGenerationSchemaMethod", "StructuredGenerationParameters", "StructuredGenerationBackend", "StructuredGenerationSchemaMethod", "STRUCTURAL_TAG_COMPATIBLE_BACKENDS", + "SUPPORTED_STRUCTURED_GENERATION_BACKENDS", "ValidationParameters", "resolve_structured_generation_schema_method", "structural_tag_backend_error_message", @@ -265,8 +282,8 @@ class GenerateParameters(Parameters, BaseModel): Field( title="attention_backend", description=( - "The attention backend for the vLLM engine. Common values: 'FLASHINFER', " - "'FLASH_ATTN', 'TRITON_ATTN', 'FLEX_ATTENTION'. " + "The attention backend for the vLLM engine. Common values: " + f"{', '.join(repr(backend) for backend in COMMON_ATTENTION_BACKENDS)}. " "If ``None`` or 'auto', vLLM will auto-select the best available backend." ), ), diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py index a2d0153d5..e568bb62e 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_analysis.py @@ -42,10 +42,11 @@ from pydantic import BaseModel, ConfigDict, Field from .vllm_benchmark import BenchmarkOutput, CandidateMetrics +from .vllm_benchmark_presets import DEFAULT_BRACKETED_AB_N ClusterSignal = Literal["wall_seconds", "acceptance_rate", "auto"] -MIN_CANDIDATE_RUNS_PER_CONDITION: int = 6 +MIN_CANDIDATE_RUNS_PER_CONDITION: int = DEFAULT_BRACKETED_AB_N """Minimum candidate runs per condition before delta-style aggregates are computed. Matches :data:`DEFAULT_BRACKETED_AB_N`. Below this threshold the diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py index b0ce05e02..4b68e85e3 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark_presets.py @@ -20,8 +20,14 @@ import json from collections.abc import Callable +from dataclasses import dataclass from typing import Any +from ..config.generate import ( + COMMON_ATTENTION_BACKENDS, + SUPPORTED_STRUCTURED_GENERATION_BACKENDS, + StructuredGenerationBackend, +) from .vllm_benchmark import BatchDispatchMode, BenchmarkCandidate, BenchmarkEngineConfig PresetFn = Callable[[BenchmarkEngineConfig], list[BenchmarkCandidate]] @@ -45,22 +51,41 @@ tabular workloads we benchmark. """ -ATTENTION_BACKENDS: tuple[str, ...] = ( - "FLASHINFER", - "FLASH_ATTN", - "TRITON_ATTN", +_ATTENTION_BACKENDS_EXCLUDED_FROM_SWEEP: frozenset[str] = frozenset({"TORCH_SDPA", "FLEX_ATTENTION"}) +"""Common attention backends that are not direct benchmark sweep entries.""" + +ATTENTION_BACKENDS: tuple[str, ...] = tuple( + backend for backend in COMMON_ATTENTION_BACKENDS if backend not in _ATTENTION_BACKENDS_EXCLUDED_FROM_SWEEP ) -"""CUDA attention backends the sweep covers; excludes ROCm/XPU/MLA variants.""" +"""Common attention backends covered by the benchmark sweep.""" + +_STRUCTURED_BACKENDS_EXCLUDED_FROM_SWEEP: frozenset[StructuredGenerationBackend] = frozenset( + {"auto", "lm-format-enforcer"}, +) +"""Supported structured backends that are not direct benchmark sweep entries.""" + +STRUCTURED_BACKENDS: tuple[StructuredGenerationBackend, ...] = tuple( + backend + for backend in SUPPORTED_STRUCTURED_GENERATION_BACKENDS + if backend not in _STRUCTURED_BACKENDS_EXCLUDED_FROM_SWEEP +) +"""Structured-output backends covered by the benchmark sweep.""" + + +@dataclass(frozen=True) +class BatchingStep: + """One scheduler sizing point for :func:`batching_sweep`.""" + + max_num_seqs: int + max_num_batched_tokens: int -STRUCTURED_BACKENDS: tuple[str, ...] = ("xgrammar", "outlines", "guidance") -"""Structured-output backends the sweep covers.""" -BATCHING_STEPS: tuple[tuple[int, int], ...] = ( - (128, 4096), - (256, 8192), - (512, 16384), +BATCHING_STEPS: tuple[BatchingStep, ...] = ( + BatchingStep(max_num_seqs=128, max_num_batched_tokens=4096), + BatchingStep(max_num_seqs=256, max_num_batched_tokens=8192), + BatchingStep(max_num_seqs=512, max_num_batched_tokens=16384), ) -"""(max_num_seqs, max_num_batched_tokens) steps for the batching sweep.""" +"""Scheduler sizing steps for the batching sweep.""" MAX_MODEL_LEN_STEPS: tuple[int, ...] = (2048, 4096, 8192) """``max_model_len`` steps for the max-model-len sweep.""" @@ -141,11 +166,11 @@ def batching_sweep(base: BenchmarkEngineConfig) -> list[BenchmarkCandidate]: return [ _named_copy( base, - f"batch_seqs={seqs}_tokens={tokens}", - max_num_seqs=seqs, - max_num_batched_tokens=tokens, + f"batch_seqs={step.max_num_seqs}_tokens={step.max_num_batched_tokens}", + max_num_seqs=step.max_num_seqs, + max_num_batched_tokens=step.max_num_batched_tokens, ) - for seqs, tokens in BATCHING_STEPS + for step in BATCHING_STEPS ] diff --git a/tests/generation/test_vllm_benchmark.py b/tests/generation/test_vllm_benchmark.py index 7a2657fdc..e7d5ecc5f 100644 --- a/tests/generation/test_vllm_benchmark.py +++ b/tests/generation/test_vllm_benchmark.py @@ -21,6 +21,7 @@ import nemo_safe_synthesizer.generation.processors as processors_mod import nemo_safe_synthesizer.generation.vllm_benchmark as benchmark_mod +from nemo_safe_synthesizer.config.generate import COMMON_ATTENTION_BACKENDS, SUPPORTED_STRUCTURED_GENERATION_BACKENDS from nemo_safe_synthesizer.generation.vllm_benchmark import ( BenchmarkCandidate, BenchmarkCandidateDocument, @@ -40,13 +41,16 @@ run_benchmark, ) from nemo_safe_synthesizer.generation.vllm_benchmark_presets import ( + BATCHING_STEPS, DEFAULT_BENCHMARK_SEED, DEFAULT_BRACKETED_AB_N, + BatchingStep, attention_backend_sweep, baseline, bracketed_ab, bracketed_ab_spec_ngram, default_matrix, + structured_backend_sweep, ) from nemo_safe_synthesizer.generation.vllm_observability import GenerationObservability @@ -492,10 +496,25 @@ def test_default_matrix_dedupes(self, empty_base: BenchmarkEngineConfig) -> None keys = {(c.engine_config.model_dump_json(), json.dumps(c.sampling_overrides, sort_keys=True)) for c in cands} assert len(keys) == len(cands) - def test_attention_backend_sweep_covers_known_backends(self, empty_base: BenchmarkEngineConfig) -> None: + def test_attention_backend_sweep_uses_common_backends(self, empty_base: BenchmarkEngineConfig) -> None: cands = attention_backend_sweep(empty_base) backends = {c.engine_config.attention_backend for c in cands} - assert {"FLASHINFER", "FLASH_ATTN", "TRITON_ATTN"}.issubset(backends) + assert backends == set(COMMON_ATTENTION_BACKENDS) - {"TORCH_SDPA", "FLEX_ATTENTION"} + + def test_structured_backend_sweep_uses_supported_explicit_backends( + self, + empty_base: BenchmarkEngineConfig, + ) -> None: + cands = structured_backend_sweep(empty_base) + backends = {c.engine_config.structured_generation_backend for c in cands} + assert backends == set(SUPPORTED_STRUCTURED_GENERATION_BACKENDS) - {"auto", "lm-format-enforcer"} + + def test_batching_steps_are_named_scheduler_settings(self) -> None: + assert BATCHING_STEPS == ( + BatchingStep(max_num_seqs=128, max_num_batched_tokens=4096), + BatchingStep(max_num_seqs=256, max_num_batched_tokens=8192), + BatchingStep(max_num_seqs=512, max_num_batched_tokens=16384), + ) class TestBracketedAb: From 6b6d5286d28fb212c748a623bfd9cac1c1b1c2c8 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Tue, 16 Jun 2026 20:15:18 +0000 Subject: [PATCH 16/17] Move vLLM benchmark subprocess runner to tool Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark.py | 104 +------------ .../generation/vllm_benchmark_single_run.py | 50 ------- tests/generation/test_vllm_benchmark.py | 40 ----- tests/generation/test_vllm_benchmark_cli.py | 89 ++++++++++- tools/vllm_benchmark.py | 140 +++++++++++++++++- 5 files changed, 223 insertions(+), 200 deletions(-) delete mode 100644 src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index 80e44661e..9c2422a49 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -6,8 +6,7 @@ The harness replays a captured workload corpus (one ``GenerationTrace`` JSONL) under varying engine + sampling configurations and reports calibrated metrics per candidate. The models in this module stay -CPU-importable; the runner + subprocess wrapper live in sibling -modules and import vLLM lazily. +CPU-importable; the runner imports vLLM lazily. Architecture: @@ -26,17 +25,13 @@ - :class:`BenchmarkOutput` - JSON-serialised result of one matrix invocation, with skip records for candidates that failed. -The runner (next commit) is in ``vllm_benchmark.py`` alongside these -models; the subprocess wrapper + single-run entry point are split into -``vllm_benchmark_single_run.py``. +The repo-local tool owns subprocess isolation; this module owns the +reusable schemas and in-process runner. """ from __future__ import annotations import json -import subprocess -import sys -import tempfile import threading import time from collections.abc import Sequence @@ -69,9 +64,6 @@ # capture-time-only metadata. _NON_SAMPLING_FIELDS: tuple[str, ...] = ("structured_outputs",) -SUBPROCESS_STDERR_LIMIT: int = 500 -"""Maximum bytes of captured stderr to record on a subprocess failure.""" - PromptAssemblyMode = Literal["multi_record", "per_record"] """Prompt-assembly regime - controls how max_tokens partitions the budget.""" @@ -826,93 +818,3 @@ def run_benchmark( bracket_position=candidate.bracket_position, observability=observability, ) - - -# --------------------------------------------------------------------------- -# Subprocess isolation -# --------------------------------------------------------------------------- - - -class SubprocessRunResult(BaseModel): - """Outcome of one subprocess-isolated candidate run.""" - - model_config = ConfigDict(extra="forbid") - - metrics: CandidateMetrics | None = Field(default=None, description="Populated when the child exited successfully.") - error: str | None = Field(default=None, description="Captured stderr summary; populated on non-zero exit.") - error_class: str | None = Field(default=None, description="Best-effort exception class name parsed from stderr.") - - -def _truncate_stderr(stderr: str, limit: int = SUBPROCESS_STDERR_LIMIT) -> str: - """Trim ``stderr`` to ``limit`` bytes, keeping the tail.""" - stderr = stderr.strip() - if len(stderr) <= limit: - return stderr - head = "...[truncated]..." - return head + stderr[-(limit - len(head)) :] - - -def _parse_error_class(stderr: str) -> str: - """Best-effort parse of the exception class name from a Python traceback.""" - for line in reversed(stderr.strip().splitlines()): - stripped = line.strip() - if not stripped: - continue - head = stripped.split(":", 1)[0] - if head.isidentifier() or "." in head: - return head - return "Error" - return "Error" - - -def run_benchmark_in_subprocess( - candidate: BenchmarkCandidate, - corpus_path: str | Path, - simulate_training_overlap_seconds: float = 0.0, -) -> SubprocessRunResult: - """Run one candidate in a child process the OS reclaims on exit. - - Spawns ``python -m nemo_safe_synthesizer.generation.vllm_benchmark_single_run`` - with the candidate JSON-encoded as argv. Child writes - ``CandidateMetrics`` JSON to a temp file; parent reads it back. - - Subprocess isolation is what makes a multi-candidate matrix - reliable on this stack: vLLM holds significant CUDA + DRAM state - in module-level globals that the in-process Python runtime can't - clean up between candidates. Each candidate runs in a fresh - interpreter; the OS reclaims everything on child exit. - """ - candidate_json = candidate.model_dump_json() - with tempfile.TemporaryDirectory(prefix="nss-vllm-benchmark-") as result_dir: - result_path = Path(result_dir) / "result.json" - completed = subprocess.run( - [ - sys.executable, - "-m", - "nemo_safe_synthesizer.generation.vllm_benchmark_single_run", - "--candidate", - candidate_json, - "--corpus", - str(corpus_path), - "--result-out", - str(result_path), - "--simulate-training-overlap-seconds", - str(simulate_training_overlap_seconds), - ], - capture_output=True, - text=True, - check=False, - ) - if completed.returncode != 0: - stderr = _truncate_stderr(completed.stderr or completed.stdout) - return SubprocessRunResult( - error=stderr or f"subprocess exit {completed.returncode}", - error_class=_parse_error_class(completed.stderr or completed.stdout), - ) - if not result_path.exists() or result_path.stat().st_size == 0: - return SubprocessRunResult( - error="subprocess exited 0 but produced no result file", - error_class="RuntimeError", - ) - metrics = CandidateMetrics.model_validate_json(result_path.read_text(encoding="utf-8")) - return SubprocessRunResult(metrics=metrics) diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py deleted file mode 100644 index 1b8e546c8..000000000 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark_single_run.py +++ /dev/null @@ -1,50 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Subprocess entry point for one isolated benchmark candidate run. - -Invoked by :func:`vllm_benchmark.run_benchmark_in_subprocess` via -``python -m nemo_safe_synthesizer.generation.vllm_benchmark_single_run``. -Loads the candidate from the ``--candidate`` JSON argv, loads the -corpus from ``--corpus``, runs :func:`vllm_benchmark.run_benchmark`, -and writes the resulting ``CandidateMetrics`` JSON to ``--result-out``. - -Subprocess isolation is what makes a multi-candidate matrix reliable -on this stack — vLLM holds significant CUDA + DRAM state in -module-level globals; running each candidate in a fresh interpreter -lets the OS reclaim everything on exit. -""" - -from __future__ import annotations - -import argparse -from pathlib import Path - -from .vllm_benchmark import BenchmarkCandidate, BenchmarkCorpus, run_benchmark - - -def main() -> None: - parser = argparse.ArgumentParser(description="Run one benchmark candidate in isolation.") - parser.add_argument("--candidate", required=True, help="JSON-serialised BenchmarkCandidate.") - parser.add_argument("--corpus", required=True, help="Path to the corpus JSONL.") - parser.add_argument("--result-out", required=True, help="Path to write the CandidateMetrics JSON.") - parser.add_argument( - "--simulate-training-overlap-seconds", - type=float, - default=0.0, - help="Seconds to sleep after kicking off engine init (simulates concurrent training).", - ) - args = parser.parse_args() - - candidate = BenchmarkCandidate.model_validate_json(args.candidate) - corpus = BenchmarkCorpus.from_trace_jsonl(args.corpus) - metrics = run_benchmark( - candidate=candidate, - corpus=corpus, - simulate_training_overlap_seconds=args.simulate_training_overlap_seconds, - ) - Path(args.result_out).write_text(metrics.model_dump_json(indent=2), encoding="utf-8") - - -if __name__ == "__main__": - main() diff --git a/tests/generation/test_vllm_benchmark.py b/tests/generation/test_vllm_benchmark.py index e7d5ecc5f..c3bbdc02d 100644 --- a/tests/generation/test_vllm_benchmark.py +++ b/tests/generation/test_vllm_benchmark.py @@ -30,14 +30,11 @@ BenchmarkOutput, BenchmarkPrompt, CandidateMetrics, - SubprocessRunResult, TraceHeader, TracePromptRecord, _build_vllm_kwargs, _extract_ttft_ms, - _parse_error_class, _percentile, - _truncate_stderr, run_benchmark, ) from nemo_safe_synthesizer.generation.vllm_benchmark_presets import ( @@ -375,16 +372,6 @@ def test_extract_ttft_ms(self, metrics_obj: Any, expected: float | None) -> None output = type("Out", (), {"metrics": metrics_obj})() assert _extract_ttft_ms(output) == expected - def test_truncate_stderr_keeps_tail(self) -> None: - long = "x" * 600 - truncated = _truncate_stderr(long, limit=100) - assert len(truncated) <= 100 - assert truncated.endswith("xxxxx") # tail-preserving - - def test_parse_error_class_finds_exception_name(self) -> None: - assert _parse_error_class("Traceback (most recent call last):\n...\nValueError: bad") == "ValueError" - assert _parse_error_class("") == "Error" - # --------------------------------------------------------------------------- # Engine kwargs builder @@ -550,30 +537,3 @@ def test_all_candidate_runs_seed_pinned(self, empty_base: BenchmarkEngineConfig) candidate_runs = bracketed_ab_spec_ngram(empty_base) for c in candidate_runs: assert c.sampling_overrides.get("seed") == DEFAULT_BENCHMARK_SEED - - -# --------------------------------------------------------------------------- -# SubprocessRunResult -# --------------------------------------------------------------------------- - - -class TestSubprocessRunResult: - def test_success_shape(self) -> None: - m = CandidateMetrics( - name="t", - raw_tok_s=1.0, - acceptance_rate=0.9, - effective_tok_s=0.9, - ttft_p50_ms=0.0, - ttft_p99_ms=0.0, - prompts_attempted=10, - prompts_accepted=9, - total_output_tokens=100, - total_wall_seconds=1.0, - ) - r = SubprocessRunResult(metrics=m) - assert r.metrics is not None and r.error is None - - def test_failure_shape(self) -> None: - r = SubprocessRunResult(error="exit 1", error_class="RuntimeError") - assert r.metrics is None and r.error_class == "RuntimeError" diff --git a/tests/generation/test_vllm_benchmark_cli.py b/tests/generation/test_vllm_benchmark_cli.py index b6dc3ba14..9ad23baf0 100644 --- a/tests/generation/test_vllm_benchmark_cli.py +++ b/tests/generation/test_vllm_benchmark_cli.py @@ -8,13 +8,16 @@ import importlib.util import sys from pathlib import Path +from types import SimpleNamespace +from typing import Any +import pytest from click.testing import CliRunner -from nemo_safe_synthesizer.generation.vllm_benchmark import BenchmarkOutput, CandidateMetrics +from nemo_safe_synthesizer.generation.vllm_benchmark import BenchmarkCandidate, BenchmarkOutput, CandidateMetrics -def _load_cli(): +def _load_tool(): tool_path = Path(__file__).resolve().parents[2] / "tools" / "vllm_benchmark.py" spec = importlib.util.spec_from_file_location("vllm_benchmark_tool", tool_path) if spec is None or spec.loader is None: @@ -22,10 +25,11 @@ def _load_cli(): module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module spec.loader.exec_module(module) - return module.cli + return module -cli = _load_cli() +tool = _load_tool() +cli = tool.cli def _metric(name: str, condition: str, *, eff: float, wall: float, bracket: int) -> CandidateMetrics: @@ -71,3 +75,80 @@ def test_analyze_accepts_min_runs_option(tmp_path: Path) -> None: assert result.exit_code == 0, result.output assert "Cluster-conditioned analysis" in result.output + + +def test_hidden_run_candidate_command_is_available() -> None: + """The subprocess wrapper should target the repo-local tool command.""" + result = CliRunner().invoke(cli, ["_run-candidate", "--help"], color=False) + + assert result.exit_code == 0, result.output + assert "Run one benchmark candidate" in result.output + + +def test_truncate_stderr_keeps_tail() -> None: + long = "x" * 600 + truncated = tool._truncate_stderr(long, limit=100) + + assert len(truncated) <= 100 + assert truncated.endswith("xxxxx") + + +def test_parse_error_class_finds_exception_name() -> None: + assert tool._parse_error_class("Traceback (most recent call last):\n...\nValueError: bad") == "ValueError" + assert tool._parse_error_class("") == "Error" + + +def test_subprocess_result_success_shape() -> None: + metrics = CandidateMetrics( + name="t", + raw_tok_s=1.0, + acceptance_rate=0.9, + effective_tok_s=0.9, + ttft_p50_ms=0.0, + ttft_p99_ms=0.0, + prompts_attempted=10, + prompts_accepted=9, + total_output_tokens=100, + total_wall_seconds=1.0, + ) + + result = tool.SubprocessRunResult(metrics=metrics) + + assert result.metrics is not None and result.error is None + + +def test_subprocess_result_failure_shape() -> None: + result = tool.SubprocessRunResult(error="exit 1", error_class="RuntimeError") + + assert result.metrics is None and result.error_class == "RuntimeError" + + +def test_subprocess_runner_targets_tool_command(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + calls: list[list[str]] = [] + + def fake_run(args: list[str], **_kwargs: Any) -> Any: + calls.append(args) + result_path = Path(args[args.index("--result-out") + 1]) + metrics = CandidateMetrics( + name="baseline", + raw_tok_s=1.0, + acceptance_rate=1.0, + effective_tok_s=1.0, + ttft_p50_ms=0.0, + ttft_p99_ms=0.0, + prompts_attempted=1, + prompts_accepted=1, + total_output_tokens=1, + total_wall_seconds=1.0, + ) + result_path.write_text(metrics.model_dump_json(), encoding="utf-8") + return SimpleNamespace(returncode=0, stderr="", stdout="") + + monkeypatch.setattr(tool.subprocess, "run", fake_run) + + result = tool.run_benchmark_in_subprocess(BenchmarkCandidate(name="baseline"), tmp_path / "trace.jsonl") + + assert result.metrics is not None + assert calls[0][0] == sys.executable + assert calls[0][1].endswith("tools/vllm_benchmark.py") + assert calls[0][2] == "_run-candidate" diff --git a/tools/vllm_benchmark.py b/tools/vllm_benchmark.py index 9547f79c4..e44356363 100644 --- a/tools/vllm_benchmark.py +++ b/tools/vllm_benchmark.py @@ -27,11 +27,15 @@ from __future__ import annotations +import subprocess +import sys +import tempfile from datetime import datetime, timezone from pathlib import Path from typing import cast import click +from pydantic import BaseModel, ConfigDict, Field from rich.console import Console from rich.table import Table @@ -53,6 +57,63 @@ console = Console() +SUBPROCESS_STDERR_LIMIT: int = 500 +"""Maximum bytes of captured stderr to record on a subprocess failure.""" + + +class SubprocessRunResult(BaseModel): + """Outcome of one subprocess-isolated candidate run.""" + + model_config = ConfigDict(extra="forbid") + + metrics: CandidateMetrics | None = Field(default=None, description="Populated when the child exited successfully.") + error: str | None = Field(default=None, description="Captured stderr summary; populated on non-zero exit.") + error_class: str | None = Field(default=None, description="Best-effort exception class name parsed from stderr.") + + +def _truncate_stderr(stderr: str, limit: int = SUBPROCESS_STDERR_LIMIT) -> str: + """Trim ``stderr`` to ``limit`` bytes, keeping the tail.""" + stderr = stderr.strip() + if len(stderr) <= limit: + return stderr + head = "...[truncated]..." + return head + stderr[-(limit - len(head)) :] + + +def _parse_error_class(stderr: str) -> str: + """Best-effort parse of the exception class name from a Python traceback.""" + for line in reversed(stderr.strip().splitlines()): + stripped = line.strip() + if not stripped: + continue + head = stripped.split(":", 1)[0] + if head.isidentifier() or "." in head: + return head + return "Error" + return "Error" + + +def _run_candidate_command_args( + candidate: BenchmarkCandidate, + corpus_path: str | Path, + result_path: Path, + simulate_training_overlap_seconds: float, +) -> list[str]: + """Build argv for this tool's isolated candidate runner.""" + return [ + sys.executable, + str(Path(__file__).resolve()), + "_run-candidate", + "--candidate", + candidate.model_dump_json(), + "--corpus", + str(corpus_path), + "--result-out", + str(result_path), + "--simulate-training-overlap-seconds", + str(simulate_training_overlap_seconds), + ] + @click.group() def cli() -> None: @@ -90,6 +151,40 @@ def _resolve_candidates( raise click.UsageError("One of --candidates or --candidates-file is required.") +def run_benchmark_in_subprocess( + candidate: BenchmarkCandidate, + corpus_path: str | Path, + simulate_training_overlap_seconds: float = 0.0, +) -> SubprocessRunResult: + """Run one candidate in a child process the OS reclaims on exit.""" + with tempfile.TemporaryDirectory(prefix="nss-vllm-benchmark-") as result_dir: + result_path = Path(result_dir) / "result.json" + completed = subprocess.run( + _run_candidate_command_args( + candidate, + corpus_path, + result_path, + simulate_training_overlap_seconds, + ), + capture_output=True, + text=True, + check=False, + ) + if completed.returncode != 0: + stderr = _truncate_stderr(completed.stderr or completed.stdout) + return SubprocessRunResult( + error=stderr or f"subprocess exit {completed.returncode}", + error_class=_parse_error_class(completed.stderr or completed.stdout), + ) + if not result_path.exists() or result_path.stat().st_size == 0: + return SubprocessRunResult( + error="subprocess exited 0 but produced no result file", + error_class="RuntimeError", + ) + metrics = CandidateMetrics.model_validate_json(result_path.read_text(encoding="utf-8")) + return SubprocessRunResult(metrics=metrics) + + @cli.command("run") @click.argument("corpus_path", type=click.Path(exists=True, dir_okay=False, path_type=Path)) @click.option( @@ -128,11 +223,6 @@ def run_cmd( simulate_training_overlap_seconds: float, ) -> None: """Replay CORPUS_PATH against the chosen candidates and persist results.""" - # Lazy import so ``list`` and ``compare`` work without spinning up vLLM. - from nemo_safe_synthesizer.generation.vllm_benchmark import ( - run_benchmark_in_subprocess, - ) - corpus = BenchmarkCorpus.from_trace_jsonl(corpus_path) base = BenchmarkEngineConfig.model_validate(corpus.header.engine_parameters or {}) candidates = _resolve_candidates(base, preset_name, candidates_file) @@ -190,6 +280,46 @@ def run_cmd( console.print(f"[green]wrote[/green] {output_path} ({len(results)}/{len(candidates)} ok, {len(skipped)} skipped)") +@cli.command("_run-candidate", hidden=True) +@click.option("--candidate", required=True, help="JSON-serialised BenchmarkCandidate.") +@click.option( + "--corpus", + "corpus_path", + type=click.Path(path_type=Path, exists=True, dir_okay=False), + required=True, + help="Path to the corpus JSONL.", +) +@click.option( + "--result-out", + "result_out", + type=click.Path(path_type=Path, dir_okay=False), + required=True, + help="Path to write the CandidateMetrics JSON.", +) +@click.option( + "--simulate-training-overlap-seconds", + type=float, + default=0.0, + show_default=True, + help="Seconds to sleep after kicking off engine init (simulates concurrent training).", +) +def run_candidate_cmd( + candidate: str, + corpus_path: Path, + result_out: Path, + simulate_training_overlap_seconds: float, +) -> None: + """Run one benchmark candidate in an isolated child process.""" + from nemo_safe_synthesizer.generation.vllm_benchmark import run_benchmark + + metrics = run_benchmark( + candidate=BenchmarkCandidate.model_validate_json(candidate), + corpus=BenchmarkCorpus.from_trace_jsonl(corpus_path), + simulate_training_overlap_seconds=simulate_training_overlap_seconds, + ) + result_out.write_text(metrics.model_dump_json(indent=2), encoding="utf-8") + + @cli.command("compare") @click.argument("output_path", type=click.Path(exists=True, dir_okay=False, path_type=Path)) def compare_cmd(output_path: Path) -> None: From 9707c60533cbe883e8d342bc269dae85c7c5211c Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Tue, 16 Jun 2026 21:15:19 +0000 Subject: [PATCH 17/17] Fix vLLM benchmark replay fidelity Signed-off-by: Aaron Gonzales --- .../generation/vllm_benchmark.py | 54 ++++++--- tests/generation/test_vllm_benchmark.py | 113 +++++++++++++++++- tests/generation/test_vllm_benchmark_cli.py | 53 ++++++++ tools/vllm_benchmark.py | 4 +- 4 files changed, 206 insertions(+), 18 deletions(-) diff --git a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py index 9c2422a49..96a5783e7 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_benchmark.py +++ b/src/nemo_safe_synthesizer/generation/vllm_benchmark.py @@ -51,6 +51,7 @@ read_loadavg, read_vllm_runtime_metrics, ) +from ..llm.utils import ModelRef from ..observability import get_logger if TYPE_CHECKING: @@ -192,6 +193,19 @@ def from_trace_jsonl(cls, path: str | Path) -> BenchmarkCorpus: raise ValueError(f"{path}: missing header line") return cls(header=header, prompts=prompts) + def common_sampling_params(self) -> dict[str, Any]: + """Return the corpus-wide sampling params, rejecting heterogeneous traces.""" + if not self.prompts: + return {} + first = self.prompts[0].original_sampling_params + for prompt in self.prompts[1:]: + if prompt.original_sampling_params != first: + raise ValueError( + "benchmark corpus contains heterogeneous sampling_params; " + "capture a homogeneous trace or split the corpus before replay", + ) + return dict(first) + class BenchmarkEngineConfig(BaseModel): """Engine-construction kwargs the harness forwards to ``vllm.LLM(...)``. @@ -219,11 +233,12 @@ class BenchmarkEngineConfig(BaseModel): default=None, description="vLLM attention backend (``FLASHINFER``, ``FLASH_ATTN``, ``TRITON_ATTN``, etc.). ``None`` or ``'auto'`` leaves it unset.", ) - structured_generation_backend: str = Field( - default="xgrammar", + structured_generation_backend: str | None = Field( + default=None, description=( - "Structured-outputs backend used by benchmark sweeps. The preset " - "matrix covers ``'xgrammar'``, ``'outlines'``, and ``'guidance'``." + "Structured-outputs backend used by benchmark sweeps. ``None`` preserves " + "the trace header's captured backend. The preset matrix covers " + "``'xgrammar'``, ``'outlines'``, and ``'guidance'``." ), ) max_model_len: int | None = Field( @@ -284,6 +299,14 @@ class BenchmarkEngineConfig(BaseModel): ), ) + def with_trace_defaults(self, header: TraceHeader) -> Self: + """Apply trace-level sizing defaults that are not direct vLLM kwargs.""" + if self.max_model_len is not None: + return self + if header.max_tokens_per_example is None or header.max_tokens_per_example <= 0: + return self + return self.model_copy(update={"max_model_len": header.max_tokens_per_example}) + class BenchmarkCandidate(BaseModel): """One configuration to benchmark - engine kwargs + sampling overrides + identity. @@ -569,24 +592,25 @@ def _build_vllm_kwargs(header: TraceHeader, engine_config: BenchmarkEngineConfig Drops ``None``-valued candidate fields explicitly so vLLM treats them as "not configured" rather than "override to None". """ - overlay = engine_config.model_dump(exclude_none=True) + overlay = engine_config.with_trace_defaults(header).model_dump(exclude_none=True) base = dict(header.engine_parameters) base.update(overlay) # Required-positional kwargs that aren't in BenchmarkEngineConfig: - base["model"] = header.pretrained_model + model_ref = ModelRef.parse(header.pretrained_model) + base["model"] = model_ref.target() + base["trust_remote_code"] = model_ref.trust_remote_code base.setdefault("enable_lora", header.lora_path is not None) # ``attention_backend`` -> ``attention_config`` translation. vLLM's # public API takes a config dict rather than a bare string. attention_backend = base.pop("attention_backend", None) if attention_backend not in (None, "auto"): base["attention_config"] = {"backend": attention_backend} - # ``structured_generation_backend`` -> ``structured_outputs_config``. - from vllm.config import ( - StructuredOutputsConfig, # noqa: PLC0415 - lazy, vLLM is heavy - ) - sg_backend = base.pop("structured_generation_backend", None) if sg_backend is not None: + from vllm.config import ( + StructuredOutputsConfig, # noqa: PLC0415 - lazy, vLLM is heavy + ) + base["structured_outputs_config"] = StructuredOutputsConfig(backend=sg_backend) return base @@ -669,6 +693,9 @@ def run_benchmark( when the engine's effective runtime config disagrees with the candidate's intended ``engine_config`` on any checked field. """ + base_sampling = corpus.common_sampling_params() + dispatch = candidate.dispatch_plan([p.prompt for p in corpus.prompts], base_sampling) + # Lazy imports - keep this module CPU-importable. from vllm.lora.request import LoRARequest # noqa: PLC0415 from vllm.sampling_params import SamplingParams # noqa: PLC0415 @@ -721,11 +748,6 @@ def run_benchmark( LoRARequest("lora", 1, str(corpus.header.lora_path)) if corpus.header.lora_path is not None else None ) - # Build SamplingParams from corpus default + candidate overrides. - base_sampling: dict[str, Any] = {} - if corpus.prompts: - base_sampling = corpus.prompts[0].original_sampling_params - dispatch = candidate.dispatch_plan([p.prompt for p in corpus.prompts], base_sampling) sampling_params = SamplingParams(**dispatch.sampling_kwargs) # Dispatch. diff --git a/tests/generation/test_vllm_benchmark.py b/tests/generation/test_vllm_benchmark.py index c3bbdc02d..37d5937ce 100644 --- a/tests/generation/test_vllm_benchmark.py +++ b/tests/generation/test_vllm_benchmark.py @@ -163,6 +163,17 @@ def _install_fake_vllm_modules(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setitem(sys.modules, "vllm.sampling_params", sampling_mod) +def _install_fake_vllm_config(monkeypatch: pytest.MonkeyPatch) -> None: + config_mod = types.ModuleType("vllm.config") + + class FakeStructuredOutputsConfig: + def __init__(self, *, backend: str) -> None: + self.backend = backend + + setattr(config_mod, "StructuredOutputsConfig", FakeStructuredOutputsConfig) + monkeypatch.setitem(sys.modules, "vllm.config", config_mod) + + def _install_runner_fakes( monkeypatch: pytest.MonkeyPatch, llm: _FakeLLM, @@ -204,6 +215,20 @@ def _benchmark_corpus(prompts: list[str]) -> BenchmarkCorpus: ) +def _benchmark_corpus_with_sampling(sampling_params: list[dict[str, Any]]) -> BenchmarkCorpus: + return BenchmarkCorpus( + header=TraceHeader(run_id="r", pretrained_model="m", dataset_schema={"col": "string"}), + prompts=[ + BenchmarkPrompt( + row_index=i, + prompt=f"p{i}", + original_sampling_params=params, + ) + for i, params in enumerate(sampling_params) + ], + ) + + # --------------------------------------------------------------------------- # Data model contracts # --------------------------------------------------------------------------- @@ -321,6 +346,17 @@ def test_rejects_missing_header(self, tmp_path: Any) -> None: with pytest.raises(ValueError, match="record on line 1 before any header"): BenchmarkCorpus.from_trace_jsonl(path) + def test_common_sampling_params_rejects_heterogeneous_records(self) -> None: + corpus = _benchmark_corpus_with_sampling( + [ + {"temperature": 0.0, "max_tokens": 8}, + {"temperature": 0.7, "max_tokens": 8}, + ], + ) + + with pytest.raises(ValueError, match="heterogeneous sampling_params"): + corpus.common_sampling_params() + # --------------------------------------------------------------------------- # Helper contracts @@ -379,14 +415,72 @@ def test_extract_ttft_ms(self, metrics_obj: Any, expected: float | None) -> None class TestBuildVllmKwargs: - def test_overlays_engine_config_on_header(self, header: TraceHeader) -> None: + def test_resolves_model_ref_target_and_trust_remote_code( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + class FakeModelRef: + trust_remote_code = True + + def target(self) -> str: + return "/cache/nvidia-model" + + parse_calls: list[str] = [] + + def fake_parse(value: str) -> FakeModelRef: + parse_calls.append(value) + return FakeModelRef() + + monkeypatch.setattr(benchmark_mod.ModelRef, "parse", staticmethod(fake_parse)) + header = TraceHeader(run_id="r", pretrained_model="nvidia/Nemotron-Mini-4B-Instruct", dataset_schema={}) + + kwargs = _build_vllm_kwargs(header, BenchmarkEngineConfig()) + + assert parse_calls == ["nvidia/Nemotron-Mini-4B-Instruct"] + assert kwargs["model"] == "/cache/nvidia-model" + assert kwargs["trust_remote_code"] is True + + def test_overlays_engine_config_on_header(self, monkeypatch: pytest.MonkeyPatch, header: TraceHeader) -> None: """Candidate engine_config overlays the header's engine_parameters.""" + monkeypatch.setattr( + benchmark_mod.ModelRef, + "parse", + staticmethod( + lambda value: SimpleNamespace( + target=lambda: value, + trust_remote_code=False, + ), + ), + ) cfg = BenchmarkEngineConfig(attention_backend="FLASHINFER", max_model_len=4096) kwargs = _build_vllm_kwargs(header, cfg) assert kwargs["model"] == "mistralai/Mistral-7B-Instruct-v0.3" assert kwargs["max_lora_rank"] == 32 # from header assert kwargs["max_model_len"] == 4096 # from cfg + def test_empty_engine_config_preserves_header_structured_backend( + self, + monkeypatch: pytest.MonkeyPatch, + header: TraceHeader, + ) -> None: + _install_fake_vllm_config(monkeypatch) + + kwargs = _build_vllm_kwargs(header, BenchmarkEngineConfig()) + + assert kwargs["structured_outputs_config"].backend == "outlines" + + def test_engine_config_applies_trace_max_token_hint_when_unset(self) -> None: + header = TraceHeader( + run_id="r", + pretrained_model="m", + dataset_schema={}, + max_tokens_per_example=8192, + ) + + cfg = BenchmarkEngineConfig().with_trace_defaults(header) + + assert cfg.max_model_len == 8192 + def test_translates_attention_backend_to_attention_config( self, header: TraceHeader, empty_base: BenchmarkEngineConfig ) -> None: @@ -466,6 +560,23 @@ def test_overlap_savings_are_capped_by_actual_engine_init_time(self, monkeypatch assert metrics.startup_seconds == 0.0 assert metrics.startup_overlap_savings_seconds == 1.25 + def test_rejects_heterogeneous_sampling_params(self, monkeypatch: pytest.MonkeyPatch) -> None: + llm = _FakeLLM() + _install_runner_fakes(monkeypatch, llm) + + with pytest.raises(ValueError, match="heterogeneous sampling_params"): + run_benchmark( + BenchmarkCandidate(name="baseline"), + _benchmark_corpus_with_sampling( + [ + {"temperature": 0.0, "max_tokens": 8}, + {"temperature": 0.7, "max_tokens": 8}, + ], + ), + ) + + assert llm.calls == [] + # --------------------------------------------------------------------------- # Preset contracts diff --git a/tests/generation/test_vllm_benchmark_cli.py b/tests/generation/test_vllm_benchmark_cli.py index 9ad23baf0..b548bebe1 100644 --- a/tests/generation/test_vllm_benchmark_cli.py +++ b/tests/generation/test_vllm_benchmark_cli.py @@ -77,6 +77,59 @@ def test_analyze_accepts_min_runs_option(tmp_path: Path) -> None: assert "Cluster-conditioned analysis" in result.output +def test_run_presets_receive_trace_max_token_hint(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + trace_path = tmp_path / "trace.jsonl" + trace_path.write_text( + "\n".join( + [ + ( + '{"kind": "header", "run_id": "r", "pretrained_model": "m", ' + '"dataset_schema": {}, "max_tokens_per_example": 8192}' + ), + '{"kind": "record", "row_index": 0, "prompt": "p", "sampling_params": {"temperature": 0.0}}', + ], + ), + encoding="utf-8", + ) + output_path = tmp_path / "benchmark.json" + seen_candidates: list[BenchmarkCandidate] = [] + + def fake_run_benchmark_in_subprocess( + candidate: BenchmarkCandidate, + _corpus_path: str | Path, + **_kwargs: Any, + ) -> Any: + seen_candidates.append(candidate) + return tool.SubprocessRunResult( + metrics=CandidateMetrics( + name=candidate.name, + raw_tok_s=1.0, + acceptance_rate=1.0, + effective_tok_s=1.0, + ttft_p50_ms=0.0, + ttft_p99_ms=0.0, + prompts_attempted=1, + prompts_accepted=1, + total_output_tokens=1, + total_wall_seconds=1.0, + ), + ) + + monkeypatch.setattr(tool, "resolve_sweep_id", lambda: None) + monkeypatch.setattr(tool, "init_candidate_run", lambda **_kwargs: None) + monkeypatch.setattr(tool, "log_and_finish", lambda *_args, **_kwargs: None) + monkeypatch.setattr(tool, "run_benchmark_in_subprocess", fake_run_benchmark_in_subprocess) + + result = CliRunner().invoke( + cli, + ["run", str(trace_path), "--output", str(output_path), "--candidates", "baseline"], + color=False, + ) + + assert result.exit_code == 0, result.output + assert seen_candidates[0].engine_config.max_model_len == 8192 + + def test_hidden_run_candidate_command_is_available() -> None: """The subprocess wrapper should target the repo-local tool command.""" result = CliRunner().invoke(cli, ["_run-candidate", "--help"], color=False) diff --git a/tools/vllm_benchmark.py b/tools/vllm_benchmark.py index e44356363..c6cc0cb34 100644 --- a/tools/vllm_benchmark.py +++ b/tools/vllm_benchmark.py @@ -224,7 +224,9 @@ def run_cmd( ) -> None: """Replay CORPUS_PATH against the chosen candidates and persist results.""" corpus = BenchmarkCorpus.from_trace_jsonl(corpus_path) - base = BenchmarkEngineConfig.model_validate(corpus.header.engine_parameters or {}) + base = BenchmarkEngineConfig.model_validate(corpus.header.engine_parameters or {}).with_trace_defaults( + corpus.header + ) candidates = _resolve_candidates(base, preset_name, candidates_file) if not candidates: raise click.UsageError("Resolved candidate list is empty.")