diff --git a/AGENTS.md b/AGENTS.md index eb0ee43a..3777e83c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -177,7 +177,10 @@ src/inference_endpoint/ │ ├── benchmark/ │ │ ├── __init__.py │ │ ├── cli.py # benchmark_app: offline, online, from-config subcommands -│ │ └── execute.py # Phased execution: setup/run_threaded/finalize + BenchmarkContext; run_benchmark runs the main benchmark (cli._run dispatches run_audit when audit: is set) +│ │ ├── execute.py # Phased orchestration: setup_benchmark/run_benchmark_async/finalize_benchmark + BenchmarkContext; run_benchmark runs the main benchmark (cli._run dispatches run_audit when audit: is set) +│ │ ├── profiling.py # Profiler-trigger protocol (vLLM /start_profile,/stop_profile) + ProfileController (URL derivation + start/stop/payload lifecycle) +│ │ ├── accuracy.py # AccuracyConfiguration + per-dataset scoring (_score_accuracy, OSL/response-count rollups, write_accuracy_results) +│ │ └── pipeline.py # MetricsPipeline: ZMQ + metrics-aggregator/event-logger subprocess lifecycle (start/drain_and_build_report/abort/close) + snapshot→Report │ ├── audit.py # run_audit() — compliance audit orchestrator (phases → verify → result) │ ├── probe.py # ProbeConfig + execute_probe() │ ├── info.py # execute_info() diff --git a/src/inference_endpoint/commands/benchmark/accuracy.py b/src/inference_endpoint/commands/benchmark/accuracy.py new file mode 100644 index 00000000..abd0e75a --- /dev/null +++ b/src/inference_endpoint/commands/benchmark/accuracy.py @@ -0,0 +1,405 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Accuracy scoring for benchmark finalization. + +Off the hot path: scores each accuracy dataset (and the optional inline +perf-scored entry) into one ``accuracy_scores`` list entry, computes per-phase +response accounting and output-token-length rollups from ``events.jsonl``, and +writes the ``accuracy/accuracy_results.json`` artifact. +""" + +from __future__ import annotations + +import json +import logging +import numbers +import time +from collections.abc import Iterable +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import msgspec.json + +from inference_endpoint.async_utils.services.metrics_aggregator.token_metrics import ( + encode_lengths, + load_reference_backend, +) +from inference_endpoint.config.schema import DatasetType +from inference_endpoint.dataset_manager.dataset import Dataset +from inference_endpoint.evaluation import Extractor +from inference_endpoint.evaluation.accuracy_results import average_accuracy +from inference_endpoint.evaluation.scoring import Scorer +from inference_endpoint.exceptions import ExecutionError, InputValidationError +from inference_endpoint.load_generator.session import SessionResult +from inference_endpoint.metrics.report import series_metric_dict +from inference_endpoint.utils.atomic_write import atomic_write_bytes + +if TYPE_CHECKING: + from inference_endpoint.commands.benchmark.execute import BenchmarkContext + +logger = logging.getLogger(__name__) + + +@dataclass +class AccuracyConfiguration: + scorer: type[Scorer] + extractor: type[Extractor] | None + dataset_name: str + dataset: Dataset + report_dir: Path + ground_truth_column: str | None + num_repeats: int + extras: dict[str, Any] = field(default_factory=dict) + # Discriminates the inline perf-scored entry (PERFORMANCE) from real accuracy + # datasets (ACCURACY). Branch on this, not on dataset_name == "performance": + # a dataset legitimately named "performance" must not be misclassified. + dataset_type: DatasetType = DatasetType.ACCURACY + + +def _phase_osl_stats( + sample_uuids: Iterable[str], + uuid_to_text: dict[str, str], + backend: Any, + batch_size: int = 256, +) -> dict[str, Any] | None: + """Output-token-length rollup over one accuracy phase's completions. + + Counts tokens on each sample's response text via the shared reference + tokenizer backend — the server's ``completion_tokens`` is not persisted, only + the text is (in ``events.jsonl``) — then shapes the lengths via + ``series_metric_dict`` so the block matches the perf report's + ``output_sequence_lengths`` exactly. Returns ``None`` when the phase has no + completed outputs. + + ``batch_size`` bounds each ``encode_batch`` pass: accuracy outputs can be tens + of thousands of tokens each (e.g. gpt-oss lcb at 32768), so counting the whole + population in one call would hold every Encoding in memory at once. + """ + # Skip empty/failed completions (a failed request still logs a COMPLETE + # event with output == ""). The perf-side OslTrigger does the same + # (metrics_table.OslTrigger._extract_text returns None for empty text), so + # accuracy OSL matches its population and a failure isn't counted as a + # 0-token sample that would drag min/avg down. + texts = [ + uuid_to_text[u] for u in sample_uuids if u in uuid_to_text and uuid_to_text[u] + ] + if not texts: + return None + lengths: list[int] = [] + for i in range(0, len(texts), batch_size): + lengths.extend(encode_lengths(backend, texts[i : i + batch_size])) + return series_metric_dict(lengths) or None + + +def _phase_response_counts( + sample_uuids: Iterable[str], + uuid_to_text: dict[str, str], +) -> dict[str, int]: + """Per-phase response accounting over one accuracy phase's issued samples. + + Complements :func:`_phase_osl_stats`, which reports token lengths only over + non-empty completions — on its own that can hide a run where the server + returned blanks or dropped requests. Classifies each issued ``sample_uuid`` + as ``scored`` (COMPLETE, non-empty output — exactly the OSL population), + ``empty`` (COMPLETE with blank output: a failed request the load generator + logged as ERROR then an empty COMPLETE), or ``missing`` (no COMPLETE event). + ``issued == scored + empty + missing`` always holds. + + Emptiness uses the same truthiness test as ``_phase_osl_stats`` so ``scored`` + is byte-for-byte the OSL population — the two blocks cannot disagree. + """ + issued = scored = empty = missing = 0 + for u in sample_uuids: + issued += 1 + if u not in uuid_to_text: + missing += 1 + elif uuid_to_text[u]: + scored += 1 + else: + empty += 1 + return {"issued": issued, "scored": scored, "empty": empty, "missing": missing} + + +def _accuracy_uuid_bound( + report_dir: Path | None, eval_configs: list[AccuracyConfiguration] +) -> set[str]: + """Union of the accuracy datasets' issued uuids from ``sample_idx_map.json``. + + Bounds the finalize-side raw-output read to the accuracy population so it + never holds the whole run's (incl. perf) response-text corpus. Returns an + empty set (⇒ caller reads unbounded) when there is no report dir; a missing, + corrupt, or wrong-shape map is warned and also falls back to unbounded. + """ + if report_dir is None: + return set() + try: + idx_map = msgspec.json.decode((report_dir / "sample_idx_map.json").read_bytes()) + except (OSError, msgspec.DecodeError) as e: + logger.warning( + "Accuracy OSL uuid bound unavailable (%s); reading outputs unbounded", e + ) + return set() + # A syntactically-valid map of the wrong shape must not crash finalize: this + # runs outside the per-dataset try, so a raised AttributeError/TypeError would + # fail scoring (OSL must never do that). Skip anything not dict-shaped. + if not isinstance(idx_map, dict): + logger.warning( + "Accuracy OSL uuid bound: sample_idx_map.json is not an object; " + "reading outputs unbounded" + ) + return set() + bound: set[str] = set() + for ec in eval_configs: + if ec.dataset_type == DatasetType.ACCURACY: + per_dataset = idx_map.get(ec.dataset_name) + if isinstance(per_dataset, dict): + bound |= set(per_dataset) + return bound + + +def _load_osl_backend(has_accuracy: bool, tokenizer_name: str | None) -> Any | None: + """Load the reference tokenizer backend for accuracy OSL, or None to disable. + + Loaded only when a real accuracy dataset exists; a load failure or a tokenizer + with no fast (Rust) backend disables OSL rather than failing scoring. + """ + if not (has_accuracy and tokenizer_name is not None): + return None + try: + osl_backend = load_reference_backend(tokenizer_name) + except Exception as e: # noqa: BLE001 - OSL is optional; never fail scoring + logger.warning( + "Accuracy OSL disabled: could not load tokenizer %r: %s", + tokenizer_name, + e, + ) + return None + # A tokenizer with no fast (Rust) backend disables OSL rather than falling + # back to a slow Python-tokenizer count: the perf side + # (token_metrics._setup_shards) requires a fast backend too and raises without + # one, so OSL stays fast-only and consistent on both sides. Warn so the skip is + # visible instead of silently dropping the block. + if osl_backend is None: + logger.warning( + "Accuracy OSL disabled: tokenizer %r has no fast (Rust) backend " + "(token counting requires one, as on the perf side)", + tokenizer_name, + ) + return osl_backend + + +def _score_accuracy( + ctx: BenchmarkContext, result: SessionResult +) -> list[dict[str, Any]]: + """Score each accuracy dataset into its own list entry. + + One entry per eval_config, in order; no cross-dataset consolidation. Each + entry carries the scalar ``score`` plus sample accounting + (``unit_samples`` × ``num_repeats`` = ``total_samples``); a scorer that + returns a ``score_breakdown()`` (DeepSeek-R1, BFCL) also attaches + ``breakdown``. The ``"performance"`` inline entry totals the perf phases' + issued counts instead of unit × repeats (repeats is forced to 1 there). + """ + accuracy_scores: list[dict[str, Any]] = [] + + # Per-phase wall-clock (seconds) keyed by phase name. The accuracy phase name + # is the dataset_name; the inline-scored perf entry keys on "performance". + phase_durations: dict[str, float] = {} + for pr in result.phase_results: + phase_durations[pr.name] = phase_durations.get(pr.name, 0.0) + max( + 0.0, (pr.end_time_ns - pr.start_time_ns) / 1e9 + ) + + # Accuracy-phase output-token lengths (finalize-side, off the hot path): the + # aggregator only tokenizes perf-window samples, so count the accuracy + # responses (already in events.jsonl) here, using the same reference tokenizer + # as the perf side. (Counts still differ from perf for tool-call responses — + # client-side OSL is approximate for structured output.) + has_accuracy = any( + ec.dataset_type == DatasetType.ACCURACY for ec in ctx.eval_configs + ) + osl_backend = _load_osl_backend(has_accuracy, ctx.tokenizer_name) + # Bound the raw-output read to the accuracy population so finalize never holds + # the whole run's (incl. perf) response-text corpus. + accuracy_uuids = ( + _accuracy_uuid_bound(ctx.report_dir, ctx.eval_configs) + if has_accuracy + else set() + ) + uuid_to_text: dict[str, str] | None = None + + for eval_cfg in ctx.eval_configs: + try: + scorer_instance = eval_cfg.scorer( + eval_cfg.dataset_name, + eval_cfg.dataset, + eval_cfg.report_dir, + extractor=eval_cfg.extractor, + ground_truth_column=eval_cfg.ground_truth_column, + **eval_cfg.extras, + ) + except TypeError as e: + raise InputValidationError( + f"Dataset '{eval_cfg.dataset_name}': invalid accuracy_config.extras " + f"for scorer '{eval_cfg.scorer.__name__}': {e}" + ) from e + score, n_repeats = scorer_instance.score() + # Coerce a numpy scalar score (np.float32/64, numpy ints — e.g. np.mean + # from the base Scorer) to a native Python float so the entry stays + # serializable by both msgspec (result_summary.json) and json + # (accuracy_results.json). numbers.Real catches every numpy scalar (not + # just np.float64, which isinstance(..., float) alone would miss) while + # leaving None / dict (RougeScorer) untouched; bool is excluded. + if isinstance(score, numbers.Real) and not isinstance(score, bool): + score = float(score) + unit_samples = eval_cfg.dataset.num_samples() + num_repeats = eval_cfg.num_repeats + if eval_cfg.dataset_type == DatasetType.PERFORMANCE: + # A performance dataset always scores its already-issued outputs once + # (enforced by the num_repeats == 1 guard in _load_datasets), so make + # that locally provable rather than relying on eval_cfg carrying 1. + num_repeats = 1 + total_samples = sum(phase.issued_count for phase in result.perf_results) + else: + total_samples = unit_samples * num_repeats + entry: dict[str, Any] = { + "dataset_name": eval_cfg.dataset_name, + "extractor": ( + eval_cfg.extractor.__name__ if eval_cfg.extractor is not None else None + ), + "ground_truth_column": eval_cfg.ground_truth_column, + "score": score, + "unit_samples": unit_samples, + "num_repeats": num_repeats, + "total_samples": total_samples, + # Wall-clock of this dataset's issue phase (seconds); 0.0 if the + # phase left no timing (e.g. a scored-but-not-issued dataset). + "duration_s": round(phase_durations.get(eval_cfg.dataset_name, 0.0), 3), + # False when the scorer produced only a partial headline (e.g. + # LegacyMLPerfDeepSeekR1Scorer when lcb-service was unreachable), so a + # partial number is never mistaken for a complete one. + "complete": scorer_instance.complete, + # Persist the same DatasetType discriminator carried on the eval config + # so consumers filter the inline perf-scored entry by type, not by + # matching dataset_name == "performance". + "dataset_type": eval_cfg.dataset_type.value, + } + breakdown = scorer_instance.score_breakdown() + if breakdown is not None: + entry["breakdown"] = breakdown + + # Response accounting + avg/min/max output-token length. Skipped for the + # perf entry (its OSL / failure counts live in result_summary.json). The + # counts are computed independent of the tokenizer and of OSL returning a + # block — an all-failed phase must still publish scored=0 rather than + # silently omitting everything. OSL stays tokenizer-gated. A read/tokenize + # failure only drops these blocks — it never fails scoring. + if eval_cfg.dataset_type == DatasetType.ACCURACY: + try: + if uuid_to_text is None: + # Built once from the first scorer and reused for every + # dataset. get_raw_outputs() returns the model's actual + # completion text (not the scorer's scoring-normalized form) + # for *all* phases' COMPLETE events, bounded to the accuracy + # population; intersecting it with each dataset's + # sample_index_map yields correct per-dataset counts. + out_df = scorer_instance.get_raw_outputs(accuracy_uuids or None) + uuid_to_text = dict( + zip(out_df["sample_uuid"], out_df["output"], strict=False) + ) + # Drop the DataFrame so finalize doesn't hold both it and the + # dict (each carrying the response-text corpus). + del out_df + entry["response_counts"] = _phase_response_counts( + scorer_instance.sample_index_map, uuid_to_text + ) + if osl_backend is not None: + t0 = time.perf_counter() + osl = _phase_osl_stats( + scorer_instance.sample_index_map, uuid_to_text, osl_backend + ) + if osl is not None: + # Same shape/key as the perf report output_sequence_lengths. + entry["output_sequence_lengths"] = osl + # Wall-clock of just this phase's tokenization (seconds); + # summed across datasets for the accuracy report's total. + entry["osl_tokenize_s"] = round(time.perf_counter() - t0, 3) + except Exception as e: # noqa: BLE001 - optional blocks; never fail scoring + logger.warning( + "Accuracy response counts/OSL skipped for %s: %s", + eval_cfg.dataset_name, + e, + ) + + accuracy_scores.append(entry) + logger.info( + f"Score for {eval_cfg.dataset_name}: {score} " + f"({n_repeats} repeats, complete={scorer_instance.complete})" + ) + + return accuracy_scores + + +def write_accuracy_results( + report_dir: Path, accuracy_scores: list[dict[str, Any]] +) -> None: + """Emit ``accuracy/accuracy_results.json`` from the per-dataset score entries. + + Perf rollups (qps/tps/latency percentiles) live in + ``performance/result_summary.json`` and response/error text lives in + ``events.jsonl``, so neither is duplicated here. A no-op when ``accuracy_scores`` + is empty — a perf-only run leaves no ``accuracy/`` folder. + """ + if not accuracy_scores: + return + # Plain cross-component mean of the per-dataset scores (3 datasets for + # gpt-oss, 1 for DeepSeek-R1); None when nothing numeric was scored. + avg_accuracy = average_accuracy(accuracy_scores) + # Total finalize-time spent tokenizing accuracy outputs for OSL (seconds), + # summed across datasets. Emitted whenever OSL was computed for at least + # one dataset — gating on the key's presence, not the rounded wall-clock, + # so a sub-millisecond total (tiny outputs) still records 0.0 rather than + # silently dropping the field. + osl_computed = any("osl_tokenize_s" in e for e in accuracy_scores) + osl_tokenization_s = round( + sum(e.get("osl_tokenize_s", 0.0) for e in accuracy_scores), 3 + ) + accuracy_dir = report_dir / "accuracy" + accuracy_results_path = accuracy_dir / "accuracy_results.json" + accuracy_payload: dict[str, Any] = {} + if avg_accuracy is not None: + accuracy_payload["average_accuracy"] = avg_accuracy + if osl_computed: + accuracy_payload["osl_tokenization_s"] = osl_tokenization_s + accuracy_payload["accuracy_scores"] = accuracy_scores + # Atomic write so a crash mid-write can't leave truncated JSON the + # compliance checker would read as corrupt. Not swallowed: if scoring + # produced entries but they can't be persisted — the dir can't be made + # (OSError), the payload won't serialize (TypeError/ValueError, e.g. a + # numpy scalar left in a breakdown block), or the write fails — fail the + # run loudly rather than exit 0 with no accuracy artifact. + try: + accuracy_dir.mkdir(parents=True, exist_ok=True) + atomic_write_bytes( + accuracy_results_path, + json.dumps(accuracy_payload, indent=2).encode(), + ) + except (OSError, TypeError, ValueError) as e: + raise ExecutionError( + f"Failed to write accuracy results to {accuracy_results_path}: {e}" + ) from e + logger.info(f"Saved: {accuracy_results_path}") diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index 6b2e6d89..ecab04b6 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -19,6 +19,10 @@ 1. setup_benchmark() — load tokenizer, dataset, config (no IO) 2. run_benchmark_async() — HTTP client + async BenchmarkSession 3. finalize_benchmark() — accuracy scoring, results JSON + +Cohesive sub-concerns live in sibling modules: profiler triggers (``profiling``), +accuracy scoring (``accuracy``), and the ZMQ/metrics/event-logger service lifecycle +(``pipeline``). """ from __future__ import annotations @@ -26,47 +30,36 @@ import asyncio import json import logging -import numbers import random import shutil import signal import tempfile -import time import uuid -from collections.abc import Callable, Iterable +from collections.abc import Callable from dataclasses import dataclass, field from dataclasses import replace as dataclass_replace from datetime import datetime from pathlib import Path -from typing import Any, TextIO -from urllib import error as urllib_error -from urllib import request as urllib_request +from typing import Any from urllib.parse import urljoin -import msgspec import msgspec.json import msgspec.structs from huggingface_hub import model_info from tqdm import tqdm from transformers.utils import logging as transformers_logging -from inference_endpoint.async_utils.event_publisher import EventPublisherService from inference_endpoint.async_utils.loop_manager import LoopManager -from inference_endpoint.async_utils.services.launcher import ( - ServiceConfig, - ServiceLauncher, -) -from inference_endpoint.async_utils.services.metrics_aggregator.snapshot import ( - snapshot_to_dict, -) -from inference_endpoint.async_utils.services.metrics_aggregator.subscriber import ( - MetricsSnapshotSubscriber, +from inference_endpoint.commands.benchmark.accuracy import ( + AccuracyConfiguration, + _score_accuracy, + write_accuracy_results, ) -from inference_endpoint.async_utils.services.metrics_aggregator.token_metrics import ( - encode_lengths, - load_reference_backend, +from inference_endpoint.commands.benchmark.pipeline import MetricsPipeline +from inference_endpoint.commands.benchmark.profiling import ( + ProfileController, + _write_profiling_section, ) -from inference_endpoint.async_utils.transport.zmq.context import ManagedZMQContext from inference_endpoint.compliance import AuditRunSpec from inference_endpoint.config.runtime_settings import RuntimeSettings from inference_endpoint.config.schema import ( @@ -75,7 +68,6 @@ DatasetType, LoadPattern, LoadPatternType, - ProfilerEngine, StreamingMode, TestMode, TestType, @@ -90,7 +82,6 @@ from inference_endpoint.endpoint_client.http_client import HTTPEndpointClient from inference_endpoint.endpoint_client.http_sample_issuer import HttpClientSampleIssuer from inference_endpoint.evaluation import Extractor -from inference_endpoint.evaluation.accuracy_results import average_accuracy from inference_endpoint.evaluation.scoring import Scorer from inference_endpoint.exceptions import ( ExecutionError, @@ -107,8 +98,7 @@ PhaseType, SessionResult, ) -from inference_endpoint.metrics.report import Report, series_metric_dict -from inference_endpoint.utils.atomic_write import atomic_write_bytes +from inference_endpoint.metrics.report import Report transformers_logging.set_verbosity_error() @@ -170,22 +160,6 @@ class BenchmarkResult: profiling: dict[str, Any] | None = None -@dataclass -class AccuracyConfiguration: - scorer: type[Scorer] - extractor: type[Extractor] | None - dataset_name: str - dataset: Dataset - report_dir: Path - ground_truth_column: str | None - num_repeats: int - extras: dict[str, Any] = field(default_factory=dict) - # Discriminates the inline perf-scored entry (PERFORMANCE) from real accuracy - # datasets (ACCURACY). Branch on this, not on dataset_name == "performance": - # a dataset legitimately named "performance" must not be misclassified. - dataset_type: DatasetType = DatasetType.ACCURACY - - @dataclass class BenchmarkContext: """All state needed to run a benchmark, created by setup_benchmark. @@ -631,27 +605,6 @@ def _build_phases( return phases -def _load_final_snapshot_from_disk(path: Path) -> dict[str, Any] | None: - """Read the persisted ``final_snapshot.json`` written by the aggregator. - - Returns the snapshot in its dict form — the same shape produced by - ``snapshot_to_dict`` and consumed by ``Report.from_snapshot``. No - intermediate Struct decode (see ``Report.from_snapshot`` docstring - for why the dict shape is the consumer contract). - - Returns ``None`` if the file is missing (the aggregator was killed - by an uncatchable signal — SIGKILL, OOM-kill — before its handler - could write) or unreadable. - """ - if not path.exists(): - return None - try: - return json.loads(path.read_bytes()) - except Exception as e: # noqa: BLE001 — best-effort. - logger.warning("Failed to read final snapshot %s: %s", path, e) - return None - - class _PerfPhaseTimeout: """Session-stop timer that bounds the PERFORMANCE phase only. @@ -686,110 +639,92 @@ def cancel(self) -> None: self._handle = None -# (start_path, stop_path) for each supported inference engine's profiling -# protocol. Add a row when introducing a new ProfilerEngine variant. -_PROFILE_PATHS: dict[ProfilerEngine, tuple[str, str]] = { - ProfilerEngine.VLLM: ("/start_profile", "/stop_profile"), -} - +async def _create_issuer( + ctx: BenchmarkContext, loop: asyncio.AbstractEventLoop +) -> tuple[HttpClientSampleIssuer, HTTPEndpointClient]: + """Create the HTTP endpoint client + sample issuer, or raise SetupError.""" + config = ctx.config + endpoints = config.endpoint_config.endpoints + logger.info(f"Connecting: {endpoints}") + try: + api_type: APIType = config.endpoint_config.api_type + # client.api_type is propagated from endpoint_config.api_type by + # BenchmarkConfig._propagate_client_api_type — no override needed here. + client_overrides: dict = { + "endpoint_urls": [ + urljoin(e.rstrip("/") + "/", api_type.default_route()) + for e in endpoints + ], + "api_key": config.endpoint_config.api_key, + "event_logs_dir": ctx.report_dir, + "cpu_affinity": ctx.affinity_plan, + } + if ctx.accuracy_only: + # Single-stream (num_workers=1, max_connections=1) is baked into + # config in setup_benchmark so it is persisted to config.yaml; + # no runtime override needed here. + logger.info( + "Accuracy-only: single-stream (1 worker, 1 connection) for " + "deterministic ordering" + ) + http_config = config.settings.client.with_updates(**client_overrides) + http_client = await HTTPEndpointClient.create(http_config, loop) + issuer = HttpClientSampleIssuer(http_client) + return issuer, http_client + except Exception as e: + raise SetupError(f"Failed to connect to endpoint: {e}") from e -def _derive_profile_urls( - endpoints: list[str], engine: ProfilerEngine, action: str -) -> list[str]: - """One profile URL per endpoint, derived from the engine's HTTP protocol. - For vLLM: strip a trailing ``/v1`` from each endpoint and append - ``/{start,stop}_profile``. ``action`` is ``"start"`` or ``"stop"``. - """ - if not endpoints: - raise ValueError( - f"profiling.engine={engine.value} but endpoint_config.endpoints " - f"is empty; cannot derive {action} URLs" - ) - start_path, stop_path = _PROFILE_PATHS[engine] - path = start_path if action == "start" else stop_path - urls: list[str] = [] - for ep in endpoints: - base = ep.rstrip("/") - if base.endswith("/v1"): - base = base[:-3] - urls.append(f"{base.rstrip('/')}{path}") - return urls - - -def _post_profile(url: str) -> dict[str, Any]: - """POST {url} with empty body; never raises. Returns a record dict suitable - for report.txt rendering and profiling.json serialization.""" - record: dict[str, Any] = { - "url": url, - "sent_at_ns": time.monotonic_ns(), - "sent_at_iso": datetime.now().isoformat(timespec="milliseconds"), - "status": None, - "error": None, - } - req = urllib_request.Request(url, method="POST", data=b"") - try: - with urllib_request.urlopen(req, timeout=2) as resp: - record["status"] = resp.status - except urllib_error.HTTPError as e: - record["status"] = e.code - record["error"] = f"{e.code} {e.reason}" - except Exception as e: # noqa: BLE001 — profile failures must never abort a run - record["error"] = f"{type(e).__name__}: {e}" - return record - - -def _render_profile_status(rec: dict[str, Any]) -> str: - status = rec.get("status") - error = rec.get("error") - if status == 200: - return "200 OK" - if status == 404: - return ( - "404 (profiling not enabled on server — pass " - "--profiler-config.profiler=... to server)" - ) - if error: - return error - if status is not None: - return str(status) - return "ERROR" - - -def _write_profiling_section(f: TextIO, profiling: dict[str, Any]) -> None: - """Append the Profiling section to report.txt (called after report.display).""" - starts = profiling.get("starts", []) - stops = profiling.get("stops", []) - f.write("\n------------------- Profiling -------------------\n") - f.write(f"Engine: {profiling.get('engine', 'unknown')}\n") - f.write("Start:\n") - for rec in starts: - f.write( - f" POST {rec['url']} @ {rec['sent_at_iso']} → " - f"{_render_profile_status(rec)}\n" +def _build_agentic_strategy( + ctx: BenchmarkContext, +) -> AgenticInferenceStrategy | None: + """Build the agentic inference strategy when the perf dataset uses it.""" + if not isinstance(ctx.dataloader, AgenticInferenceDataset): + return None + agentic_cfg = None + if ctx.config.datasets: + perf_ds_cfg = next( + (d for d in ctx.config.datasets if d.type == DatasetType.PERFORMANCE), + None, ) - if stops: - f.write("Stop:\n") - for rec in stops: - suffix = ( - " (from abort handler)" if rec.get("stop_reason") == "abort" else "" - ) - f.write( - f" POST {rec['url']} @ {rec['sent_at_iso']} → " - f"{_render_profile_status(rec)}{suffix}\n" - ) - if starts and stops: - first_start = min(r["sent_at_ns"] for r in starts) - last_stop = max(r["sent_at_ns"] for r in stops) - f.write(f"Trigger span: {(last_stop - first_start) / 1e9:.2f} s\n") - f.write( - "\nNote: actual trace window is bounded by server-side " - "--profiler-config.delay_iterations and " - "--profiler-config.max_iterations.\n" - "Trace artifact path is in server stdout.\n" + if perf_ds_cfg is not None: + agentic_cfg = perf_ds_cfg.agentic_inference + assert ctx.dataloader.conversation_metadata is not None + return AgenticInferenceStrategy( + conversation_manager=ConversationManager(), + dataset_metadata=ctx.dataloader.conversation_metadata, + agentic_inference_config=agentic_cfg, + target_concurrency=ctx.config.settings.load_pattern.target_concurrency, ) +def _wire_on_sample_complete( + collector: ResponseCollector, + agentic_inference_strategy: AgenticInferenceStrategy | None, + publisher: Any, +) -> Callable[[QueryResult], None]: + """Compose the per-sample completion callback (agentic strategy + collector).""" + if agentic_inference_strategy is None: + return collector.on_complete_hook + + def _on_sample_complete(result: QueryResult) -> None: + try: + agentic_inference_strategy.on_sample_complete(result) + except Exception: + logger.exception( + "agentic_inference_strategy.on_sample_complete failed (result=%s)", + result.id, + ) + try: + collector.on_complete_hook(result) + except Exception: + logger.exception("collector.on_complete_hook failed (result=%s)", result.id) + + agentic_inference_strategy._session_on_sample_complete = _on_sample_complete + agentic_inference_strategy._session_publisher = publisher + return _on_sample_complete + + async def _run_benchmark_async( ctx: BenchmarkContext, loop: asyncio.AbstractEventLoop, @@ -806,218 +741,73 @@ async def _run_benchmark_async( ) collector = ResponseCollector(collect_responses=ctx.collect_responses, pbar=pbar) - # ZMQ context for event publishing + service launcher - tmpfs_dir: Path | None = None + # Tmpfs for high-frequency writes (event log); execute owns its lifecycle + # (salvage + rmtree). metrics_output_dir lives on disk under the report dir so + # the final snapshot is preserved with the rest of the run artifacts — it is + # NOT tmpfs and is never removed here. Paths are computed here (no mkdir) so the + # cleanup `except` can always reference tmpfs_dir; the mkdirs run inside the try + # so a mkdir failure that already created tmpfs_dir is still salvaged/removed. + shm = Path("/dev/shm") + tmpfs_base = shm if shm.exists() else Path(tempfile.gettempdir()) + tmpfs_dir = tmpfs_base / f"benchmark_{session_id}" + event_log_dir = tmpfs_dir / "events" + metrics_output_dir = ctx.report_dir / "metrics" + + pipe = MetricsPipeline( + config, + tokenizer_name=ctx.tokenizer_name, + enable_streaming=ctx.enable_streaming, + event_log_dir=event_log_dir, + metrics_output_dir=metrics_output_dir, + loop=loop, + ) + report: Report | None = None + profiler: ProfileController + try: - with ManagedZMQContext.scoped(io_threads=2) as zmq_ctx: - # Event publisher - publisher = EventPublisherService(zmq_ctx) - pub_socket_name = publisher.socket_name - - # Tmpfs for high-frequency writes (event log). - shm = Path("/dev/shm") - use_shm = shm.exists() - tmpfs_base = shm if use_shm else Path(tempfile.gettempdir()) - tmpfs_dir = tmpfs_base / f"benchmark_{session_id}" + try: tmpfs_dir.mkdir(parents=True, exist_ok=True) - - event_log_dir = tmpfs_dir / "events" event_log_dir.mkdir(parents=True, exist_ok=True) - - # Metrics-snapshot output (disk fallback for the final snapshot). - # Lives under the report dir so it's preserved with the rest of - # the run artifacts. - metrics_output_dir = ctx.report_dir / "metrics" metrics_output_dir.mkdir(parents=True, exist_ok=True) - - metrics_socket_name = f"metrics_pub_{uuid.uuid4().hex[:8]}" - - # Connect the metrics-snapshot subscriber BEFORE launching the - # aggregator subprocess that binds the matching PUB socket. ZMQ - # tolerates connect-before-bind on IPC (the connect resolves once - # the binder appears), and starting the SUB reader early gives - # the subscription handshake time to complete during the - # ~1-2 second subprocess-launch window. This eliminates the - # slow-joiner risk of dropping early live ticks (or the worst - # case: missing COMPLETE if the SUB handshake never warms up). - if zmq_ctx.socket_dir is None: - raise RuntimeError("ZMQ socket_dir must be set after publisher bind") - metrics_subscriber = MetricsSnapshotSubscriber( - metrics_socket_name, zmq_ctx, loop - ) - metrics_subscriber.start() - - # Launch service subprocesses - launcher = ServiceLauncher(zmq_ctx) - aggregator_args: list[str] = [ - "--socket-dir", - zmq_ctx.socket_dir, - "--socket-name", - pub_socket_name, - "--metrics-socket", - metrics_socket_name, - "--metrics-output-dir", - str(metrics_output_dir), - ] - if ctx.enable_streaming: - aggregator_args.append("--streaming") - if ctx.tokenizer_name is not None: - aggregator_args.extend(["--tokenizer", ctx.tokenizer_name]) - aggregator_args.extend( - ["--drain-timeout", str(config.settings.drain.metrics_drain_timeout_s)] - ) - aggregator_args.extend( - [ - "--tokenizer-workers", - str(config.settings.drain.metrics_tokenizer_workers), - ] + # Pre-derive profile URLs (inside the run scope so a bad config — + # engine set but no endpoints — fails before the run yet still triggers + # the tmpfs/pbar cleanup in the except/finally below). + profiler = ProfileController( + config.settings.profiling.engine, + config.endpoint_config.endpoints, + config.settings.profiling.urls, ) + await pipe.start() + # start() guarantees the publisher exists; narrow it for the type checker. + publisher = pipe.publisher + assert publisher is not None if config.settings.early_stopping.enabled: aggregator_args.append("--early-stopping") - # EventLoggerService writes events.jsonl to tmpfs (high-frequency writes) - event_logger_args: list[str] = [ - "--log-dir", - str(event_log_dir), - "--socket-dir", - zmq_ctx.socket_dir, - "--socket-name", - pub_socket_name, - "--writers", - "jsonl", - ] - - await launcher.launch( - [ - ServiceConfig( - module="inference_endpoint.async_utils.services.metrics_aggregator", - args=aggregator_args, - ), - ServiceConfig( - module="inference_endpoint.async_utils.services.event_logger", - args=event_logger_args, - ), - ], - timeout=config.settings.service_ready_timeout_s, - ) - - # Create endpoint client on the shared loop - endpoints = config.endpoint_config.endpoints - logger.info(f"Connecting: {endpoints}") - http_client: HTTPEndpointClient | None = None try: - api_type: APIType = config.endpoint_config.api_type - # client.api_type is propagated from endpoint_config.api_type by - # BenchmarkConfig._propagate_client_api_type — no override needed here. - client_overrides: dict = { - "endpoint_urls": [ - urljoin(e.rstrip("/") + "/", api_type.default_route()) - for e in endpoints - ], - "api_key": config.endpoint_config.api_key, - "event_logs_dir": ctx.report_dir, - "cpu_affinity": ctx.affinity_plan, - } - if ctx.accuracy_only: - # Single-stream (num_workers=1, max_connections=1) is baked into - # config in setup_benchmark so it is persisted to config.yaml; - # no runtime override needed here. - logger.info( - "Accuracy-only: single-stream (1 worker, 1 connection) for " - "deterministic ordering" - ) - http_config = config.settings.client.with_updates(**client_overrides) - http_client = await HTTPEndpointClient.create(http_config, loop) - issuer = HttpClientSampleIssuer(http_client) - except Exception as e: - pbar.close() - publisher.close() - launcher.kill_all() - raise SetupError(f"Failed to connect to endpoint: {e}") from e - - # Build agentic inference strategy if the performance dataset uses it. - agentic_inference_strategy: AgenticInferenceStrategy | None = None - if isinstance(ctx.dataloader, AgenticInferenceDataset): - agentic_cfg = None - if ctx.config.datasets: - perf_ds_cfg = next( - ( - d - for d in ctx.config.datasets - if d.type == DatasetType.PERFORMANCE - ), - None, - ) - if perf_ds_cfg is not None: - agentic_cfg = perf_ds_cfg.agentic_inference - assert ctx.dataloader.conversation_metadata is not None - agentic_inference_strategy = AgenticInferenceStrategy( - conversation_manager=ConversationManager(), - dataset_metadata=ctx.dataloader.conversation_metadata, - agentic_inference_config=agentic_cfg, - target_concurrency=ctx.config.settings.load_pattern.target_concurrency, - ) - - _on_sample_complete: Callable[[QueryResult], None] - if agentic_inference_strategy is not None: - - def _on_sample_complete(result: QueryResult) -> None: - try: - agentic_inference_strategy.on_sample_complete(result) - except Exception: - logger.exception( - "agentic_inference_strategy.on_sample_complete failed (result=%s)", - result.id, - ) - try: - collector.on_complete_hook(result) - except Exception: - logger.exception( - "collector.on_complete_hook failed (result=%s)", result.id - ) - - agentic_inference_strategy._session_on_sample_complete = ( - _on_sample_complete - ) - agentic_inference_strategy._session_publisher = publisher - - else: - _on_sample_complete = collector.on_complete_hook + issuer, http_client = await _create_issuer(ctx, loop) + except SetupError: + await pipe.abort() + raise + + agentic_inference_strategy = _build_agentic_strategy(ctx) + on_sample_complete = _wire_on_sample_complete( + collector, agentic_inference_strategy, publisher + ) - # Create session session = BenchmarkSession( issuer=issuer, event_publisher=publisher, loop=loop, - on_sample_complete=_on_sample_complete, + on_sample_complete=on_sample_complete, session_id=session_id, ) - phases = _build_phases(ctx, perf_strategy=agentic_inference_strategy) - report: Report | None = None - _timeout_done = False max_duration_ms = ( ctx.rt_settings.max_duration_ms if ctx.rt_settings is not None else None ) - - # Profile trigger state. Pre-derive URLs once so a bad config - # (engine set but no endpoints) fails before the run. - profiling_cfg = config.settings.profiling - profile_start_urls: list[str] = [] - profile_stop_urls: list[str] = [] - profile_starts: list[dict[str, Any]] = [] - profile_stops: list[dict[str, Any]] = [] - if profiling_cfg.engine is not None: - profile_endpoints = ( - profiling_cfg.urls or config.endpoint_config.endpoints - ) - profile_start_urls = _derive_profile_urls( - profile_endpoints, profiling_cfg.engine, "start" - ) - profile_stop_urls = _derive_profile_urls( - profile_endpoints, profiling_cfg.engine, "stop" - ) + _timeout_done = False session_completed_normally = False def _on_global_timeout() -> None: @@ -1041,20 +831,8 @@ def _on_phase_start(phase: PhaseConfig) -> None: if phase.phase_type != PhaseType.PERFORMANCE: return # Fire /start_profile sequentially before any perf request is - # issued, so the server is armed when traffic begins. Blocks - # the loop briefly (sub-100ms per URL); strategy task hasn't - # been created yet so nothing is starved. - for url in profile_start_urls: - rec = _post_profile(url) - if rec["status"] == 200: - logger.info("Profile start: %s -> 200 OK", url) - else: - logger.warning( - "Profile start: %s -> %s", - url, - rec["error"] or rec["status"], - ) - profile_starts.append(rec) + # issued, so the server is armed when traffic begins. + profiler.start() loop.add_signal_handler(signal.SIGINT, session.stop) try: @@ -1067,134 +845,40 @@ def _on_phase_start(phase: PhaseConfig) -> None: perf_timeout.cancel() loop.remove_signal_handler(signal.SIGINT) # Fire /stop_profile for URLs whose /start_profile succeeded. - # Unifies the clean phase-end path and the abort path — - # both reach this block, both fire stops. - if profile_starts: - stop_reason = "phase_end" if session_completed_normally else "abort" - for i, start_rec in enumerate(profile_starts): - if start_rec["status"] != 200 or i >= len(profile_stop_urls): - continue - rec = _post_profile(profile_stop_urls[i]) - rec["stop_reason"] = stop_reason - if rec["status"] == 200: - logger.info( - "Profile stop: %s -> 200 OK", profile_stop_urls[i] - ) - else: - logger.warning( - "Profile stop: %s -> %s", - profile_stop_urls[i], - rec["error"] or rec["status"], - ) - profile_stops.append(rec) + # Unifies the clean phase-end path and the abort path — both reach + # this block. + profiler.stop(session_completed_normally) logger.info("Cleaning up...") try: if http_client: await http_client.shutdown_async() except Exception as e: logger.warning(f"Client cleanup error: {e}") - logger.info( - "Closing publisher (buffer=%d, pending=%d)...", - publisher.buffered_count, - publisher.pending_count, - ) - publisher.close() - logger.info("Waiting for services to finish processing...") - await asyncio.to_thread(launcher.wait_for_exit, None) - - # Source the snapshot dict for Report: - # 1. Preferred: the JSON file the aggregator atomically wrote - # in publish_final (ENDED-driven or signal-handler-driven). - # 2. Fallback: convert the last live snapshot from pub/sub to - # its dict form. Only reached when the aggregator was killed - # by an uncatchable signal (SIGKILL / OOM) before its - # handler could write. Report will be marked incomplete - # because state will be LIVE / DRAINING, not "complete". - snap_dict: dict[str, Any] | None = _load_final_snapshot_from_disk( - metrics_output_dir / "final_snapshot.json" - ) - if snap_dict is not None: - logger.info("Built report from final_snapshot.json") - elif metrics_subscriber.latest is not None: - snap_dict = snapshot_to_dict(metrics_subscriber.latest) - logger.warning( - "No final_snapshot.json on disk; falling back to last " - "pub/sub snapshot (state may or may not be terminal)" - ) - else: - logger.error("No metrics snapshot available; cannot build report") - - if snap_dict is not None: - try: - load_pattern = ctx.config.settings.load_pattern - runtime_cfg = ctx.config.settings.runtime - # load_pattern + warmup config and the RNG seeds, so - # result_summary.json is self-describing and a valid run is - # identified by its settings. The full, re-runnable config - # lives in config.yaml alongside. The resolved/effective - # runtime settings (sample count + ordering, which can differ - # per audit phase) are deferred to a follow-up. endpoint_config - # (api_key/URLs) is a sibling of settings and never included, - # so no secrets. - run_config = ctx.config.settings.model_dump( - mode="json", include={"load_pattern", "warmup"} - ) - run_config["scheduler_random_seed"] = ( - runtime_cfg.scheduler_random_seed - ) - run_config["dataloader_random_seed"] = ( - runtime_cfg.dataloader_random_seed - ) - report = Report.from_snapshot( - snap_dict, - run_config=run_config, - use_legacy_loadgen_qps_metrics=( - load_pattern.type == LoadPatternType.POISSON - and load_pattern.use_legacy_loadgen_qps_metrics - ), - ) - if not report.complete: - logger.warning( - "Report is incomplete (state=%s, n_pending_tasks=%d)", - report.state, - snap_dict.get("n_pending_tasks", 0), - ) - if report.legacy_loadgen_window_duration_ns is not None: - logger.warning( - "Reporting QPS/TPS with the legacy MLPerf LoadGen Server " - "'completed' definition (deprecated; to be removed once a " - "formal tail-cutting mechanism lands). Pass " - "--no-use-legacy-loadgen-qps-metrics for endpoints-native " - "metrics." - ) - except Exception as e: # noqa: BLE001 — best-effort report build. - logger.warning(f"Failed to build report from snapshot: {e}") - - metrics_subscriber.close() + # Graceful drain runs on both the clean-finish and session-failure + # paths (BenchmarkSession.run publishes ENDED in its own finally, so + # a failed run still has a terminal snapshot worth draining). + report = await pipe.drain_and_build_report() + finally: + # Cleanup runs on every path. pipe.close() (ZMQ scope exit) must run + # even if pbar.close() raises, and a failure here propagates to the + # outer except → tmpfs salvage — matching the monolith's guarantees + # (pbar.close was inside the ZMQ `with`, whose __exit__ always ran). + try: pbar.close() + finally: + pipe.close() except BaseException: - # tmpfs_dir may still be None if the exception hit before it was - # created (e.g. ZMQ context setup), in which case there is nothing - # to clean up. - if tmpfs_dir is not None and tmpfs_dir.exists(): + if tmpfs_dir.exists(): _salvage_tmpfs(ctx.report_dir, tmpfs_dir) shutil.rmtree(tmpfs_dir, ignore_errors=True) raise - profiling_payload: dict[str, Any] | None = None - if profiling_cfg.engine is not None: - profiling_payload = { - "engine": profiling_cfg.engine.value, - "starts": profile_starts, - "stops": profile_stops, - } - return BenchmarkResult( session=result, collector=collector, report=report, tmpfs_dir=tmpfs_dir, - profiling=profiling_payload, + profiling=profiler.payload(), ) @@ -1247,280 +931,68 @@ def _salvage_tmpfs(report_dir: Path, tmpfs_dir: Path) -> None: logger.debug(f"Copied {src_events} -> {dst_events}") -def _phase_osl_stats( - sample_uuids: Iterable[str], - uuid_to_text: dict[str, str], - backend: Any, - batch_size: int = 256, -) -> dict[str, Any] | None: - """Output-token-length rollup over one accuracy phase's completions. - - Counts tokens on each sample's response text via the shared reference - tokenizer backend — the server's ``completion_tokens`` is not persisted, only - the text is (in ``events.jsonl``) — then shapes the lengths via - ``series_metric_dict`` so the block matches the perf report's - ``output_sequence_lengths`` exactly. Returns ``None`` when the phase has no - completed outputs. - - ``batch_size`` bounds each ``encode_batch`` pass: accuracy outputs can be tens - of thousands of tokens each (e.g. gpt-oss lcb at 32768), so counting the whole - population in one call would hold every Encoding in memory at once. - """ - # Skip empty/failed completions (a failed request still logs a COMPLETE - # event with output == ""). The perf-side OslTrigger does the same - # (metrics_table.OslTrigger._extract_text returns None for empty text), so - # accuracy OSL matches its population and a failure isn't counted as a - # 0-token sample that would drag min/avg down. - texts = [ - uuid_to_text[u] for u in sample_uuids if u in uuid_to_text and uuid_to_text[u] - ] - if not texts: - return None - lengths: list[int] = [] - for i in range(0, len(texts), batch_size): - lengths.extend(encode_lengths(backend, texts[i : i + batch_size])) - return series_metric_dict(lengths) or None - - -def _phase_response_counts( - sample_uuids: Iterable[str], - uuid_to_text: dict[str, str], -) -> dict[str, int]: - """Per-phase response accounting over one accuracy phase's issued samples. - - Complements :func:`_phase_osl_stats`, which reports token lengths only over - non-empty completions — on its own that can hide a run where the server - returned blanks or dropped requests. Classifies each issued ``sample_uuid`` - as ``scored`` (COMPLETE, non-empty output — exactly the OSL population), - ``empty`` (COMPLETE with blank output: a failed request the load generator - logged as ERROR then an empty COMPLETE), or ``missing`` (no COMPLETE event). - ``issued == scored + empty + missing`` always holds. - - Emptiness uses the same truthiness test as ``_phase_osl_stats`` so ``scored`` - is byte-for-byte the OSL population — the two blocks cannot disagree. - """ - issued = scored = empty = missing = 0 - for u in sample_uuids: - issued += 1 - if u not in uuid_to_text: - missing += 1 - elif uuid_to_text[u]: - scored += 1 - else: - empty += 1 - return {"issued": issued, "scored": scored, "empty": empty, "missing": missing} - - -def _accuracy_uuid_bound( - report_dir: Path | None, eval_configs: list[AccuracyConfiguration] -) -> set[str]: - """Union of the accuracy datasets' issued uuids from ``sample_idx_map.json``. +def _write_report_artifacts( + ctx: BenchmarkContext, report: Report, profiling: dict[str, Any] | None +) -> None: + """Display the report and write result_summary.json + report.txt. - Bounds the finalize-side raw-output read to the accuracy population so it - never holds the whole run's (incl. perf) response-text corpus. Returns an - empty set (⇒ caller reads unbounded) when there is no report dir; a missing, - corrupt, or wrong-shape map is warned and also falls back to unbounded. - """ - if report_dir is None: - return set() - try: - idx_map = msgspec.json.decode((report_dir / "sample_idx_map.json").read_bytes()) - except (OSError, msgspec.DecodeError) as e: - logger.warning( - "Accuracy OSL uuid bound unavailable (%s); reading outputs unbounded", e - ) - return set() - # A syntactically-valid map of the wrong shape must not crash finalize: this - # runs outside the per-dataset try, so a raised AttributeError/TypeError would - # fail scoring (OSL must never do that). Skip anything not dict-shaped. - if not isinstance(idx_map, dict): - logger.warning( - "Accuracy OSL uuid bound: sample_idx_map.json is not an object; " - "reading outputs unbounded" - ) - return set() - bound: set[str] = set() - for ec in eval_configs: - if ec.dataset_type == DatasetType.ACCURACY: - per_dataset = idx_map.get(ec.dataset_name) - if isinstance(per_dataset, dict): - bound |= set(per_dataset) - return bound - - -def _score_accuracy( - ctx: BenchmarkContext, result: SessionResult -) -> list[dict[str, Any]]: - """Score each accuracy dataset into its own list entry. - - One entry per eval_config, in order; no cross-dataset consolidation. Each - entry carries the scalar ``score`` plus sample accounting - (``unit_samples`` × ``num_repeats`` = ``total_samples``); a scorer that - returns a ``score_breakdown()`` (DeepSeek-R1, BFCL) also attaches - ``breakdown``. The ``"performance"`` inline entry totals the perf phases' - issued counts instead of unit × repeats (repeats is forced to 1 there). + result_summary.json is the self-complete machine-readable report (carries + qps/tps/seeds/accuracy via Report.to_json); report.txt is the full + human-readable dump; the console log shows the summary. """ - accuracy_scores: list[dict[str, Any]] = [] + report.display(fn=lambda s: logger.info(s), summary_only=True) + performance_dir = ctx.report_dir / "performance" + performance_dir.mkdir(parents=True, exist_ok=True) + report.to_json(save_to=performance_dir / "result_summary.json") - # Per-phase wall-clock (seconds) keyed by phase name. The accuracy phase name - # is the dataset_name; the inline-scored perf entry keys on "performance". - phase_durations: dict[str, float] = {} - for pr in result.phase_results: - phase_durations[pr.name] = phase_durations.get(pr.name, 0.0) + max( - 0.0, (pr.end_time_ns - pr.start_time_ns) / 1e9 - ) + report_txt = ctx.report_dir / "report.txt" + with report_txt.open("w") as f: + report.display(fn=lambda s: print(s, file=f)) + if profiling is not None: + _write_profiling_section(f, profiling) + logger.info("Report written to %s", report_txt) - # Accuracy-phase output-token lengths (finalize-side, off the hot path): the - # aggregator only tokenizes perf-window samples, so count the accuracy - # responses (already in events.jsonl) here, using the same reference tokenizer - # as the perf side. (Counts still differ from perf for tool-call responses — - # client-side OSL is approximate for structured output.) Loaded only when a - # real accuracy dataset exists; a load failure or a tokenizer with no fast - # backend disables OSL rather than failing scoring. - has_accuracy = any( - ec.dataset_type == DatasetType.ACCURACY for ec in ctx.eval_configs - ) - osl_backend: Any = None - if has_accuracy and ctx.tokenizer_name is not None: - try: - osl_backend = load_reference_backend(ctx.tokenizer_name) - except Exception as e: # noqa: BLE001 - OSL is optional; never fail scoring - logger.warning( - "Accuracy OSL disabled: could not load tokenizer %r: %s", - ctx.tokenizer_name, - e, - ) - else: - # A tokenizer with no fast (Rust) backend disables OSL rather than - # falling back to a slow Python-tokenizer count: the perf side - # (token_metrics._setup_shards) requires a fast backend too and raises - # without one, so OSL stays fast-only and consistent on both sides. - # Warn so the skip is visible instead of silently dropping the block. - if osl_backend is None: - logger.warning( - "Accuracy OSL disabled: tokenizer %r has no fast (Rust) backend " - "(token counting requires one, as on the perf side)", - ctx.tokenizer_name, - ) - # Bound the raw-output read to the accuracy population so finalize never holds - # the whole run's (incl. perf) response-text corpus. - accuracy_uuids = ( - _accuracy_uuid_bound(ctx.report_dir, ctx.eval_configs) - if has_accuracy - else set() - ) - uuid_to_text: dict[str, str] | None = None - for eval_cfg in ctx.eval_configs: - try: - scorer_instance = eval_cfg.scorer( - eval_cfg.dataset_name, - eval_cfg.dataset, - eval_cfg.report_dir, - extractor=eval_cfg.extractor, - ground_truth_column=eval_cfg.ground_truth_column, - **eval_cfg.extras, - ) - except TypeError as e: - raise InputValidationError( - f"Dataset '{eval_cfg.dataset_name}': invalid accuracy_config.extras " - f"for scorer '{eval_cfg.scorer.__name__}': {e}" - ) from e - score, n_repeats = scorer_instance.score() - # Coerce a numpy scalar score (np.float32/64, numpy ints — e.g. np.mean - # from the base Scorer) to a native Python float so the entry stays - # serializable by both msgspec (result_summary.json) and json - # (accuracy_results.json). numbers.Real catches every numpy scalar (not - # just np.float64, which isinstance(..., float) alone would miss) while - # leaving None / dict (RougeScorer) untouched; bool is excluded. - if isinstance(score, numbers.Real) and not isinstance(score, bool): - score = float(score) - unit_samples = eval_cfg.dataset.num_samples() - num_repeats = eval_cfg.num_repeats - if eval_cfg.dataset_type == DatasetType.PERFORMANCE: - # A performance dataset always scores its already-issued outputs once - # (enforced by the num_repeats == 1 guard in _load_datasets), so make - # that locally provable rather than relying on eval_cfg carrying 1. - num_repeats = 1 - total_samples = sum(phase.issued_count for phase in result.perf_results) +def _summarize_and_log_metrics( + ctx: BenchmarkContext, + report: Report | None, + result: SessionResult, + collector: ResponseCollector, +) -> None: + """Log the run's headline metrics, preferring Report over SessionResult.""" + if report is not None and report.duration_ns is not None: + perf_elapsed = report.duration_ns / 1e9 + total_issued = report.n_samples_issued + n_errors = report.n_samples_failed + qps = report.qps or 0.0 + else: + perf = result.perf_results[0] if result.perf_results else None + if perf: + perf_elapsed = (perf.end_time_ns - perf.start_time_ns) / 1e9 + total_issued = perf.issued_count else: - total_samples = unit_samples * num_repeats - entry: dict[str, Any] = { - "dataset_name": eval_cfg.dataset_name, - "extractor": ( - eval_cfg.extractor.__name__ if eval_cfg.extractor is not None else None - ), - "ground_truth_column": eval_cfg.ground_truth_column, - "score": score, - "unit_samples": unit_samples, - "num_repeats": num_repeats, - "total_samples": total_samples, - # Wall-clock of this dataset's issue phase (seconds); 0.0 if the - # phase left no timing (e.g. a scored-but-not-issued dataset). - "duration_s": round(phase_durations.get(eval_cfg.dataset_name, 0.0), 3), - # False when the scorer produced only a partial headline (e.g. - # LegacyMLPerfDeepSeekR1Scorer when lcb-service was unreachable), so a - # partial number is never mistaken for a complete one. - "complete": scorer_instance.complete, - # Persist the same DatasetType discriminator carried on the eval config - # so consumers filter the inline perf-scored entry by type, not by - # matching dataset_name == "performance". - "dataset_type": eval_cfg.dataset_type.value, - } - breakdown = scorer_instance.score_breakdown() - if breakdown is not None: - entry["breakdown"] = breakdown - - # Response accounting + avg/min/max output-token length. Skipped for the - # perf entry (its OSL / failure counts live in result_summary.json). The - # counts are computed independent of the tokenizer and of OSL returning a - # block — an all-failed phase must still publish scored=0 rather than - # silently omitting everything. OSL stays tokenizer-gated. A read/tokenize - # failure only drops these blocks — it never fails scoring. - if eval_cfg.dataset_type == DatasetType.ACCURACY: - try: - if uuid_to_text is None: - # Built once from the first scorer and reused for every - # dataset. get_raw_outputs() returns the model's actual - # completion text (not the scorer's scoring-normalized form) - # for *all* phases' COMPLETE events, bounded to the accuracy - # population; intersecting it with each dataset's - # sample_index_map yields correct per-dataset counts. - out_df = scorer_instance.get_raw_outputs(accuracy_uuids or None) - uuid_to_text = dict( - zip(out_df["sample_uuid"], out_df["output"], strict=False) - ) - # Drop the DataFrame so finalize doesn't hold both it and the - # dict (each carrying the response-text corpus). - del out_df - entry["response_counts"] = _phase_response_counts( - scorer_instance.sample_index_map, uuid_to_text - ) - if osl_backend is not None: - t0 = time.perf_counter() - osl = _phase_osl_stats( - scorer_instance.sample_index_map, uuid_to_text, osl_backend - ) - if osl is not None: - # Same shape/key as the perf report output_sequence_lengths. - entry["output_sequence_lengths"] = osl - # Wall-clock of just this phase's tokenization (seconds); - # summed across datasets for the accuracy report's total. - entry["osl_tokenize_s"] = round(time.perf_counter() - t0, 3) - except Exception as e: # noqa: BLE001 - optional blocks; never fail scoring - logger.warning( - "Accuracy response counts/OSL skipped for %s: %s", - eval_cfg.dataset_name, - e, - ) + perf_elapsed = (result.end_time_ns - result.start_time_ns) / 1e9 + total_issued = 0 + n_errors = len(collector.errors) + qps = total_issued / perf_elapsed if perf_elapsed > 0 else 0.0 - accuracy_scores.append(entry) + logger.info(f"Completed in {perf_elapsed:.1f}s") + if ctx.accuracy_only: + acc_total = sum(ds.num_samples() * ds.repeats for ds in ctx.accuracy_datasets) + logger.info(f"Accuracy-only: {acc_total} samples evaluated") + else: logger.info( - f"Score for {eval_cfg.dataset_name}: {score} " - f"({n_repeats} repeats, complete={scorer_instance.complete})" + f"Results: {max(0, total_issued - n_errors)}/{total_issued} successful" ) + if qps > 0: + logger.info(f"Estimated QPS: {qps:.1f}") - return accuracy_scores + if collector.errors: + logger.warning(f"Errors: {len(collector.errors)}") + for err in collector.errors[:3]: + logger.debug(f" {err}") + if len(collector.errors) > 3: + logger.debug(f" ... +{len(collector.errors) - 3} more") def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: @@ -1540,12 +1012,12 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: # sample_idx_map.json + events.jsonl from here). _write_scoring_artifacts(ctx, result, bench.tmpfs_dir) - # Accuracy scoring (one entry per accuracy dataset). - # Scoring runs before the report is written so the accuracy headline can be - # attached, but the report is written in the `finally` below so a scoring - # failure (e.g. lcb-service unreachable, missing eval subproject, bad extras) - # still leaves the perf run's result_summary.json / report.txt on disk - # instead of discarding them — then the exception propagates as before. + # Accuracy scoring (one entry per accuracy dataset). Scoring runs before the + # report is written so the accuracy headline can be attached, but the report + # is written in the `finally` below so a scoring failure (e.g. lcb-service + # unreachable, missing eval subproject, bad extras) still leaves the perf + # run's result_summary.json / report.txt on disk instead of discarding them — + # then the exception propagates as before. accuracy_scores: list[dict[str, Any]] = [] try: accuracy_scores = _score_accuracy(ctx, result) @@ -1555,102 +1027,15 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: # failure). if report is not None: report = msgspec.structs.replace(report, accuracy=accuracy_scores) - - # Display report if available (from MetricsAggregator pub/sub snapshot). - # result_summary.json is the self-complete machine-readable report - # (carries qps/tps/seeds/accuracy via Report.to_json); report.txt is the - # full human-readable dump; the console log shows the summary. + # Display the report + write result_summary.json / report.txt. if report is not None: - report.display(fn=lambda s: logger.info(s), summary_only=True) - performance_dir = ctx.report_dir / "performance" - performance_dir.mkdir(parents=True, exist_ok=True) - report.to_json(save_to=performance_dir / "result_summary.json") - - report_txt = ctx.report_dir / "report.txt" - with report_txt.open("w") as f: - report.display(fn=lambda s: print(s, file=f)) - if bench.profiling is not None: - _write_profiling_section(f, bench.profiling) - logger.info("Report written to %s", report_txt) - - # Report metrics: prefer Report from MetricsSnapshot, fall back to SessionResult - if report is not None and report.duration_ns is not None: - perf_elapsed = report.duration_ns / 1e9 - total_issued = report.n_samples_issued - n_errors = report.n_samples_failed - qps = report.qps or 0.0 - else: - perf = result.perf_results[0] if result.perf_results else None - if perf: - perf_elapsed = (perf.end_time_ns - perf.start_time_ns) / 1e9 - total_issued = perf.issued_count - else: - perf_elapsed = (result.end_time_ns - result.start_time_ns) / 1e9 - total_issued = 0 - n_errors = len(collector.errors) - qps = total_issued / perf_elapsed if perf_elapsed > 0 else 0.0 + _write_report_artifacts(ctx, report, bench.profiling) - logger.info(f"Completed in {perf_elapsed:.1f}s") - if ctx.accuracy_only: - acc_total = sum(ds.num_samples() * ds.repeats for ds in ctx.accuracy_datasets) - logger.info(f"Accuracy-only: {acc_total} samples evaluated") - else: - logger.info( - f"Results: {max(0, total_issued - n_errors)}/{total_issued} successful" - ) - if qps > 0: - logger.info(f"Estimated QPS: {qps:.1f}") + _summarize_and_log_metrics(ctx, report, result, collector) - if collector.errors: - logger.warning(f"Errors: {len(collector.errors)}") - for err in collector.errors[:3]: - logger.debug(f" {err}") - if len(collector.errors) > 3: - logger.debug(f" ... +{len(collector.errors) - 3} more") - - # Emit the accuracy results as a focused artifact under accuracy/. Perf - # rollups (qps/tps/latency percentiles) live in performance/result_summary.json - # and response/error text lives in events.jsonl, so neither is duplicated - # here. Written only when scoring produced entries — a perf-only run leaves - # no accuracy/ folder. - if accuracy_scores: - # Plain cross-component mean of the per-dataset scores (3 datasets for - # gpt-oss, 1 for DeepSeek-R1); None when nothing numeric was scored. - avg_accuracy = average_accuracy(accuracy_scores) - # Total finalize-time spent tokenizing accuracy outputs for OSL (seconds), - # summed across datasets. Emitted whenever OSL was computed for at least - # one dataset — gating on the key's presence, not the rounded wall-clock, - # so a sub-millisecond total (tiny outputs) still records 0.0 rather than - # silently dropping the field. - osl_computed = any("osl_tokenize_s" in e for e in accuracy_scores) - osl_tokenization_s = round( - sum(e.get("osl_tokenize_s", 0.0) for e in accuracy_scores), 3 - ) - accuracy_dir = ctx.report_dir / "accuracy" - accuracy_results_path = accuracy_dir / "accuracy_results.json" - accuracy_payload: dict[str, Any] = {} - if avg_accuracy is not None: - accuracy_payload["average_accuracy"] = avg_accuracy - if osl_computed: - accuracy_payload["osl_tokenization_s"] = osl_tokenization_s - accuracy_payload["accuracy_scores"] = accuracy_scores - # Atomic write so a crash mid-write can't leave truncated JSON the - # compliance checker would read as corrupt. Not swallowed: if scoring - # produced entries but they can't be persisted — the dir can't be made - # (OSError), the payload won't serialize (TypeError/ValueError, e.g. a - # numpy scalar left in a breakdown block), or the write fails — fail the - # run loudly rather than exit 0 with no accuracy artifact. - try: - accuracy_dir.mkdir(parents=True, exist_ok=True) - atomic_write_bytes( - accuracy_results_path, - json.dumps(accuracy_payload, indent=2).encode(), - ) - except (OSError, TypeError, ValueError) as e: - raise ExecutionError( - f"Failed to write accuracy results to {accuracy_results_path}: {e}" - ) from e - logger.info(f"Saved: {accuracy_results_path}") + # Emit the accuracy results as a focused artifact under accuracy/. Written + # after the report artifacts so a write failure here can't discard them. + write_accuracy_results(ctx.report_dir, accuracy_scores) def run_benchmark( diff --git a/src/inference_endpoint/commands/benchmark/pipeline.py b/src/inference_endpoint/commands/benchmark/pipeline.py new file mode 100644 index 00000000..47847268 --- /dev/null +++ b/src/inference_endpoint/commands/benchmark/pipeline.py @@ -0,0 +1,345 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Metrics/event-log service pipeline for a benchmark run. + +``MetricsPipeline`` owns the ZMQ context, the event publisher, the metrics-snapshot +subscriber, and the two service subprocesses (metrics aggregator + event logger). +It exposes explicit lifecycle methods rather than a context manager because the run +has three distinct teardown paths that a single ``__aexit__`` cannot express cleanly: + +* ``start()`` — bring the infrastructure up; unwind partially-acquired resources if + service launch fails (so a launch error can't leak the ZMQ context). +* ``drain_and_build_report()`` — the graceful end-of-run drain (close publisher → + wait for services → source the final snapshot → build the Report). Runs on both a + clean finish and a session-run failure, since ``BenchmarkSession.run`` publishes + ``ENDED`` in its own ``finally`` and the aggregator writes a terminal snapshot. +* ``abort()`` — the connect-failure fast path: kill the services without a graceful + drain (no ``ENDED`` was ever published). + +``close()`` exits the ZMQ scope and is idempotent, so the orchestrator can call it +unconditionally in its ``finally``. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import uuid +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from inference_endpoint.async_utils.event_publisher import EventPublisherService +from inference_endpoint.async_utils.services.launcher import ( + ServiceConfig, + ServiceLauncher, +) +from inference_endpoint.async_utils.services.metrics_aggregator.snapshot import ( + snapshot_to_dict, +) +from inference_endpoint.async_utils.services.metrics_aggregator.subscriber import ( + MetricsSnapshotSubscriber, +) +from inference_endpoint.async_utils.transport.zmq.context import ManagedZMQContext +from inference_endpoint.config.schema import LoadPatternType +from inference_endpoint.metrics.report import Report + +if TYPE_CHECKING: + from inference_endpoint.config.schema import BenchmarkConfig + +logger = logging.getLogger(__name__) + +_AGGREGATOR_MODULE = "inference_endpoint.async_utils.services.metrics_aggregator" +_EVENT_LOGGER_MODULE = "inference_endpoint.async_utils.services.event_logger" + + +def _load_final_snapshot_from_disk(path: Path) -> dict[str, Any] | None: + """Read the persisted ``final_snapshot.json`` written by the aggregator. + + Returns the snapshot in its dict form — the same shape produced by + ``snapshot_to_dict`` and consumed by ``Report.from_snapshot``. No + intermediate Struct decode (see ``Report.from_snapshot`` docstring + for why the dict shape is the consumer contract). + + Returns ``None`` if the file is missing (the aggregator was killed + by an uncatchable signal — SIGKILL, OOM-kill — before its handler + could write) or unreadable. + """ + if not path.exists(): + return None + try: + return json.loads(path.read_bytes()) + except Exception as e: # noqa: BLE001 — best-effort. + logger.warning("Failed to read final snapshot %s: %s", path, e) + return None + + +def _build_aggregator_args( + *, + socket_dir: str, + pub_socket_name: str, + metrics_socket_name: str, + metrics_output_dir: Path, + enable_streaming: bool, + tokenizer_name: str | None, + drain_timeout_s: float, + tokenizer_workers: int, +) -> list[str]: + """CLI args for the metrics_aggregator subprocess.""" + args: list[str] = [ + "--socket-dir", + socket_dir, + "--socket-name", + pub_socket_name, + "--metrics-socket", + metrics_socket_name, + "--metrics-output-dir", + str(metrics_output_dir), + ] + if enable_streaming: + args.append("--streaming") + if tokenizer_name is not None: + args.extend(["--tokenizer", tokenizer_name]) + args.extend(["--drain-timeout", str(drain_timeout_s)]) + args.extend(["--tokenizer-workers", str(tokenizer_workers)]) + return args + + +def _build_event_logger_args( + *, event_log_dir: Path, socket_dir: str, pub_socket_name: str +) -> list[str]: + """CLI args for the event_logger subprocess (writes events.jsonl to tmpfs).""" + return [ + "--log-dir", + str(event_log_dir), + "--socket-dir", + socket_dir, + "--socket-name", + pub_socket_name, + "--writers", + "jsonl", + ] + + +def _build_report_from_snapshot( + snap_dict: dict[str, Any], config: BenchmarkConfig +) -> Report | None: + """Build the Report from a metrics snapshot dict; best-effort (None on failure). + + A malformed snapshot must never fail an otherwise-completed benchmark, so any + build error is logged and swallowed to ``None``. + """ + try: + load_pattern = config.settings.load_pattern + runtime_cfg = config.settings.runtime + # load_pattern + warmup config and the RNG seeds, so result_summary.json is + # self-describing and a valid run is identified by its settings. The full, + # re-runnable config lives in config.yaml alongside. The resolved/effective + # runtime settings (sample count + ordering, which can differ per audit + # phase) are deferred to a follow-up. endpoint_config (api_key/URLs) is a + # sibling of settings and never included, so no secrets. + run_config = config.settings.model_dump( + mode="json", include={"load_pattern", "warmup"} + ) + run_config["scheduler_random_seed"] = runtime_cfg.scheduler_random_seed + run_config["dataloader_random_seed"] = runtime_cfg.dataloader_random_seed + report = Report.from_snapshot( + snap_dict, + run_config=run_config, + use_legacy_loadgen_qps_metrics=( + load_pattern.type == LoadPatternType.POISSON + and load_pattern.use_legacy_loadgen_qps_metrics + ), + ) + if not report.complete: + logger.warning( + "Report is incomplete (state=%s, n_pending_tasks=%d)", + report.state, + snap_dict.get("n_pending_tasks", 0), + ) + if report.legacy_loadgen_window_duration_ns is not None: + logger.warning( + "Reporting QPS/TPS with the legacy MLPerf LoadGen Server " + "'completed' definition (deprecated; to be removed once a " + "formal tail-cutting mechanism lands). Pass " + "--no-use-legacy-loadgen-qps-metrics for endpoints-native " + "metrics." + ) + return report + except Exception as e: # noqa: BLE001 — best-effort report build. + logger.warning(f"Failed to build report from snapshot: {e}") + return None + + +class MetricsPipeline: + """ZMQ + metrics/event-logger subprocess lifecycle for one benchmark run. + + ``event_log_dir`` is on tmpfs (salvaged then removed by the caller); + ``metrics_output_dir`` is on disk under the report dir (holds + ``final_snapshot.json`` — the primary Report source) and is never removed here. + """ + + def __init__( + self, + config: BenchmarkConfig, + *, + tokenizer_name: str | None, + enable_streaming: bool, + event_log_dir: Path, + metrics_output_dir: Path, + loop: asyncio.AbstractEventLoop, + ) -> None: + self._config = config + self._tokenizer_name = tokenizer_name + self._enable_streaming = enable_streaming + self._event_log_dir = event_log_dir + self._metrics_output_dir = metrics_output_dir + self._loop = loop + + self._zmq_cm: Any = None + self._launcher: ServiceLauncher | None = None + self.publisher: EventPublisherService | None = None + self.subscriber: MetricsSnapshotSubscriber | None = None + self._closed = False + + async def start(self) -> None: + """Bring up ZMQ + publisher + subscriber + service subprocesses. + + Connects the snapshot subscriber BEFORE launching the aggregator that binds + the matching PUB socket (ZMQ tolerates connect-before-bind on IPC; starting + the SUB reader early lets the subscription handshake complete during the + ~1-2s subprocess-launch window, avoiding the slow-joiner risk of dropping + early live ticks — or, worst case, missing COMPLETE). On any failure the + partially-acquired resources are unwound so nothing (notably the ZMQ + context) leaks. + """ + self._zmq_cm = ManagedZMQContext.scoped(io_threads=2) + zmq_ctx = self._zmq_cm.__enter__() + try: + self.publisher = EventPublisherService(zmq_ctx) + pub_socket_name = self.publisher.socket_name + metrics_socket_name = f"metrics_pub_{uuid.uuid4().hex[:8]}" + + if zmq_ctx.socket_dir is None: + raise RuntimeError("ZMQ socket_dir must be set after publisher bind") + self.subscriber = MetricsSnapshotSubscriber( + metrics_socket_name, zmq_ctx, self._loop + ) + self.subscriber.start() + + self._launcher = ServiceLauncher(zmq_ctx) + drain = self._config.settings.drain + aggregator_args = _build_aggregator_args( + socket_dir=zmq_ctx.socket_dir, + pub_socket_name=pub_socket_name, + metrics_socket_name=metrics_socket_name, + metrics_output_dir=self._metrics_output_dir, + enable_streaming=self._enable_streaming, + tokenizer_name=self._tokenizer_name, + drain_timeout_s=drain.metrics_drain_timeout_s, + tokenizer_workers=drain.metrics_tokenizer_workers, + ) + event_logger_args = _build_event_logger_args( + event_log_dir=self._event_log_dir, + socket_dir=zmq_ctx.socket_dir, + pub_socket_name=pub_socket_name, + ) + await self._launcher.launch( + [ + ServiceConfig(module=_AGGREGATOR_MODULE, args=aggregator_args), + ServiceConfig(module=_EVENT_LOGGER_MODULE, args=event_logger_args), + ], + timeout=self._config.settings.service_ready_timeout_s, + ) + except BaseException: + self._teardown(kill=True) + raise + + async def drain_and_build_report(self) -> Report | None: + """Graceful drain: close publisher, wait for services, build the Report. + + Runs on both the clean-finish and session-failure paths. Sources the + snapshot from the aggregator's on-disk ``final_snapshot.json`` when present, + else falls back to the last live pub/sub snapshot (only reached when the + aggregator was killed before it could write). Report construction is + best-effort — a build failure yields ``None`` rather than aborting. + """ + assert self.publisher is not None + assert self._launcher is not None + logger.info( + "Closing publisher (buffer=%d, pending=%d)...", + self.publisher.buffered_count, + self.publisher.pending_count, + ) + self.publisher.close() + logger.info("Waiting for services to finish processing...") + await asyncio.to_thread(self._launcher.wait_for_exit, None) + + snap_dict = _load_final_snapshot_from_disk( + self._metrics_output_dir / "final_snapshot.json" + ) + if snap_dict is not None: + logger.info("Built report from final_snapshot.json") + elif self.subscriber is not None and self.subscriber.latest is not None: + snap_dict = snapshot_to_dict(self.subscriber.latest) + logger.warning( + "No final_snapshot.json on disk; falling back to last " + "pub/sub snapshot (state may or may not be terminal)" + ) + else: + logger.error("No metrics snapshot available; cannot build report") + + report = ( + _build_report_from_snapshot(snap_dict, self._config) + if snap_dict is not None + else None + ) + # Null the closed handles so the caller's unconditional close() does not + # re-close them (close() is idempotent, but this keeps the invariant clean). + self.publisher = None + if self.subscriber is not None: + self.subscriber.close() + self.subscriber = None + return report + + async def abort(self) -> None: + """Connect-failure fast path: kill services without a graceful drain.""" + self._teardown(kill=True) + + def close(self) -> None: + """Exit the ZMQ scope (idempotent) — safe to call unconditionally.""" + self._teardown(kill=False) + + def _teardown(self, *, kill: bool) -> None: + if self._closed: + return + self._closed = True + if self.publisher is not None: + try: + self.publisher.close() + except Exception as e: # noqa: BLE001 — teardown best-effort + logger.warning("Publisher close error: %s", e) + if kill and self._launcher is not None: + self._launcher.kill_all() + if self.subscriber is not None: + try: + self.subscriber.close() + except Exception as e: # noqa: BLE001 — teardown best-effort + logger.warning("Subscriber close error: %s", e) + self.subscriber = None + if self._zmq_cm is not None: + self._zmq_cm.__exit__(None, None, None) + self._zmq_cm = None diff --git a/src/inference_endpoint/commands/benchmark/profiling.py b/src/inference_endpoint/commands/benchmark/profiling.py new file mode 100644 index 00000000..3a5f8595 --- /dev/null +++ b/src/inference_endpoint/commands/benchmark/profiling.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Profiler trigger protocol for benchmark runs (vLLM /start_profile, /stop_profile). + +``ProfileController`` owns the whole trigger lifecycle: it pre-derives the per-endpoint +URLs up front (so a misconfigured run fails before any traffic), fires ``/start_profile`` +when the performance phase begins, fires ``/stop_profile`` for every start that succeeded, +and produces the ``{engine, starts, stops}`` payload rendered into ``report.txt`` and a +sibling ``profiling.json``. +""" + +from __future__ import annotations + +import logging +import time +from datetime import datetime +from typing import Any, TextIO +from urllib import error as urllib_error +from urllib import request as urllib_request + +from inference_endpoint.config.schema import ProfilerEngine + +logger = logging.getLogger(__name__) + + +# (start_path, stop_path) for each supported inference engine's profiling +# protocol. Add a row when introducing a new ProfilerEngine variant. +_PROFILE_PATHS: dict[ProfilerEngine, tuple[str, str]] = { + ProfilerEngine.VLLM: ("/start_profile", "/stop_profile"), +} + + +def _derive_profile_urls( + endpoints: list[str], engine: ProfilerEngine, action: str +) -> list[str]: + """One profile URL per endpoint, derived from the engine's HTTP protocol. + + For vLLM: strip a trailing ``/v1`` from each endpoint and append + ``/{start,stop}_profile``. ``action`` is ``"start"`` or ``"stop"``. + """ + if not endpoints: + raise ValueError( + f"profiling.engine={engine.value} but endpoint_config.endpoints " + f"is empty; cannot derive {action} URLs" + ) + start_path, stop_path = _PROFILE_PATHS[engine] + path = start_path if action == "start" else stop_path + urls: list[str] = [] + for ep in endpoints: + base = ep.rstrip("/") + if base.endswith("/v1"): + base = base[:-3] + urls.append(f"{base.rstrip('/')}{path}") + return urls + + +def _post_profile(url: str) -> dict[str, Any]: + """POST {url} with empty body; never raises. Returns a record dict suitable + for report.txt rendering and profiling.json serialization.""" + record: dict[str, Any] = { + "url": url, + "sent_at_ns": time.monotonic_ns(), + "sent_at_iso": datetime.now().isoformat(timespec="milliseconds"), + "status": None, + "error": None, + } + req = urllib_request.Request(url, method="POST", data=b"") + try: + with urllib_request.urlopen(req, timeout=2) as resp: + record["status"] = resp.status + except urllib_error.HTTPError as e: + record["status"] = e.code + record["error"] = f"{e.code} {e.reason}" + except Exception as e: # noqa: BLE001 — profile failures must never abort a run + record["error"] = f"{type(e).__name__}: {e}" + return record + + +def _render_profile_status(rec: dict[str, Any]) -> str: + status = rec.get("status") + error = rec.get("error") + if status == 200: + return "200 OK" + if status == 404: + return ( + "404 (profiling not enabled on server — pass " + "--profiler-config.profiler=... to server)" + ) + if error: + return error + if status is not None: + return str(status) + return "ERROR" + + +def _write_profiling_section(f: TextIO, profiling: dict[str, Any]) -> None: + """Append the Profiling section to report.txt (called after report.display).""" + starts = profiling.get("starts", []) + stops = profiling.get("stops", []) + f.write("\n------------------- Profiling -------------------\n") + f.write(f"Engine: {profiling.get('engine', 'unknown')}\n") + f.write("Start:\n") + for rec in starts: + f.write( + f" POST {rec['url']} @ {rec['sent_at_iso']} → " + f"{_render_profile_status(rec)}\n" + ) + if stops: + f.write("Stop:\n") + for rec in stops: + suffix = ( + " (from abort handler)" if rec.get("stop_reason") == "abort" else "" + ) + f.write( + f" POST {rec['url']} @ {rec['sent_at_iso']} → " + f"{_render_profile_status(rec)}{suffix}\n" + ) + if starts and stops: + first_start = min(r["sent_at_ns"] for r in starts) + last_stop = max(r["sent_at_ns"] for r in stops) + f.write(f"Trigger span: {(last_stop - first_start) / 1e9:.2f} s\n") + f.write( + "\nNote: actual trace window is bounded by server-side " + "--profiler-config.delay_iterations and " + "--profiler-config.max_iterations.\n" + "Trace artifact path is in server stdout.\n" + ) + + +class ProfileController: + """Owns the profiler start/stop trigger lifecycle for one benchmark run. + + Disabled (a no-op) when ``engine is None``. When enabled it derives the start + and stop URLs at construction — so an engine set with no endpoints fails before + the run — and only fires ``/stop_profile`` for the ``/start_profile`` calls that + returned 200. + """ + + def __init__( + self, + engine: ProfilerEngine | None, + endpoints: list[str], + urls_override: list[str] | None, + ) -> None: + self._engine = engine + self._start_urls: list[str] = [] + self._stop_urls: list[str] = [] + self._starts: list[dict[str, Any]] = [] + self._stops: list[dict[str, Any]] = [] + if engine is not None: + # Preserve the truthiness fallback: an explicit empty urls override + # falls through to endpoint_config.endpoints. + profile_endpoints = urls_override or endpoints + self._start_urls = _derive_profile_urls(profile_endpoints, engine, "start") + self._stop_urls = _derive_profile_urls(profile_endpoints, engine, "stop") + + @property + def enabled(self) -> bool: + return self._engine is not None + + def start(self) -> None: + """Fire /start_profile sequentially before any perf request is issued.""" + for url in self._start_urls: + rec = _post_profile(url) + if rec["status"] == 200: + logger.info("Profile start: %s -> 200 OK", url) + else: + logger.warning( + "Profile start: %s -> %s", url, rec["error"] or rec["status"] + ) + self._starts.append(rec) + + def stop(self, completed_normally: bool) -> None: + """Fire /stop_profile for every start that returned 200. + + Unifies the clean phase-end path and the abort path — both call this. + """ + if not self._starts: + return + stop_reason = "phase_end" if completed_normally else "abort" + for i, start_rec in enumerate(self._starts): + if start_rec["status"] != 200 or i >= len(self._stop_urls): + continue + rec = _post_profile(self._stop_urls[i]) + rec["stop_reason"] = stop_reason + if rec["status"] == 200: + logger.info("Profile stop: %s -> 200 OK", self._stop_urls[i]) + else: + logger.warning( + "Profile stop: %s -> %s", + self._stop_urls[i], + rec["error"] or rec["status"], + ) + self._stops.append(rec) + + def payload(self) -> dict[str, Any] | None: + """The {engine, starts, stops} record, or None when profiling is disabled. + + A payload is emitted whenever an engine is configured, even if no + performance phase fired (no starts/stops) — matching the pre-refactor + behavior where the payload's presence tracks the engine, not the triggers. + """ + if self._engine is None: + return None + return { + "engine": self._engine.value, + "starts": self._starts, + "stops": self._stops, + } diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index 491dc72b..83c3ad74 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -40,14 +40,16 @@ BenchmarkContext, ResponseCollector, _build_phases, - _derive_profile_urls, _load_datasets, _PerfPhaseTimeout, + _run_benchmark_async, + setup_benchmark, +) +from inference_endpoint.commands.benchmark.profiling import ( + _derive_profile_urls, _post_profile, _render_profile_status, - _run_benchmark_async, _write_profiling_section, - setup_benchmark, ) from inference_endpoint.config.runtime_settings import RuntimeSettings from inference_endpoint.config.schema import ( @@ -677,16 +679,16 @@ async def _capture_launch(service_configs, *, timeout): with ( patch( - "inference_endpoint.commands.benchmark.execute.ManagedZMQContext" + "inference_endpoint.commands.benchmark.pipeline.ManagedZMQContext" ) as MockZMQ, patch( - "inference_endpoint.commands.benchmark.execute.EventPublisherService" + "inference_endpoint.commands.benchmark.pipeline.EventPublisherService" ) as MockPub, patch( - "inference_endpoint.commands.benchmark.execute.MetricsSnapshotSubscriber" + "inference_endpoint.commands.benchmark.pipeline.MetricsSnapshotSubscriber" ) as MockSub, patch( - "inference_endpoint.commands.benchmark.execute.ServiceLauncher" + "inference_endpoint.commands.benchmark.pipeline.ServiceLauncher" ) as MockLauncher, patch("inference_endpoint.commands.benchmark.execute.tqdm"), ): @@ -727,16 +729,16 @@ async def _capture_launch(service_configs, *, timeout): with ( patch( - "inference_endpoint.commands.benchmark.execute.ManagedZMQContext" + "inference_endpoint.commands.benchmark.pipeline.ManagedZMQContext" ) as MockZMQ, patch( - "inference_endpoint.commands.benchmark.execute.EventPublisherService" + "inference_endpoint.commands.benchmark.pipeline.EventPublisherService" ) as MockPub, patch( - "inference_endpoint.commands.benchmark.execute.MetricsSnapshotSubscriber" + "inference_endpoint.commands.benchmark.pipeline.MetricsSnapshotSubscriber" ) as MockSub, patch( - "inference_endpoint.commands.benchmark.execute.ServiceLauncher" + "inference_endpoint.commands.benchmark.pipeline.ServiceLauncher" ) as MockLauncher, patch("inference_endpoint.commands.benchmark.execute.tqdm"), ): @@ -781,16 +783,16 @@ async def _capture_launch(service_configs, *, timeout): with ( patch( - "inference_endpoint.commands.benchmark.execute.ManagedZMQContext" + "inference_endpoint.commands.benchmark.pipeline.ManagedZMQContext" ) as MockZMQ, patch( - "inference_endpoint.commands.benchmark.execute.EventPublisherService" + "inference_endpoint.commands.benchmark.pipeline.EventPublisherService" ) as MockPub, patch( - "inference_endpoint.commands.benchmark.execute.MetricsSnapshotSubscriber" + "inference_endpoint.commands.benchmark.pipeline.MetricsSnapshotSubscriber" ) as MockSub, patch( - "inference_endpoint.commands.benchmark.execute.ServiceLauncher" + "inference_endpoint.commands.benchmark.pipeline.ServiceLauncher" ) as MockLauncher, patch("inference_endpoint.commands.benchmark.execute.tqdm"), ): @@ -1855,7 +1857,7 @@ def test_post_profile_200(self): resp = MagicMock() resp.__enter__.return_value.status = 200 with patch( - "inference_endpoint.commands.benchmark.execute.urllib_request.urlopen", + "inference_endpoint.commands.benchmark.profiling.urllib_request.urlopen", return_value=resp, ): rec = _post_profile("http://h/start_profile") @@ -1868,7 +1870,7 @@ def test_post_profile_200(self): def test_post_profile_http_error(self): err = urllib_error.HTTPError("http://h", 404, "Not Found", {}, None) with patch( - "inference_endpoint.commands.benchmark.execute.urllib_request.urlopen", + "inference_endpoint.commands.benchmark.profiling.urllib_request.urlopen", side_effect=err, ): rec = _post_profile("http://h/start_profile") @@ -1878,7 +1880,7 @@ def test_post_profile_http_error(self): @pytest.mark.unit def test_post_profile_connection_failure_never_raises(self): with patch( - "inference_endpoint.commands.benchmark.execute.urllib_request.urlopen", + "inference_endpoint.commands.benchmark.profiling.urllib_request.urlopen", side_effect=OSError("refused"), ): rec = _post_profile("http://h/start_profile") diff --git a/tests/unit/commands/test_benchmark_final_snapshot.py b/tests/unit/commands/test_benchmark_final_snapshot.py index e3eac7b5..897c7b27 100644 --- a/tests/unit/commands/test_benchmark_final_snapshot.py +++ b/tests/unit/commands/test_benchmark_final_snapshot.py @@ -33,7 +33,7 @@ from inference_endpoint.async_utils.services.metrics_aggregator.snapshot import ( SessionState, ) -from inference_endpoint.commands.benchmark.execute import ( +from inference_endpoint.commands.benchmark.pipeline import ( _load_final_snapshot_from_disk, ) from inference_endpoint.metrics.report import Report diff --git a/tests/unit/commands/test_score_accuracy.py b/tests/unit/commands/test_score_accuracy.py index 1ca7a0be..aaef0d0f 100644 --- a/tests/unit/commands/test_score_accuracy.py +++ b/tests/unit/commands/test_score_accuracy.py @@ -22,7 +22,7 @@ import msgspec.json import pytest -from inference_endpoint.commands.benchmark.execute import ( +from inference_endpoint.commands.benchmark.accuracy import ( AccuracyConfiguration, _accuracy_uuid_bound, _phase_osl_stats, @@ -31,11 +31,11 @@ ) from inference_endpoint.config.schema import DatasetType -# Module object for the tests that monkeypatch execute's own module-level symbols -# (load_reference_backend) so _score_accuracy resolves the patched one. Taken from -# sys.modules to avoid importing execute under both the ``from ... import`` and -# ``import ...`` styles. -execute_mod = sys.modules[_score_accuracy.__module__] +# The module that defines _score_accuracy (commands.benchmark.accuracy), resolved +# via __module__ so the tests monkeypatch its module-level load_reference_backend +# and _score_accuracy picks up the patched one — without importing the module under +# both the ``from ... import`` and ``import ...`` styles. +scoring_mod = sys.modules[_score_accuracy.__module__] class _FakeDataset: @@ -246,7 +246,7 @@ def test_osl_attached_with_tokenizer(self, tmp_path, monkeypatch): """With a tokenizer, each accuracy entry gets an output_sequence_lengths block (same shape as the perf report) from the phase's completions.""" monkeypatch.setattr( - execute_mod, "load_reference_backend", lambda name: _WordBackend() + scoring_mod, "load_reference_backend", lambda name: _WordBackend() ) cfg = _cfg("aime25::gptoss", 2, 0.8, tmp_path, scorer=_FakeOSLScorer) entry = _by_name(_score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ @@ -265,7 +265,7 @@ def test_osl_attached_with_tokenizer(self, tmp_path, monkeypatch): def test_osl_skipped_for_performance_entry(self, tmp_path, monkeypatch): monkeypatch.setattr( - execute_mod, "load_reference_backend", lambda name: _WordBackend() + scoring_mod, "load_reference_backend", lambda name: _WordBackend() ) cfg = _cfg("performance", 2, 0.6, tmp_path, scorer=_FakeOSLScorer) entry = _by_name(_score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ @@ -276,7 +276,7 @@ def test_osl_skipped_for_performance_entry(self, tmp_path, monkeypatch): def test_osl_dropped_when_get_raw_outputs_raises(self, tmp_path, monkeypatch): """A read/tokenize failure only drops OSL — scoring still succeeds.""" monkeypatch.setattr( - execute_mod, "load_reference_backend", lambda name: _WordBackend() + scoring_mod, "load_reference_backend", lambda name: _WordBackend() ) class _RaisingScorer(_FakeOSLScorer): @@ -308,7 +308,7 @@ def test_response_counts_published_when_all_empty(self, tmp_path, monkeypatch): """Masking regression: every response empty => OSL returns None, but response_counts must still publish scored=0 rather than omitting all.""" monkeypatch.setattr( - execute_mod, "load_reference_backend", lambda name: _WordBackend() + scoring_mod, "load_reference_backend", lambda name: _WordBackend() ) cfg = _cfg("aime25::gptoss", 2, 0.8, tmp_path, scorer=_EmptyOSLScorer) entry = _by_name(_score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ @@ -326,7 +326,7 @@ def test_response_counts_classifies_missing(self, tmp_path, monkeypatch): """scored/empty/missing partition the issued samples; OSL tokenizes only the one scored (non-empty) response.""" monkeypatch.setattr( - execute_mod, "load_reference_backend", lambda name: _WordBackend() + scoring_mod, "load_reference_backend", lambda name: _WordBackend() ) cfg = _cfg("ds", 1, 0.8, tmp_path, scorer=_MixedOSLScorer) entry = _by_name(_score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ @@ -349,7 +349,7 @@ def test_all_missing_still_publishes_response_counts(self, tmp_path, monkeypatch """An accuracy phase with no COMPLETE rows (empty, columned frame) must still publish response_counts (all missing) rather than dropping them.""" monkeypatch.setattr( - execute_mod, "load_reference_backend", lambda name: _WordBackend() + scoring_mod, "load_reference_backend", lambda name: _WordBackend() ) cfg = _cfg("ds", 1, 0.8, tmp_path, scorer=_AllMissingScorer) entry = _by_name(_score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ @@ -366,7 +366,7 @@ def test_all_missing_still_publishes_response_counts(self, tmp_path, monkeypatch def test_no_fast_backend_disables_osl_keeps_counts(self, tmp_path, monkeypatch): """A tokenizer with no fast backend (load_reference_backend -> None) skips OSL but still publishes response_counts.""" - monkeypatch.setattr(execute_mod, "load_reference_backend", lambda name: None) + monkeypatch.setattr(scoring_mod, "load_reference_backend", lambda name: None) cfg = _cfg("aime25::gptoss", 2, 0.8, tmp_path, scorer=_FakeOSLScorer) entry = _by_name(_score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ "aime25::gptoss" @@ -382,7 +382,7 @@ def test_tokenizer_load_failure_disables_osl_not_scoring( def _boom(name): raise OSError("no tokenizer") - monkeypatch.setattr(execute_mod, "load_reference_backend", _boom) + monkeypatch.setattr(scoring_mod, "load_reference_backend", _boom) cfg = _cfg("aime25::gptoss", 2, 0.8, tmp_path, scorer=_FakeOSLScorer) entry = _by_name(_score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ "aime25::gptoss" @@ -395,7 +395,7 @@ def test_uuid_bound_passed_to_get_raw_outputs(self, tmp_path, monkeypatch): """End-to-end: _score_accuracy computes the accuracy uuid bound from sample_idx_map.json and passes it into get_raw_outputs.""" monkeypatch.setattr( - execute_mod, "load_reference_backend", lambda name: _WordBackend() + scoring_mod, "load_reference_backend", lambda name: _WordBackend() ) (tmp_path / "sample_idx_map.json").write_bytes( msgspec.json.encode({"aime25::gptoss": {"u1": 0, "u2": 1}})