From 57d9a66777471cda833f393b40973112d4d35a38 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:55:34 -0500 Subject: [PATCH 01/13] refactor(benchmark): split execute.py into execute/profiling/accuracy/pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit execute.py concentrated all CLI benchmark orchestration; the large _run_benchmark_async and _score_accuracy dominated its complexity. Split along natural seams into four cohesive modules: - profiling.py (new): profile-trigger protocol (vLLM /start_profile,/stop_profile) + ProfileController (URL derivation + start/stop/payload lifecycle) - accuracy.py (new): AccuracyConfiguration, _score_accuracy, _load_osl_backend, write_accuracy_results - pipeline.py (new): MetricsPipeline — ZMQ + metrics-aggregator/event-logger subprocess lifecycle (start/drain_and_build_report/abort/close) + snapshot→Report - execute.py: thin orchestrator Everything audit.py/cli.py import stays in execute (TestMode, BenchmarkResult, _salvage_tmpfs, setup/run/finalize, run_benchmark, resolve_report_dir); one-way import graph (execute → {profiling, accuracy, pipeline}). Three unit test files repointed. Rebased onto main after the SWE-bench scorer (#342). #342's execute.py accuracy changes are routed into their post-split homes: AccuracyConfiguration model_params/endpoint_config fields, _effective_external_sample_count, _accuracy_uuid_bound (None-on-unavailable), and SKIP_ENDPOINT_PHASE handling in _score_accuracy land in accuracy.py; _validate_accuracy_config_for_scorer, the external-scorer skips/logging, and the dataset-loader/preflight wiring land in execute.py. Also restores the early-stopping aggregator flag the split had dropped: "--early-stopping" is now threaded through pipeline._build_aggregator_args (gated on settings.early_stopping.enabled), replacing an orphaned append to a variable that no longer existed in _run_benchmark_async after the split. Verified: unit suite green except pre-existing macOS-only failures (sched_setaffinity / IPC-socket-path in cpu_affinity/token_metrics/metrics_aggregator, none in commands/benchmark). ruff/ruff-format/prettier/license clean; mypy clean on the changed files (remaining mypy errors are the macOS-only os.sched_* attr-defined false-positives in untouched files). Co-Authored-By: Claude Opus 4.8 (1M context) --- AGENTS.md | 5 +- .../commands/benchmark/accuracy.py | 434 +++++++ .../commands/benchmark/execute.py | 1093 ++++------------- .../commands/benchmark/pipeline.py | 349 ++++++ .../commands/benchmark/profiling.py | 222 ++++ tests/unit/commands/test_benchmark.py | 42 +- .../commands/test_benchmark_final_snapshot.py | 2 +- tests/unit/commands/test_score_accuracy.py | 30 +- 8 files changed, 1275 insertions(+), 902 deletions(-) create mode 100644 src/inference_endpoint/commands/benchmark/accuracy.py create mode 100644 src/inference_endpoint/commands/benchmark/pipeline.py create mode 100644 src/inference_endpoint/commands/benchmark/profiling.py diff --git a/AGENTS.md b/AGENTS.md index 829083cc9..5f4ddfb8c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -178,7 +178,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 000000000..dab1de24c --- /dev/null +++ b/src/inference_endpoint/commands/benchmark/accuracy.py @@ -0,0 +1,434 @@ +# 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, ScorerMethod, TestMode +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 + from inference_endpoint.config.schema import EndpointConfig, ModelParams + +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) + model_params: ModelParams | None = None + endpoint_config: EndpointConfig | None = None + # 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 _effective_external_sample_count( + eval_cfg: AccuracyConfiguration, +) -> int | None: + """Clamp an external scorer's requested count to its loaded dataset size.""" + count = eval_cfg.scorer.external_sample_count(eval_cfg.extras) + if count is None: + return None + return min(count, eval_cfg.dataset.num_samples()) + + +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] | None: + """Return issued accuracy UUIDs, or ``None`` when the map is unavailable. + + 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 + ``None`` (⇒ caller reads unbounded) when there is no report dir or the map is + missing, corrupt, or the wrong shape. An empty set is a valid bound for + scorers that skip endpoint issuance. + """ + if report_dir is None: + return None + 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 None + # 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 None + 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]]: + """Run configured scorers and return reportable accuracy entries. + + 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]] = [] + eval_configs = [ + eval_cfg + for eval_cfg in ctx.eval_configs + if not (ctx.test_mode == TestMode.PERF and eval_cfg.scorer.SKIP_ENDPOINT_PHASE) + ] + + # 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 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, eval_configs) if has_accuracy else set() + ) + uuid_to_text: dict[str, str] | None = None + + for eval_cfg in eval_configs: + try: + scorer_kwargs = dict(eval_cfg.extras) + if ( + getattr(eval_cfg.scorer, "SCORER_ID", None) + == ScorerMethod.SWE_BENCH.value + ): + scorer_kwargs.update( + model_params=eval_cfg.model_params, + endpoint_config=eval_cfg.endpoint_config, + ) + 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, + **scorer_kwargs, + ) + 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 + if eval_cfg.scorer.SKIP_ENDPOINT_PHASE: + ext = _effective_external_sample_count(eval_cfg) + if ext is not None: + unit_samples = ext + total_samples = ext + 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) + 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 f6fedc681..e45fbe253 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,59 +30,45 @@ 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, + _effective_external_sample_count, + _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 ( APIType, BenchmarkConfig, DatasetType, - EndpointConfig, LoadPattern, LoadPatternType, - ModelParams, - ProfilerEngine, - ScorerMethod, StreamingMode, TestMode, TestType, @@ -93,7 +83,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, @@ -110,8 +99,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() @@ -173,34 +161,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) - model_params: ModelParams | None = None - endpoint_config: EndpointConfig | None = None - # 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 _effective_external_sample_count( - eval_cfg: AccuracyConfiguration, -) -> int | None: - """Clamp an external scorer's requested count to its loaded dataset size.""" - count = eval_cfg.scorer.external_sample_count(eval_cfg.extras) - if count is None: - return None - return min(count, eval_cfg.dataset.num_samples()) - - @dataclass class BenchmarkContext: """All state needed to run a benchmark, created by setup_benchmark. @@ -378,6 +338,9 @@ def _load_datasets( num_repeats=acc_cfg.accuracy_config.num_repeats, **loader_kwargs, ) + # Value/api-type validity of the override is already enforced at config + # construction (BenchmarkConfig validates each dataset's effective params), + # so this cannot raise for a validated config. ds_model_params = acc_cfg.effective_generation_config(config.model_params) ds.load(api_type=config.endpoint_config.api_type, model_params=ds_model_params) logger.info(f"Loaded {ds} - {ds.num_samples()} samples") @@ -705,27 +668,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. @@ -760,110 +702,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, @@ -880,218 +804,71 @@ 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), - ] - ) - 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, + # 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 - # 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: @@ -1115,20 +892,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: @@ -1141,134 +906,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(), ) @@ -1324,290 +995,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] | None: - """Return issued accuracy UUIDs, or ``None`` when the map is unavailable. +def _write_report_artifacts( + ctx: BenchmarkContext, report: Report, profiling: dict[str, Any] | None +) -> None: + """Display the report and write result_summary.json + report.txt. - An empty set is a valid bound for scorers that skip endpoint issuance. + 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. """ - if report_dir is None: - return None - 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 None - # Invalid maps fall back to an unbounded read instead of failing scoring. - if not isinstance(idx_map, dict): - logger.warning( - "Accuracy OSL uuid bound: sample_idx_map.json is not an object; " - "reading outputs unbounded" - ) - return None - 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]]: - """Run configured scorers and return reportable accuracy entries. - - 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]] = [] - eval_configs = [ - eval_cfg - for eval_cfg in ctx.eval_configs - if not (ctx.test_mode == TestMode.PERF and eval_cfg.scorer.SKIP_ENDPOINT_PHASE) - ] + 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 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, eval_configs) if has_accuracy else set() - ) - uuid_to_text: dict[str, str] | None = None - for eval_cfg in eval_configs: - try: - scorer_kwargs = dict(eval_cfg.extras) - if ( - getattr(eval_cfg.scorer, "SCORER_ID", None) - == ScorerMethod.SWE_BENCH.value - ): - scorer_kwargs.update( - model_params=eval_cfg.model_params, - endpoint_config=eval_cfg.endpoint_config, - ) - 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, - **scorer_kwargs, - ) - 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 - if eval_cfg.scorer.SKIP_ENDPOINT_PHASE: - ext = _effective_external_sample_count(eval_cfg) - if ext is not None: - unit_samples = ext - total_samples = ext - 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) - 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: @@ -1627,12 +1076,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) @@ -1642,101 +1091,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 - - 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}") + _write_report_artifacts(ctx, report, bench.profiling) - 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") + _summarize_and_log_metrics(ctx, report, result, collector) - # 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 configured scoring produced entries. - 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 000000000..4d3944c65 --- /dev/null +++ b/src/inference_endpoint/commands/benchmark/pipeline.py @@ -0,0 +1,349 @@ +# 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, + early_stopping: bool, +) -> 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 early_stopping: + args.append("--early-stopping") + 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, + early_stopping=self._config.settings.early_stopping.enabled, + ) + 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 000000000..3a5f85957 --- /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 ad4a7de84..0ed51eeb4 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -42,16 +42,18 @@ BenchmarkResult, ResponseCollector, _build_phases, - _derive_profile_urls, _load_datasets, _PerfPhaseTimeout, - _post_profile, - _render_profile_status, _run_benchmark_async, - _write_profiling_section, finalize_benchmark, setup_benchmark, ) +from inference_endpoint.commands.benchmark.profiling import ( + _derive_profile_urls, + _post_profile, + _render_profile_status, + _write_profiling_section, +) from inference_endpoint.config.runtime_settings import RuntimeSettings from inference_endpoint.config.schema import ( BenchmarkConfig, @@ -684,7 +686,7 @@ def test_preflight_error_propagates(self, tmp_path): @pytest.mark.unit @pytest.mark.parametrize( - ("datasets, expected_scorer, expected_type, " "expected_accuracy_datasets"), + ("datasets, expected_scorer, expected_type, expected_accuracy_datasets"), [ ( [ @@ -1142,16 +1144,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"), ): @@ -1192,16 +1194,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"), ): @@ -1246,16 +1248,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"), ): @@ -2619,7 +2621,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") @@ -2632,7 +2634,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") @@ -2642,7 +2644,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 e3eac7b56..897c7b278 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 408791f93..ae21b7499 100644 --- a/tests/unit/commands/test_score_accuracy.py +++ b/tests/unit/commands/test_score_accuracy.py @@ -24,7 +24,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, @@ -39,11 +39,11 @@ ScorerMethod, ) -# 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: @@ -334,7 +334,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))[ @@ -353,7 +353,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))[ @@ -364,7 +364,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): @@ -396,7 +396,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))[ @@ -414,7 +414,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))[ @@ -437,7 +437,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))[ @@ -454,7 +454,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" @@ -470,7 +470,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" @@ -483,7 +483,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}}) From 3070db5b349fe0fd223e25b776a29889efa8aa6b Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:38:13 -0500 Subject: [PATCH 02/13] fix(benchmark): kill metrics services on setup failure after launch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A setup failure between MetricsPipeline.start() (which launches the aggregator + event-logger subprocesses) and session.run — e.g. _build_phases rejecting an agentic accuracy dataset, or BenchmarkSession construction raising — never reached the graceful drain, and the finally called only pipe.close() (_teardown(kill=False)). The launched subprocesses were thus neither drained nor killed; with the aggregator drain-timeout defaulting to unlimited they could linger indefinitely waiting for an ENDED that never arrives. Track whether the drain ran; if not, the finally now aborts (kill=True) so the services are killed. The connect-failure path already did this via pipe.abort() for SetupError; this extends the guarantee to every other post-start() setup failure. _teardown is idempotent (_closed guard), so the existing SetupError->abort path double-calls abort() harmlessly. Pre-existing latent gap (the monolith's ZMQ `with` __exit__ likewise never killed the launcher), surfaced by the review council on the execute.py split. Adds a regression test that drives a post-launch setup failure and asserts kill_all() runs. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../commands/benchmark/execute.py | 11 ++++ tests/unit/commands/test_benchmark.py | 57 ++++++++++++++++++- 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index e45fbe253..4301fc286 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -826,6 +826,7 @@ async def _run_benchmark_async( ) report: Report | None = None profiler: ProfileController + drained = False try: try: @@ -919,6 +920,7 @@ def _on_phase_start(phase: PhaseConfig) -> None: # 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() + drained = True 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 @@ -927,6 +929,15 @@ def _on_phase_start(phase: PhaseConfig) -> None: try: pbar.close() finally: + # A setup failure between pipe.start() and session.run never + # reaches the drain above, so the aggregator/event-logger + # subprocesses were launched but never signalled to end. Kill them + # (abort) so they can't linger — the aggregator's drain-timeout + # defaults to unlimited. On the clean-finish and session-failure + # paths the drain already waited for their exit, so only the ZMQ + # scope needs closing. + if not drained: + await pipe.abort() pipe.close() except BaseException: if tmpfs_dir.exists(): diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index 0ed51eeb4..6b48636af 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -24,7 +24,7 @@ import tempfile from pathlib import Path from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch from urllib import error as urllib_error import inference_endpoint.commands.benchmark.execute as execute_mod @@ -1276,6 +1276,61 @@ async def _capture_launch(service_configs, *, timeout): tmpfs_dir = tmpfs_base / "benchmark_cli_benchmark_deadbeef" assert not tmpfs_dir.exists() + @pytest.mark.unit + @pytest.mark.asyncio + async def test_setup_failure_after_launch_kills_services(self, tmp_path): + """A setup failure *after* the services launch but *before* session.run + must kill the aggregator/event-logger subprocesses, not orphan them. + + The graceful drain (which waits for the services to self-exit on ENDED) + is only reached once session.run runs; a failure before it never sends + ENDED, and with drain-timeout defaulting to unlimited the services would + otherwise linger. Here launch succeeds and _create_issuer then raises a + non-SetupError, so the explicit `except SetupError: pipe.abort()` is + bypassed — the finally must still abort (kill=True). + """ + config = OfflineConfig(**_OFFLINE_KWARGS, settings=OfflineSettings()) + ctx = self._make_ctx(config, tmp_path) + + async def _launch_ok(service_configs, *, timeout): + return None # services are now "running" + + mock_zmq = MagicMock() + mock_zmq.socket_dir = str(tmp_path / "sockets") + + with ( + patch( + "inference_endpoint.commands.benchmark.pipeline.ManagedZMQContext" + ) as MockZMQ, + patch( + "inference_endpoint.commands.benchmark.pipeline.EventPublisherService" + ) as MockPub, + patch( + "inference_endpoint.commands.benchmark.pipeline.MetricsSnapshotSubscriber" + ) as MockSub, + patch( + "inference_endpoint.commands.benchmark.pipeline.ServiceLauncher" + ) as MockLauncher, + patch("inference_endpoint.commands.benchmark.execute.tqdm"), + patch( + "inference_endpoint.commands.benchmark.execute._create_issuer", + new=AsyncMock(side_effect=RuntimeError("setup boom")), + ), + ): + MockZMQ.scoped.return_value.__enter__ = MagicMock(return_value=mock_zmq) + MockZMQ.scoped.return_value.__exit__ = MagicMock(return_value=False) + MockPub.return_value.socket_name = "test_pub" + MockSub.return_value.start = MagicMock() + MockLauncher.return_value.launch = _launch_ok + + loop = asyncio.get_event_loop() + with pytest.raises(RuntimeError, match="setup boom"): + await _run_benchmark_async(ctx, loop) + + # abort() → _teardown(kill=True) → launcher.kill_all(); called exactly + # once (the SetupError branch never ran) and never for a clean run. + MockLauncher.return_value.kill_all.assert_called_once() + class TestAccuracyOnlyDatasetLoading: """`--accuracy-only` must skip the performance dataset even when the config From d0956dfc65e015c224c26c050d5795a2e0373d58 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:46:45 -0500 Subject: [PATCH 03/13] fix(benchmark): make kill_all best-effort in MetricsPipeline teardown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The setup-failure abort added in the previous commit runs `pipe.abort()` → `_teardown(kill=True)` from a `finally` while an exception is propagating. `kill_all()` sat unguarded between the try/except-wrapped publisher and subscriber closes; a raise there (e.g. a child exiting between `poll()` and `kill()`) would mask the original setup error and skip the ZMQ `__exit__` — and since `_teardown` already set `_closed`, the trailing `pipe.close()` would no-op, leaking the ZMQ scope (violating the finally's "close must always run" invariant). Wrap `kill_all()` in the same best-effort try/except as the neighbouring closes so teardown always reaches the ZMQ `__exit__` and never masks the in-flight exception. Surfaced by the review council (Cursor + Claude) on the prior fix. Also adds a companion test: a SetupError after launch hits both abort() sites (explicit except + finally); asserts kill_all runs exactly once, locking in the `_closed` idempotency guard the finally-abort depends on. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../commands/benchmark/pipeline.py | 9 +++- tests/unit/commands/test_benchmark.py | 51 +++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/src/inference_endpoint/commands/benchmark/pipeline.py b/src/inference_endpoint/commands/benchmark/pipeline.py index 4d3944c65..fee21c660 100644 --- a/src/inference_endpoint/commands/benchmark/pipeline.py +++ b/src/inference_endpoint/commands/benchmark/pipeline.py @@ -337,7 +337,14 @@ def _teardown(self, *, kill: bool) -> None: 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() + try: + self._launcher.kill_all() + except Exception as e: # noqa: BLE001 — teardown best-effort + # Must not propagate: _teardown runs from a finally during + # exception unwind, and the ZMQ __exit__ below must still run + # (a raised kill_all would mask the original error and leak the + # ZMQ scope, since _closed is already set so a later close() no-ops). + logger.warning("Service kill_all error: %s", e) if self.subscriber is not None: try: self.subscriber.close() diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index 6b48636af..80255710f 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -1331,6 +1331,57 @@ async def _launch_ok(service_configs, *, timeout): # once (the SetupError branch never ran) and never for a clean run. MockLauncher.return_value.kill_all.assert_called_once() + @pytest.mark.unit + @pytest.mark.asyncio + async def test_setup_error_after_launch_aborts_exactly_once(self, tmp_path): + """A SetupError after launch hits BOTH abort() sites — the explicit + `except SetupError: pipe.abort()` and the finally's `if not drained: + pipe.abort()`. _teardown's `_closed` guard must make the second a no-op, + so kill_all() runs exactly once (not twice). Locks in the idempotency the + finally-abort relies on. + """ + config = OfflineConfig(**_OFFLINE_KWARGS, settings=OfflineSettings()) + ctx = self._make_ctx(config, tmp_path) + + async def _launch_ok(service_configs, *, timeout): + return None + + mock_zmq = MagicMock() + mock_zmq.socket_dir = str(tmp_path / "sockets") + + with ( + patch( + "inference_endpoint.commands.benchmark.pipeline.ManagedZMQContext" + ) as MockZMQ, + patch( + "inference_endpoint.commands.benchmark.pipeline.EventPublisherService" + ) as MockPub, + patch( + "inference_endpoint.commands.benchmark.pipeline.MetricsSnapshotSubscriber" + ) as MockSub, + patch( + "inference_endpoint.commands.benchmark.pipeline.ServiceLauncher" + ) as MockLauncher, + patch("inference_endpoint.commands.benchmark.execute.tqdm"), + patch( + "inference_endpoint.commands.benchmark.execute._create_issuer", + new=AsyncMock(side_effect=SetupError("connect boom")), + ), + ): + MockZMQ.scoped.return_value.__enter__ = MagicMock(return_value=mock_zmq) + MockZMQ.scoped.return_value.__exit__ = MagicMock(return_value=False) + MockPub.return_value.socket_name = "test_pub" + MockSub.return_value.start = MagicMock() + MockLauncher.return_value.launch = _launch_ok + + loop = asyncio.get_event_loop() + with pytest.raises(SetupError, match="connect boom"): + await _run_benchmark_async(ctx, loop) + + # Both abort() sites fire, but _teardown's _closed guard collapses them: + # kill_all runs once, the ZMQ scope is exited once. + MockLauncher.return_value.kill_all.assert_called_once() + class TestAccuracyOnlyDatasetLoading: """`--accuracy-only` must skip the performance dataset even when the config From 3174d26c90db4d2b78bdf2e04fa9f9161a3de3f3 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:06:14 -0500 Subject: [PATCH 04/13] refactor(benchmark): promote split helpers to public names + small cleanups Address PR #421 review comments (renames + one-liners; no behavior change): - Rename cross-module helpers public (drop the `_` prefix now that the execute.py split imports them across a package boundary): effective_external_sample_count, score_accuracy, write_profiling_section - profiling: type `action` as Literal["start", "stop"]; name the profile-POST timeout constant (_PROFILE_POST_TIMEOUT_S); drop the dead ProfileController.enabled property - execute: annotate _wire_on_sample_complete `publisher` as EventPublisherService (TYPE_CHECKING import) - accuracy: mark AccuracyConfiguration frozen=True (build-once/read-many) Co-Authored-By: Claude Opus 4.8 (1M context) --- .../commands/benchmark/accuracy.py | 8 +-- .../commands/benchmark/execute.py | 19 ++++--- .../commands/benchmark/profiling.py | 15 +++--- tests/unit/commands/test_benchmark.py | 4 +- tests/unit/commands/test_score_accuracy.py | 54 +++++++++---------- 5 files changed, 50 insertions(+), 50 deletions(-) diff --git a/src/inference_endpoint/commands/benchmark/accuracy.py b/src/inference_endpoint/commands/benchmark/accuracy.py index dab1de24c..713932079 100644 --- a/src/inference_endpoint/commands/benchmark/accuracy.py +++ b/src/inference_endpoint/commands/benchmark/accuracy.py @@ -55,7 +55,7 @@ logger = logging.getLogger(__name__) -@dataclass +@dataclass(frozen=True) class AccuracyConfiguration: scorer: type[Scorer] extractor: type[Extractor] | None @@ -73,7 +73,7 @@ class AccuracyConfiguration: dataset_type: DatasetType = DatasetType.ACCURACY -def _effective_external_sample_count( +def effective_external_sample_count( eval_cfg: AccuracyConfiguration, ) -> int | None: """Clamp an external scorer's requested count to its loaded dataset size.""" @@ -216,7 +216,7 @@ def _load_osl_backend(has_accuracy: bool, tokenizer_name: str | None) -> Any | N return osl_backend -def _score_accuracy( +def score_accuracy( ctx: BenchmarkContext, result: SessionResult ) -> list[dict[str, Any]]: """Run configured scorers and return reportable accuracy entries. @@ -301,7 +301,7 @@ def _score_accuracy( else: total_samples = unit_samples * num_repeats if eval_cfg.scorer.SKIP_ENDPOINT_PHASE: - ext = _effective_external_sample_count(eval_cfg) + ext = effective_external_sample_count(eval_cfg) if ext is not None: unit_samples = ext total_samples = ext diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index 4301fc286..7a7f69937 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -40,7 +40,7 @@ from dataclasses import replace as dataclass_replace from datetime import datetime from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from urllib.parse import urljoin import msgspec.json @@ -52,14 +52,14 @@ from inference_endpoint.async_utils.loop_manager import LoopManager from inference_endpoint.commands.benchmark.accuracy import ( AccuracyConfiguration, - _effective_external_sample_count, - _score_accuracy, + effective_external_sample_count, + score_accuracy, write_accuracy_results, ) from inference_endpoint.commands.benchmark.pipeline import MetricsPipeline from inference_endpoint.commands.benchmark.profiling import ( ProfileController, - _write_profiling_section, + write_profiling_section, ) from inference_endpoint.compliance import AuditRunSpec from inference_endpoint.config.runtime_settings import RuntimeSettings @@ -101,6 +101,9 @@ ) from inference_endpoint.metrics.report import Report +if TYPE_CHECKING: + from inference_endpoint.async_utils.event_publisher import EventPublisherService + transformers_logging.set_verbosity_error() logger = logging.getLogger(__name__) @@ -537,7 +540,7 @@ def setup_benchmark( logger.info(f"Accuracy-only mode, Expected samples: {total_samples}") for ec in eval_configs: if ec.scorer.SKIP_ENDPOINT_PHASE: - n = _effective_external_sample_count(ec) + n = effective_external_sample_count(ec) if n is not None: logger.info( "Accuracy dataset '%s' (%s): %d instances evaluated externally", @@ -764,7 +767,7 @@ def _build_agentic_strategy( def _wire_on_sample_complete( collector: ResponseCollector, agentic_inference_strategy: AgenticInferenceStrategy | None, - publisher: Any, + publisher: EventPublisherService, ) -> Callable[[QueryResult], None]: """Compose the per-sample completion callback (agentic strategy + collector).""" if agentic_inference_strategy is None: @@ -1024,7 +1027,7 @@ def _write_report_artifacts( 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) + write_profiling_section(f, profiling) logger.info("Report written to %s", report_txt) @@ -1095,7 +1098,7 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: # then the exception propagates as before. accuracy_scores: list[dict[str, Any]] = [] try: - accuracy_scores = _score_accuracy(ctx, result) + accuracy_scores = score_accuracy(ctx, result) finally: # Attach the per-dataset accuracy list so result_summary.json, the # console summary, and report.txt all carry it (stays [] on a scoring diff --git a/src/inference_endpoint/commands/benchmark/profiling.py b/src/inference_endpoint/commands/benchmark/profiling.py index 3a5f85957..60db572e3 100644 --- a/src/inference_endpoint/commands/benchmark/profiling.py +++ b/src/inference_endpoint/commands/benchmark/profiling.py @@ -27,7 +27,7 @@ import logging import time from datetime import datetime -from typing import Any, TextIO +from typing import Any, Literal, TextIO from urllib import error as urllib_error from urllib import request as urllib_request @@ -42,9 +42,12 @@ ProfilerEngine.VLLM: ("/start_profile", "/stop_profile"), } +# Timeout (seconds) for a single /start_profile or /stop_profile POST. +_PROFILE_POST_TIMEOUT_S = 2 + def _derive_profile_urls( - endpoints: list[str], engine: ProfilerEngine, action: str + endpoints: list[str], engine: ProfilerEngine, action: Literal["start", "stop"] ) -> list[str]: """One profile URL per endpoint, derived from the engine's HTTP protocol. @@ -79,7 +82,7 @@ def _post_profile(url: str) -> dict[str, Any]: } req = urllib_request.Request(url, method="POST", data=b"") try: - with urllib_request.urlopen(req, timeout=2) as resp: + with urllib_request.urlopen(req, timeout=_PROFILE_POST_TIMEOUT_S) as resp: record["status"] = resp.status except urllib_error.HTTPError as e: record["status"] = e.code @@ -106,7 +109,7 @@ def _render_profile_status(rec: dict[str, Any]) -> str: return "ERROR" -def _write_profiling_section(f: TextIO, profiling: dict[str, Any]) -> None: +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", []) @@ -167,10 +170,6 @@ def __init__( 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: diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index 80255710f..959a5f79a 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -52,7 +52,7 @@ _derive_profile_urls, _post_profile, _render_profile_status, - _write_profiling_section, + write_profiling_section, ) from inference_endpoint.config.runtime_settings import RuntimeSettings from inference_endpoint.config.schema import ( @@ -2791,7 +2791,7 @@ def test_write_section_and_json_roundtrip(self): ], } buf = io.StringIO() - _write_profiling_section(buf, payload) + write_profiling_section(buf, payload) text = buf.getvalue() assert "Profiling" in text assert "http://h/start_profile" in text diff --git a/tests/unit/commands/test_score_accuracy.py b/tests/unit/commands/test_score_accuracy.py index ae21b7499..fd5e23284 100644 --- a/tests/unit/commands/test_score_accuracy.py +++ b/tests/unit/commands/test_score_accuracy.py @@ -29,7 +29,7 @@ _accuracy_uuid_bound, _phase_osl_stats, _phase_response_counts, - _score_accuracy, + score_accuracy, ) from inference_endpoint.config import schema as config_schema from inference_endpoint.config.schema import ( @@ -39,11 +39,11 @@ ScorerMethod, ) -# The module that defines _score_accuracy (commands.benchmark.accuracy), resolved +# 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 +# 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__] +scoring_mod = sys.modules[score_accuracy.__module__] class _FakeDataset: @@ -240,7 +240,7 @@ def test_swebench_receives_typed_runtime_model_and_endpoint( endpoint_config=endpoint_config, ) - _score_accuracy(_ctx([cfg], test_mode=test_mode), _RESULT) + score_accuracy(_ctx([cfg], test_mode=test_mode), _RESULT) assert _FakeSWEBenchScorer.received_kwargs == { "extractor": None, @@ -256,7 +256,7 @@ def test_each_dataset_gets_its_own_entry(self, tmp_path): _cfg("gpqa::gptoss", 198, 0.9, tmp_path, repeats=5), _cfg("cnn_dailymail::llama3_8b", 100, 0.5, tmp_path), ] - scores = _score_accuracy(_ctx(cfgs), _RESULT) + scores = score_accuracy(_ctx(cfgs), _RESULT) assert isinstance(scores, list) by = _by_name(scores) assert set(by) == { @@ -277,7 +277,7 @@ def test_breakdown_attached_only_when_scorer_provides_it(self, tmp_path): _cfg("plain", 10, 0.7, tmp_path), _cfg("with_bd", 10, 0.83, tmp_path, scorer=_FakeBreakdownScorer), ] - by = _by_name(_score_accuracy(_ctx(cfgs), _RESULT)) + by = _by_name(score_accuracy(_ctx(cfgs), _RESULT)) assert "breakdown" not in by["plain"] assert by["with_bd"]["breakdown"]["overall_accuracy"] == 80.0 @@ -298,20 +298,20 @@ def test_performance_entry_uses_issued_count_for_total(self, tmp_path): ), ], ) - by = _by_name(_score_accuracy(_ctx([cfg]), result)) + by = _by_name(score_accuracy(_ctx([cfg]), result)) assert by["performance"]["unit_samples"] == 3 assert by["performance"]["num_repeats"] == 1 assert by["performance"]["total_samples"] == 128 assert by["performance"]["duration_s"] == 3.0 def test_empty_when_no_datasets(self, tmp_path): - assert _score_accuracy(_ctx([]), _RESULT) == [] + assert score_accuracy(_ctx([]), _RESULT) == [] def test_perf_mode_skips_only_external_scoring(self, tmp_path): ordinary_cfg = _cfg("aime25::gptoss", 30, 0.8, tmp_path) external_cfg = _cfg("swe_bench", 1, 1.0, tmp_path, scorer=_FakeSWEBenchScorer) _FakeSWEBenchScorer.received_kwargs = {} - scores = _score_accuracy( + scores = score_accuracy( _ctx( [ordinary_cfg, external_cfg], test_mode=config_schema.TestMode.PERF, @@ -327,7 +327,7 @@ def test_no_osl_without_tokenizer(self, tmp_path): # tokenizer_name None (the default) => no output_sequence_lengths attached, # and the OSL path is never entered (fake scorers have no get_raw_outputs). cfg = _cfg("aime25::gptoss", 30, 0.8, tmp_path) - entry = _by_name(_score_accuracy(_ctx([cfg]), _RESULT))["aime25::gptoss"] + entry = _by_name(score_accuracy(_ctx([cfg]), _RESULT))["aime25::gptoss"] assert "output_sequence_lengths" not in entry def test_osl_attached_with_tokenizer(self, tmp_path, monkeypatch): @@ -337,7 +337,7 @@ def test_osl_attached_with_tokenizer(self, tmp_path, monkeypatch): 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))[ + entry = _by_name(score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ "aime25::gptoss" ] # Outputs "a b" (2) and "a b c d" (4); "other" is not in sample_index_map. @@ -356,7 +356,7 @@ def test_osl_skipped_for_performance_entry(self, tmp_path, monkeypatch): 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))[ + entry = _by_name(score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ "performance" ] assert "output_sequence_lengths" not in entry @@ -372,7 +372,7 @@ def get_raw_outputs(self, wanted_uuids=None): raise RuntimeError("events.jsonl unreadable") cfg = _cfg("ds", 1, 0.8, tmp_path, scorer=_RaisingScorer) - entry = _by_name(_score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ + entry = _by_name(score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ "ds" ] assert "output_sequence_lengths" not in entry @@ -383,7 +383,7 @@ def test_response_counts_without_tokenizer(self, tmp_path): """response_counts must be published even with no tokenizer configured — failure visibility cannot depend on OSL being enabled.""" cfg = _cfg("aime25::gptoss", 2, 0.8, tmp_path, scorer=_FakeOSLScorer) - entry = _by_name(_score_accuracy(_ctx([cfg]), _RESULT))["aime25::gptoss"] + entry = _by_name(score_accuracy(_ctx([cfg]), _RESULT))["aime25::gptoss"] assert entry["response_counts"] == { "issued": 2, "scored": 2, @@ -399,7 +399,7 @@ def test_response_counts_published_when_all_empty(self, tmp_path, monkeypatch): 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))[ + entry = _by_name(score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ "aime25::gptoss" ] assert "output_sequence_lengths" not in entry # all empty -> OSL None @@ -417,7 +417,7 @@ def test_response_counts_classifies_missing(self, tmp_path, monkeypatch): 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))[ + entry = _by_name(score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ "ds" ] assert entry["response_counts"] == { @@ -430,7 +430,7 @@ def test_response_counts_classifies_missing(self, tmp_path, monkeypatch): def test_response_counts_skipped_for_performance_entry(self, tmp_path): cfg = _cfg("performance", 2, 0.6, tmp_path, scorer=_FakeOSLScorer) - entry = _by_name(_score_accuracy(_ctx([cfg]), _RESULT))["performance"] + entry = _by_name(score_accuracy(_ctx([cfg]), _RESULT))["performance"] assert "response_counts" not in entry def test_all_missing_still_publishes_response_counts(self, tmp_path, monkeypatch): @@ -440,7 +440,7 @@ def test_all_missing_still_publishes_response_counts(self, tmp_path, monkeypatch 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))[ + entry = _by_name(score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ "ds" ] assert entry["response_counts"] == { @@ -456,7 +456,7 @@ def test_no_fast_backend_disables_osl_keeps_counts(self, tmp_path, monkeypatch): OSL but still publishes response_counts.""" 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))[ + entry = _by_name(score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ "aime25::gptoss" ] assert "output_sequence_lengths" not in entry @@ -472,7 +472,7 @@ def _boom(name): 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))[ + entry = _by_name(score_accuracy(_ctx([cfg], tokenizer_name="fake"), _RESULT))[ "aime25::gptoss" ] assert "output_sequence_lengths" not in entry @@ -480,7 +480,7 @@ def _boom(name): assert entry["score"] == pytest.approx(0.8) def test_uuid_bound_passed_to_get_raw_outputs(self, tmp_path, monkeypatch): - """End-to-end: _score_accuracy computes the accuracy uuid bound from + """End-to-end: score_accuracy computes the accuracy uuid bound from sample_idx_map.json and passes it into get_raw_outputs.""" monkeypatch.setattr( scoring_mod, "load_reference_backend", lambda name: _WordBackend() @@ -496,9 +496,7 @@ def get_raw_outputs(self, wanted_uuids=None): return super().get_raw_outputs(wanted_uuids) cfg = _cfg("aime25::gptoss", 2, 0.8, tmp_path, scorer=_CaptureScorer) - _score_accuracy( - _ctx([cfg], tokenizer_name="fake", report_dir=tmp_path), _RESULT - ) + score_accuracy(_ctx([cfg], tokenizer_name="fake", report_dir=tmp_path), _RESULT) assert captured["wanted"] == {"u1", "u2"} def test_empty_uuid_bound_stays_bounded(self, tmp_path): @@ -513,7 +511,7 @@ def get_raw_outputs(self, wanted_uuids=None): return super().get_raw_outputs(wanted_uuids) cfg = _cfg("swe_bench", 1, 0.8, tmp_path, scorer=_CaptureScorer) - _score_accuracy(_ctx([cfg], report_dir=tmp_path), _RESULT) + score_accuracy(_ctx([cfg], report_dir=tmp_path), _RESULT) assert captured["wanted"] == set() @@ -659,7 +657,7 @@ def test_accuracy_entry_has_phase_duration(self, tmp_path): ), ], ) - entry = _by_name(_score_accuracy(_ctx([cfg]), result))["aime25::gptoss"] + entry = _by_name(score_accuracy(_ctx([cfg]), result))["aime25::gptoss"] assert entry["duration_s"] == 5.5 def test_numpy_score_coerced_to_serializable(self, tmp_path): @@ -677,7 +675,7 @@ def score(self): return np.float64(0.5), 1 cfg = _cfg("np::ds", 10, 0.0, tmp_path, scorer=_NumpyScorer) - entry = _by_name(_score_accuracy(_ctx([cfg]), _RESULT))["np::ds"] + entry = _by_name(score_accuracy(_ctx([cfg]), _RESULT))["np::ds"] # np.floating is a float subclass, so isinstance(..., float) is not # enough — assert it is specifically NOT a numpy scalar. assert not isinstance(entry["score"], np.floating) From 0b227617153060b8340dbacd5e6fc5fc2cf9b9ab Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:31:22 -0500 Subject: [PATCH 05/13] fix(benchmark): make pbar.close() best-effort so it can't mask run errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In run_benchmark_async's teardown finally, an exception raised by pbar.close() (e.g. BrokenPipeError on a closed stderr) would replace the exception being unwound — and on an otherwise-clean run would propagate to the outer except → tmpfs salvage, turning a successful benchmark into a failure that never returns its BenchmarkResult. Catch and log it; the surrounding finally still guarantees pipe.abort()/pipe.close() run. Addresses PR #421 Codex review comment. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/inference_endpoint/commands/benchmark/execute.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index 7a7f69937..cd653b1aa 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -925,12 +925,13 @@ def _on_phase_start(phase: PhaseConfig) -> None: report = await pipe.drain_and_build_report() drained = True 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). + # pbar.close() is cosmetic: a failure here (e.g. BrokenPipeError on a + # closed stderr) must not mask the exception being unwound or turn a + # clean run into a failure. Swallow it; pipe teardown still runs. try: pbar.close() + except Exception as e: # noqa: BLE001 — progress bar is cosmetic + logger.warning("Progress bar close error: %s", e) finally: # A setup failure between pipe.start() and session.run never # reaches the drain above, so the aggregator/event-logger From 3204f52071d6e0e7c7793d8fe46160a55fbb1d9c Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 22:21:53 -0500 Subject: [PATCH 06/13] refactor(benchmark): make MetricsPipeline an async context manager MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the hand-rolled start()/abort()/close()/_teardown() lifecycle with a real async context manager, addressing PR #421 review (r3677374272 CM smell + the _teardown BaseException-leak it hides). pipeline.py: - __aenter__ = start(); __aexit__ releases via a contextlib.ExitStack that owns the ZMQ scope + publisher/subscriber closers, so every step runs even under a BaseException (Ctrl-C/SystemExit) mid-teardown and exceptions chain — no more skipped ZMQ __exit__ or leaked service children. No hand-called __enter__/__exit__. - drain_and_build_report() nulls self.publisher on success; __aexit__ reads that as the drained-cleanly signal to release-only, else kills the services (they'd otherwise linger on the aggregator's unlimited drain-timeout). - Removed abort()/close()/_teardown()/_zmq_cm/_closed; added _close_publisher/ _close_subscriber ExitStack callbacks. execute.py: - Consumer uses `async with pipe:`; drops the `drained` flag, the explicit `except SetupError: pipe.abort()`, and the `if not drained: abort(); close()` finally. Setup/connect/session failures skip the drain, so __aexit__ kills. Closes the hvagadia _teardown/BaseException high, the Gemini subscriber.close-in- drain and publisher-leak items, and the CM-smell design comment. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../commands/benchmark/execute.py | 203 +++++++++--------- .../commands/benchmark/pipeline.py | 114 ++++++---- tests/unit/commands/test_benchmark.py | 21 +- 3 files changed, 175 insertions(+), 163 deletions(-) diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index cd653b1aa..c90869088 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -829,120 +829,115 @@ async def _run_benchmark_async( ) report: Report | None = None profiler: ProfileController - drained = False try: - try: - tmpfs_dir.mkdir(parents=True, exist_ok=True) - event_log_dir.mkdir(parents=True, exist_ok=True) - metrics_output_dir.mkdir(parents=True, exist_ok=True) - # 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 - + tmpfs_dir.mkdir(parents=True, exist_ok=True) + event_log_dir.mkdir(parents=True, exist_ok=True) + metrics_output_dir.mkdir(parents=True, exist_ok=True) + # Pre-derive profile URLs (before the run so a bad config — engine set but + # no endpoints — fails here yet still triggers the tmpfs/pbar cleanup below). + profiler = ProfileController( + config.settings.profiling.engine, + config.endpoint_config.endpoints, + config.settings.profiling.urls, + ) + # MetricsPipeline is an async context manager: __aenter__ brings up ZMQ + + # services; __aexit__ kills the services iff the run never drained (publisher + # still set) else just releases the ZMQ scope. So the setup/connect-failure + # paths need no explicit abort — skipping the drain below leaves the + # publisher set and __aexit__ kills. + async with pipe: try: + # start() guarantees the publisher exists; narrow it for the checker. + publisher = pipe.publisher + assert publisher is not None + 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 - ) + agentic_inference_strategy = _build_agentic_strategy(ctx) + on_sample_complete = _wire_on_sample_complete( + collector, agentic_inference_strategy, publisher + ) - session = BenchmarkSession( - issuer=issuer, - event_publisher=publisher, - loop=loop, - on_sample_complete=on_sample_complete, - session_id=session_id, - ) - phases = _build_phases(ctx, perf_strategy=agentic_inference_strategy) + session = BenchmarkSession( + issuer=issuer, + event_publisher=publisher, + loop=loop, + on_sample_complete=on_sample_complete, + session_id=session_id, + ) + phases = _build_phases(ctx, perf_strategy=agentic_inference_strategy) - max_duration_ms = ( - ctx.rt_settings.max_duration_ms if ctx.rt_settings is not None else None - ) - _timeout_done = False - session_completed_normally = False - - def _on_global_timeout() -> None: - if not _timeout_done: - logger.warning( - "Performance phase max_duration reached (%d ms); " - "ending performance phase.", - max_duration_ms, - ) - # Stop only the perf phase, not the whole session, so a combined - # perf+accuracy run still runs accuracy after the perf cap. - session.stop_current_phase() - - perf_timeout = _PerfPhaseTimeout(loop, max_duration_ms, _on_global_timeout) - - def _on_phase_start(phase: PhaseConfig) -> None: - # _PerfPhaseTimeout arms the perf cap on PERFORMANCE and cancels it - # when any later phase starts, so a combined perf+accuracy run can - # never have its accuracy phase truncated by the perf cap. - perf_timeout.on_phase_start(phase.phase_type) - 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. - profiler.start() - - loop.add_signal_handler(signal.SIGINT, session.stop) - try: - result = await session.run(phases, on_phase_start=_on_phase_start) - session_completed_normally = True - except Exception as e: - raise ExecutionError(f"Benchmark execution failed: {e}") from e - finally: - _timeout_done = True - 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. - profiler.stop(session_completed_normally) - logger.info("Cleaning up...") + max_duration_ms = ( + ctx.rt_settings.max_duration_ms + if ctx.rt_settings is not None + else None + ) + _timeout_done = False + session_completed_normally = False + + def _on_global_timeout() -> None: + if not _timeout_done: + logger.warning( + "Performance phase max_duration reached (%d ms); " + "ending performance phase.", + max_duration_ms, + ) + # Stop only the perf phase, not the whole session, so a + # combined perf+accuracy run still runs accuracy after the + # perf cap. + session.stop_current_phase() + + perf_timeout = _PerfPhaseTimeout( + loop, max_duration_ms, _on_global_timeout + ) + + def _on_phase_start(phase: PhaseConfig) -> None: + # _PerfPhaseTimeout arms the perf cap on PERFORMANCE and cancels + # it when any later phase starts, so a combined perf+accuracy run + # can never have its accuracy phase truncated by the perf cap. + perf_timeout.on_phase_start(phase.phase_type) + 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. + profiler.start() + + loop.add_signal_handler(signal.SIGINT, session.stop) try: - if http_client: - await http_client.shutdown_async() + result = await session.run(phases, on_phase_start=_on_phase_start) + session_completed_normally = True except Exception as e: - logger.warning(f"Client cleanup error: {e}") - # 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() - drained = True - finally: - # pbar.close() is cosmetic: a failure here (e.g. BrokenPipeError on a - # closed stderr) must not mask the exception being unwound or turn a - # clean run into a failure. Swallow it; pipe teardown still runs. - try: - pbar.close() - except Exception as e: # noqa: BLE001 — progress bar is cosmetic - logger.warning("Progress bar close error: %s", e) + raise ExecutionError(f"Benchmark execution failed: {e}") from e + finally: + _timeout_done = True + 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. + 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}") + # 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). Nulls pipe.publisher so __aexit__ releases the ZMQ + # scope without killing the services. + report = await pipe.drain_and_build_report() finally: - # A setup failure between pipe.start() and session.run never - # reaches the drain above, so the aggregator/event-logger - # subprocesses were launched but never signalled to end. Kill them - # (abort) so they can't linger — the aggregator's drain-timeout - # defaults to unlimited. On the clean-finish and session-failure - # paths the drain already waited for their exit, so only the ZMQ - # scope needs closing. - if not drained: - await pipe.abort() - pipe.close() + # pbar.close() is cosmetic: a failure here (e.g. BrokenPipeError on a + # closed stderr) must not mask the exception being unwound or turn a + # clean run into a failure. Swallow it; pipe teardown runs via the + # async-with __aexit__. + try: + pbar.close() + except Exception as e: # noqa: BLE001 — progress bar is cosmetic + logger.warning("Progress bar close error: %s", e) except BaseException: if tmpfs_dir.exists(): _salvage_tmpfs(ctx.report_dir, tmpfs_dir) diff --git a/src/inference_endpoint/commands/benchmark/pipeline.py b/src/inference_endpoint/commands/benchmark/pipeline.py index fee21c660..19cd69eee 100644 --- a/src/inference_endpoint/commands/benchmark/pipeline.py +++ b/src/inference_endpoint/commands/benchmark/pipeline.py @@ -17,29 +17,33 @@ ``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: +It is an async context manager. ``__aenter__`` brings the infrastructure up; +``__aexit__`` releases it — killing the service subprocesses iff the run never +drained (see below), then unwinding publisher, subscriber, and the ZMQ scope via a +``contextlib.ExitStack`` (which runs every release even under ``BaseException`` and +chains them, so a signal mid-teardown can't leak a child or the ZMQ context). + +The one value-producing step stays an explicit method: -* ``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``. + ``ENDED`` in its own ``finally`` and the aggregator writes a terminal snapshot. It + nulls ``self.publisher`` on success — the signal ``__aexit__`` reads to choose + release-only vs. kill: a still-set publisher means the drain never ran (setup / + connect / session error before it), so the services are killed rather than left to + linger on the aggregator's unlimited drain-timeout. """ from __future__ import annotations import asyncio +import contextlib import json import logging import uuid from pathlib import Path +from types import TracebackType from typing import TYPE_CHECKING, Any from inference_endpoint.async_utils.event_publisher import EventPublisherService @@ -212,11 +216,38 @@ def __init__( self._metrics_output_dir = metrics_output_dir self._loop = loop - self._zmq_cm: Any = None + self._stack: contextlib.ExitStack | None = None self._launcher: ServiceLauncher | None = None self.publisher: EventPublisherService | None = None self.subscriber: MetricsSnapshotSubscriber | None = None - self._closed = False + + async def __aenter__(self) -> MetricsPipeline: + await self.start() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """Release the pipeline. Kill the services iff the run never drained. + + ``drain_and_build_report`` nulls ``self.publisher`` on success; a still-set + publisher means we never drained (setup / connect / session error before it), + so the service subprocesses are killed rather than left on the aggregator's + unlimited drain-timeout. The ``ExitStack`` then releases publisher, subscriber + and the ZMQ scope — running every step even under ``BaseException``. + """ + if self._stack is None: + return + stack, self._stack = self._stack, None + if self.publisher is not None and self._launcher is not None: + try: + self._launcher.kill_all() + except Exception as e: # noqa: BLE001 — teardown best-effort + logger.warning("Service kill_all error: %s", e) + stack.close() async def start(self) -> None: """Bring up ZMQ + publisher + subscriber + service subprocesses. @@ -225,14 +256,16 @@ async def start(self) -> None: 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. + early live ticks — or, worst case, missing COMPLETE). An ``ExitStack`` owns + the ZMQ scope + socket closers; on any failure it is unwound (and any + already-spawned children killed) so nothing — notably the ZMQ context — + leaks, and ``self._stack`` is committed only on full success. """ - self._zmq_cm = ManagedZMQContext.scoped(io_threads=2) - zmq_ctx = self._zmq_cm.__enter__() + stack = contextlib.ExitStack() try: + zmq_ctx = stack.enter_context(ManagedZMQContext.scoped(io_threads=2)) self.publisher = EventPublisherService(zmq_ctx) + stack.callback(self._close_publisher) pub_socket_name = self.publisher.socket_name metrics_socket_name = f"metrics_pub_{uuid.uuid4().hex[:8]}" @@ -242,6 +275,7 @@ async def start(self) -> None: metrics_socket_name, zmq_ctx, self._loop ) self.subscriber.start() + stack.callback(self._close_subscriber) self._launcher = ServiceLauncher(zmq_ctx) drain = self._config.settings.drain @@ -269,8 +303,14 @@ async def start(self) -> None: timeout=self._config.settings.service_ready_timeout_s, ) except BaseException: - self._teardown(kill=True) + if self._launcher is not None: # launch may have spawned children + try: + self._launcher.kill_all() + except Exception as e: # noqa: BLE001 — teardown best-effort + logger.warning("Service kill_all error: %s", e) + stack.close() raise + self._stack = stack async def drain_and_build_report(self) -> Report | None: """Graceful drain: close publisher, wait for services, build the Report. @@ -311,46 +351,26 @@ async def drain_and_build_report(self) -> Report | None: 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). + # Null the publisher so __aexit__ sees "drained cleanly" and releases the + # scope without killing the services. The subscriber and ZMQ scope are + # released by the ExitStack in __aexit__. 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 + def _close_publisher(self) -> None: + """Best-effort publisher close (ExitStack callback).""" 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: - try: - self._launcher.kill_all() - except Exception as e: # noqa: BLE001 — teardown best-effort - # Must not propagate: _teardown runs from a finally during - # exception unwind, and the ZMQ __exit__ below must still run - # (a raised kill_all would mask the original error and leak the - # ZMQ scope, since _closed is already set so a later close() no-ops). - logger.warning("Service kill_all error: %s", e) + self.publisher = None + + def _close_subscriber(self) -> None: + """Best-effort subscriber close (ExitStack callback).""" 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/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index 959a5f79a..25519f28b 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -1285,9 +1285,9 @@ async def test_setup_failure_after_launch_kills_services(self, tmp_path): The graceful drain (which waits for the services to self-exit on ENDED) is only reached once session.run runs; a failure before it never sends ENDED, and with drain-timeout defaulting to unlimited the services would - otherwise linger. Here launch succeeds and _create_issuer then raises a - non-SetupError, so the explicit `except SetupError: pipe.abort()` is - bypassed — the finally must still abort (kill=True). + otherwise linger. Here launch succeeds and _create_issuer then raises, so + the run never drains — MetricsPipeline.__aexit__ sees the publisher still + set and kills the services (kill_all). """ config = OfflineConfig(**_OFFLINE_KWARGS, settings=OfflineSettings()) ctx = self._make_ctx(config, tmp_path) @@ -1327,18 +1327,16 @@ async def _launch_ok(service_configs, *, timeout): with pytest.raises(RuntimeError, match="setup boom"): await _run_benchmark_async(ctx, loop) - # abort() → _teardown(kill=True) → launcher.kill_all(); called exactly - # once (the SetupError branch never ran) and never for a clean run. + # __aexit__ kills the services (drain never ran) → launcher.kill_all(); + # called exactly once, and never for a clean run. MockLauncher.return_value.kill_all.assert_called_once() @pytest.mark.unit @pytest.mark.asyncio async def test_setup_error_after_launch_aborts_exactly_once(self, tmp_path): - """A SetupError after launch hits BOTH abort() sites — the explicit - `except SetupError: pipe.abort()` and the finally's `if not drained: - pipe.abort()`. _teardown's `_closed` guard must make the second a no-op, - so kill_all() runs exactly once (not twice). Locks in the idempotency the - finally-abort relies on. + """A SetupError after launch never drains, so MetricsPipeline.__aexit__ + kills the services exactly once — kill_all() runs a single time (there is + one teardown path now, not an explicit abort plus a finally abort). """ config = OfflineConfig(**_OFFLINE_KWARGS, settings=OfflineSettings()) ctx = self._make_ctx(config, tmp_path) @@ -1378,8 +1376,7 @@ async def _launch_ok(service_configs, *, timeout): with pytest.raises(SetupError, match="connect boom"): await _run_benchmark_async(ctx, loop) - # Both abort() sites fire, but _teardown's _closed guard collapses them: - # kill_all runs once, the ZMQ scope is exited once. + # The setup-error path never drains, so __aexit__ kills the services once. MockLauncher.return_value.kill_all.assert_called_once() From 693a816ab1194b4391cca1c9b1f7a21746332743 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 22:33:56 -0500 Subject: [PATCH 07/13] fix(benchmark): guard tmpfs salvage + profiling.json write; cover profiler/pipeline gaps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the remaining PR #421 review items. Correctness (execute.py): - _salvage_tmpfs is now best-effort at both sites (the run_benchmark_async except-BaseException and the run_benchmark finally): a raising salvage no longer masks the exception being unwound or skips the rmtree cleanup. - profiling.json is written after the report artifacts and guarded (OSError → warn), so it can no longer be first-and-unguarded and sink the built perf report on an IO error. Tests: - ProfileController.start/stop: start→stop index mapping, non-200 start skips its stop, abort vs phase_end stop_reason, disabled no-op, stop-without-start. - _build_report_from_snapshot: incomplete-report and legacy-QPS warnings, and the swallow-to-None on a malformed snapshot. - drain_and_build_report snapshot fallback: disk-missing → subscriber.latest, and no-snapshot → None (the SIGKILL/OOM recovery path). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../commands/benchmark/execute.py | 28 ++-- tests/unit/commands/test_benchmark.py | 73 ++++++++++ .../commands/test_benchmark_final_snapshot.py | 131 ++++++++++++++++++ 3 files changed, 223 insertions(+), 9 deletions(-) diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index c90869088..e8b1dc96b 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -940,7 +940,10 @@ def _on_phase_start(phase: PhaseConfig) -> None: logger.warning("Progress bar close error: %s", e) except BaseException: if tmpfs_dir.exists(): - _salvage_tmpfs(ctx.report_dir, tmpfs_dir) + try: + _salvage_tmpfs(ctx.report_dir, tmpfs_dir) + except Exception as e: # noqa: BLE001 — salvage best-effort; keep original exc + logger.warning("Failed to salvage tmpfs: %s", e) shutil.rmtree(tmpfs_dir, ignore_errors=True) raise @@ -1075,13 +1078,6 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: collector = bench.collector report = bench.report - # Sibling profiling.json — kept separate so Report stays a pure - # snapshot-derived struct. - if bench.profiling is not None: - (ctx.report_dir / "profiling.json").write_text( - json.dumps(bench.profiling, indent=2) - ) - # Write scoring artifacts + copy event log from tmpfs to disk (scorers read # sample_idx_map.json + events.jsonl from here). _write_scoring_artifacts(ctx, result, bench.tmpfs_dir) @@ -1107,6 +1103,17 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: _summarize_and_log_metrics(ctx, report, result, collector) + # Sibling profiling.json — kept separate so Report stays a pure snapshot- + # derived struct. Written after the report artifacts (and best-effort) so an + # OSError here can't discard the already-written perf report. + if bench.profiling is not None: + try: + (ctx.report_dir / "profiling.json").write_text( + json.dumps(bench.profiling, indent=2) + ) + except OSError as e: + logger.warning("Failed to write profiling.json: %s", e) + # 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) @@ -1145,7 +1152,10 @@ def run_benchmark( finally: if bench: if bench.tmpfs_dir.exists(): - _salvage_tmpfs(ctx.report_dir, bench.tmpfs_dir) + try: + _salvage_tmpfs(ctx.report_dir, bench.tmpfs_dir) + except Exception as e: # noqa: BLE001 — salvage best-effort + logger.warning("Failed to salvage tmpfs: %s", e) shutil.rmtree(bench.tmpfs_dir, ignore_errors=True) logger.info(f"Partial results saved to {ctx.report_dir}") diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index 25519f28b..a60db9248 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -49,6 +49,7 @@ setup_benchmark, ) from inference_endpoint.commands.benchmark.profiling import ( + ProfileController, _derive_profile_urls, _post_profile, _render_profile_status, @@ -2797,6 +2798,78 @@ def test_write_section_and_json_roundtrip(self): # Mirrors what finalize_benchmark dumps to profiling.json assert json.loads(json.dumps(payload))["engine"] == "vllm" + @pytest.mark.unit + def test_controller_start_then_stop_maps_indices(self): + """start() posts each /start_profile; stop() posts /stop_profile only for + the starts that returned 200, mapped by the same index, tagging stop_reason.""" + + def _fake_post(url): + # b's /start_profile fails (500); everything else succeeds. + status = 500 if url == "http://b/start_profile" else 200 + return { + "url": url, + "status": status, + "error": None, + "sent_at_ns": 1, + "sent_at_iso": "x", + } + + with patch( + "inference_endpoint.commands.benchmark.profiling._post_profile", + side_effect=_fake_post, + ): + ctrl = ProfileController( + ProfilerEngine.VLLM, ["http://a/v1", "http://b/v1"], None + ) + ctrl.start() + ctrl.stop(completed_normally=True) + + payload = ctrl.payload() + assert payload["engine"] == "vllm" + assert [s["url"] for s in payload["starts"]] == [ + "http://a/start_profile", + "http://b/start_profile", + ] + # Only a's start returned 200 → only a's stop fires; b (500) is skipped. + assert [s["url"] for s in payload["stops"]] == ["http://a/stop_profile"] + assert payload["stops"][0]["stop_reason"] == "phase_end" + + @pytest.mark.unit + def test_controller_stop_reason_abort_when_not_completed(self): + with patch( + "inference_endpoint.commands.benchmark.profiling._post_profile", + side_effect=lambda url: { + "url": url, + "status": 200, + "error": None, + "sent_at_ns": 1, + "sent_at_iso": "x", + }, + ): + ctrl = ProfileController(ProfilerEngine.VLLM, ["http://a/v1"], None) + ctrl.start() + ctrl.stop(completed_normally=False) + assert ctrl.payload()["stops"][0]["stop_reason"] == "abort" + + @pytest.mark.unit + def test_controller_disabled_is_noop(self): + """engine=None → no URLs derived, start/stop do nothing, payload is None.""" + ctrl = ProfileController(None, ["http://a/v1"], None) + ctrl.start() + ctrl.stop(completed_normally=True) + assert ctrl.payload() is None + + @pytest.mark.unit + def test_controller_stop_without_start_posts_nothing(self): + """stop() before any start() records nothing (empty _starts, early return).""" + with patch( + "inference_endpoint.commands.benchmark.profiling._post_profile", + ) as mock_post: + ctrl = ProfileController(ProfilerEngine.VLLM, ["http://a/v1"], None) + ctrl.stop(completed_normally=True) + mock_post.assert_not_called() + assert ctrl.payload()["stops"] == [] + class _OverrideTestBase: """Shared helpers for the two end-to-end ``_load_datasets`` override classes diff --git a/tests/unit/commands/test_benchmark_final_snapshot.py b/tests/unit/commands/test_benchmark_final_snapshot.py index 897c7b278..298fb65be 100644 --- a/tests/unit/commands/test_benchmark_final_snapshot.py +++ b/tests/unit/commands/test_benchmark_final_snapshot.py @@ -26,16 +26,23 @@ from __future__ import annotations +import asyncio import json from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock +import inference_endpoint.commands.benchmark.pipeline as pipeline_mod import pytest from inference_endpoint.async_utils.services.metrics_aggregator.snapshot import ( SessionState, ) from inference_endpoint.commands.benchmark.pipeline import ( + MetricsPipeline, + _build_report_from_snapshot, _load_final_snapshot_from_disk, ) +from inference_endpoint.config.schema import LoadPatternType from inference_endpoint.metrics.report import Report @@ -159,3 +166,127 @@ def test_missing_file_path_fallback_yields_no_loaded_snapshot(self, tmp_path: Pa the precondition the fallback chain depends on.""" result = _load_final_snapshot_from_disk(tmp_path / "nope.json") assert result is None + + +def _fake_config(*, poisson: bool = False, use_legacy: bool = False) -> SimpleNamespace: + """Minimal stand-in exposing only what ``_build_report_from_snapshot`` reads.""" + return SimpleNamespace( + settings=SimpleNamespace( + load_pattern=SimpleNamespace( + type=( + LoadPatternType.POISSON + if poisson + else LoadPatternType.MAX_THROUGHPUT + ), + use_legacy_loadgen_qps_metrics=use_legacy, + ), + runtime=SimpleNamespace(scheduler_random_seed=1, dataloader_random_seed=2), + model_dump=lambda **kw: {"load_pattern": {}, "warmup": {}}, + ) + ) + + +@pytest.mark.unit +class TestBuildReportFromSnapshot: + """``_build_report_from_snapshot`` warning branches + swallow-to-None. + + ``Report.from_snapshot`` is stubbed so the branches under test (which key off + the returned Report) are exercised without a fully valid snapshot payload. + """ + + def test_incomplete_report_logs_warning(self, monkeypatch, caplog): + fake_report = SimpleNamespace( + complete=False, state="complete", legacy_loadgen_window_duration_ns=None + ) + monkeypatch.setattr( + pipeline_mod, + "Report", + SimpleNamespace(from_snapshot=lambda *a, **k: fake_report), + ) + with caplog.at_level("WARNING"): + out = _build_report_from_snapshot( + _snapshot_dict(n_pending_tasks=3), _fake_config() + ) + assert out is fake_report + assert any("incomplete" in r.message.lower() for r in caplog.records) + + def test_legacy_loadgen_qps_warning(self, monkeypatch, caplog): + fake_report = SimpleNamespace( + complete=True, + state="complete", + legacy_loadgen_window_duration_ns=5_000_000_000, + ) + monkeypatch.setattr( + pipeline_mod, + "Report", + SimpleNamespace(from_snapshot=lambda *a, **k: fake_report), + ) + with caplog.at_level("WARNING"): + out = _build_report_from_snapshot( + _snapshot_dict(), _fake_config(poisson=True, use_legacy=True) + ) + assert out is fake_report + assert any("legacy" in r.message.lower() for r in caplog.records) + + def test_malformed_snapshot_swallowed_to_none(self, monkeypatch, caplog): + def _boom(*a, **k): + raise ValueError("bad snapshot") + + monkeypatch.setattr( + pipeline_mod, "Report", SimpleNamespace(from_snapshot=_boom) + ) + with caplog.at_level("WARNING"): + out = _build_report_from_snapshot(_snapshot_dict(), _fake_config()) + assert out is None + assert any( + "Failed to build report from snapshot" in r.message for r in caplog.records + ) + + +def _make_pipe(tmp_path: Path) -> MetricsPipeline: + return MetricsPipeline( + MagicMock(), + tokenizer_name=None, + enable_streaming=False, + event_log_dir=tmp_path / "events", + metrics_output_dir=tmp_path / "metrics", + loop=asyncio.get_event_loop(), + ) + + +@pytest.mark.unit +class TestDrainFallback: + """``drain_and_build_report`` snapshot-sourcing fallback — the SIGKILL/OOM + recovery path where ``final_snapshot.json`` never made it to disk.""" + + @pytest.mark.asyncio + async def test_falls_back_to_subscriber_when_no_disk_snapshot( + self, tmp_path, monkeypatch, caplog + ): + pipe = _make_pipe(tmp_path) # metrics dir has no final_snapshot.json + pipe.publisher = MagicMock(buffered_count=0, pending_count=0) + pipe._launcher = MagicMock() + pipe.subscriber = MagicMock(latest=object()) + monkeypatch.setattr(pipeline_mod, "snapshot_to_dict", lambda s: {"k": "v"}) + monkeypatch.setattr( + pipeline_mod, "_build_report_from_snapshot", lambda d, c: "REPORT" + ) + with caplog.at_level("WARNING"): + report = await pipe.drain_and_build_report() + assert report == "REPORT" + assert any( + "No final_snapshot.json on disk" in r.message for r in caplog.records + ) + # Drained cleanly ⇒ publisher nulled (the signal __aexit__ reads). + assert pipe.publisher is None + + @pytest.mark.asyncio + async def test_no_snapshot_available_returns_none(self, tmp_path, caplog): + pipe = _make_pipe(tmp_path) + pipe.publisher = MagicMock(buffered_count=0, pending_count=0) + pipe._launcher = MagicMock() + pipe.subscriber = MagicMock(latest=None) # no disk, no live snapshot + with caplog.at_level("ERROR"): + report = await pipe.drain_and_build_report() + assert report is None + assert any("No metrics snapshot available" in r.message for r in caplog.records) From 257d50bc80e59eef3895f8604fcca90e898a228c Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Wed, 29 Jul 2026 23:54:35 -0500 Subject: [PATCH 08/13] fix(benchmark): shut down HTTP client on setup error; guard drain; cover __aexit__ policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review-council findings on the execute.py/pipeline.py split. - execute.py: bind http_client at function scope and shut it down in the outer finally (which always runs) instead of only in the session finally. A setup error after _create_issuer (e.g. _build_agentic_strategy/_build_phases raising) previously skipped the session finally and leaked the client's worker subprocesses; shutdown_async() is idempotent so the clean-path call is a no-op. - execute.py: make the end-of-run drain best-effort so a teardown error (publisher.close / wait_for_exit) can't replace the run's in-flight exception on the session-failure path (matches the sibling client-shutdown guard). - AGENTS.md: MetricsPipeline is an async context manager now — drop the stale abort/close from its Code Organization entry. Tests: - HTTP client is shut down when setup raises after _create_issuer. - __aexit__ kill policy: clean drain (publisher nulled) does NOT kill the services; never-drained (publisher still set) kills exactly once. Co-Authored-By: Claude Opus 4.8 (1M context) --- AGENTS.md | 2 +- .../commands/benchmark/execute.py | 37 ++++++++---- tests/unit/commands/test_benchmark.py | 58 +++++++++++++++++++ .../commands/test_benchmark_final_snapshot.py | 29 ++++++++++ 4 files changed, 113 insertions(+), 13 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 5f4ddfb8c..281e08fc1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -181,7 +181,7 @@ src/inference_endpoint/ │ │ ├── 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 +│ │ └── pipeline.py # MetricsPipeline: async context manager for the ZMQ + metrics-aggregator/event-logger subprocess lifecycle (__aenter__/__aexit__/start/drain_and_build_report) + 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/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index e8b1dc96b..a3db0218c 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -829,6 +829,11 @@ async def _run_benchmark_async( ) report: Report | None = None profiler: ProfileController + # Bound up front so the outer finally can always shut it down: a setup error + # after _create_issuer (e.g. _build_agentic_strategy/_build_phases raising) + # would otherwise leak the client's worker subprocesses. shutdown_async() is + # idempotent, so the clean-path shutdown below is a harmless second call. + http_client: HTTPEndpointClient | None = None try: tmpfs_dir.mkdir(parents=True, exist_ok=True) @@ -917,27 +922,35 @@ def _on_phase_start(phase: PhaseConfig) -> None: # 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}") # 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). Nulls pipe.publisher so __aexit__ releases the ZMQ - # scope without killing the services. - report = await pipe.drain_and_build_report() + # scope without killing the services. Best-effort: a drain error + # (publisher.close / wait_for_exit) must not replace the run's + # in-flight exception on the session-failure path. + try: + report = await pipe.drain_and_build_report() + except Exception as e: # noqa: BLE001 — drain best-effort; keep run error + logger.warning("Drain/report build error: %s", e) finally: - # pbar.close() is cosmetic: a failure here (e.g. BrokenPipeError on a - # closed stderr) must not mask the exception being unwound or turn a - # clean run into a failure. Swallow it; pipe teardown runs via the - # async-with __aexit__. + # Runs on every path, including a setup error before session.run + # (which never reaches the session finally above). pbar.close() is + # cosmetic — a failure here (e.g. BrokenPipeError on a closed stderr) + # must not mask the exception being unwound; swallow it. The HTTP + # client is shut down here rather than in the session finally so a + # setup error can't leak its worker subprocesses; shutdown_async() + # is idempotent, so the clean-path call after the drain is a no-op. + # pipe teardown runs via the async-with __aexit__. try: pbar.close() except Exception as e: # noqa: BLE001 — progress bar is cosmetic logger.warning("Progress bar close error: %s", e) + if http_client is not None: + try: + await http_client.shutdown_async() + except Exception as e: # noqa: BLE001 — best-effort; idempotent + logger.warning(f"Client cleanup error: {e}") except BaseException: if tmpfs_dir.exists(): try: diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index a60db9248..ccea2d2a3 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -1380,6 +1380,64 @@ async def _launch_ok(service_configs, *, timeout): # The setup-error path never drains, so __aexit__ kills the services once. MockLauncher.return_value.kill_all.assert_called_once() + @pytest.mark.unit + @pytest.mark.asyncio + async def test_http_client_shut_down_on_setup_error_after_create_issuer( + self, tmp_path + ): + """A setup error AFTER _create_issuer (client already created) but before + session.run must still shut the HTTP client down — otherwise its worker + subprocesses leak (the session finally that used to own the shutdown is + never reached on this path).""" + config = OfflineConfig(**_OFFLINE_KWARGS, settings=OfflineSettings()) + ctx = self._make_ctx(config, tmp_path) + + async def _launch_ok(service_configs, *, timeout): + return None + + mock_zmq = MagicMock() + mock_zmq.socket_dir = str(tmp_path / "sockets") + + mock_client = MagicMock() + mock_client.shutdown_async = AsyncMock() + + with ( + patch( + "inference_endpoint.commands.benchmark.pipeline.ManagedZMQContext" + ) as MockZMQ, + patch( + "inference_endpoint.commands.benchmark.pipeline.EventPublisherService" + ) as MockPub, + patch( + "inference_endpoint.commands.benchmark.pipeline.MetricsSnapshotSubscriber" + ) as MockSub, + patch( + "inference_endpoint.commands.benchmark.pipeline.ServiceLauncher" + ) as MockLauncher, + patch("inference_endpoint.commands.benchmark.execute.tqdm"), + patch( + "inference_endpoint.commands.benchmark.execute._create_issuer", + new=AsyncMock(return_value=(MagicMock(), mock_client)), + ), + patch( + "inference_endpoint.commands.benchmark.execute._build_agentic_strategy", + side_effect=RuntimeError("setup boom after client create"), + ), + ): + MockZMQ.scoped.return_value.__enter__ = MagicMock(return_value=mock_zmq) + MockZMQ.scoped.return_value.__exit__ = MagicMock(return_value=False) + MockPub.return_value.socket_name = "test_pub" + MockSub.return_value.start = MagicMock() + MockLauncher.return_value.launch = _launch_ok + + loop = asyncio.get_event_loop() + with pytest.raises(RuntimeError, match="setup boom"): + await _run_benchmark_async(ctx, loop) + + # Client was created; the run failed before the session finally, so the + # outer finally must have shut it down (idempotent, called once here). + mock_client.shutdown_async.assert_awaited_once() + class TestAccuracyOnlyDatasetLoading: """`--accuracy-only` must skip the performance dataset even when the config diff --git a/tests/unit/commands/test_benchmark_final_snapshot.py b/tests/unit/commands/test_benchmark_final_snapshot.py index 298fb65be..0bfaaae2c 100644 --- a/tests/unit/commands/test_benchmark_final_snapshot.py +++ b/tests/unit/commands/test_benchmark_final_snapshot.py @@ -27,6 +27,7 @@ from __future__ import annotations import asyncio +import contextlib import json from pathlib import Path from types import SimpleNamespace @@ -290,3 +291,31 @@ async def test_no_snapshot_available_returns_none(self, tmp_path, caplog): report = await pipe.drain_and_build_report() assert report is None assert any("No metrics snapshot available" in r.message for r in caplog.records) + + +@pytest.mark.unit +class TestAexitKillPolicy: + """``__aexit__`` kills the service subprocesses iff the run never drained. A + clean drain nulls ``pipe.publisher``; killing then would tear down a still- + writing aggregator. The negative arm (clean drain ⇒ no kill) is the core + invariant of the async-context-manager rewrite and must be pinned.""" + + @pytest.mark.asyncio + async def test_clean_drain_does_not_kill(self, tmp_path): + pipe = _make_pipe(tmp_path) + pipe._stack = contextlib.ExitStack() + pipe._launcher = MagicMock() + pipe.publisher = None # drain_and_build_report nulled it on a clean drain + pipe.subscriber = None + await pipe.__aexit__(None, None, None) + pipe._launcher.kill_all.assert_not_called() + + @pytest.mark.asyncio + async def test_never_drained_kills_once(self, tmp_path): + pipe = _make_pipe(tmp_path) + pipe._stack = contextlib.ExitStack() + pipe._launcher = MagicMock() + pipe.publisher = MagicMock() # still set ⇒ setup/session error before drain + pipe.subscriber = None + await pipe.__aexit__(None, None, None) + pipe._launcher.kill_all.assert_called_once() From cca7eee052a20d50168a9c928387ffb2c0dea64c Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:39:50 -0500 Subject: [PATCH 09/13] fix(benchmark): propagate drain error on clean runs; dedup test import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review-council follow-ups on the execute.py split. - execute.py: the end-of-run drain except is now conditional. On a clean run (session_completed_normally) a drain / report-build failure propagates so the run fails loudly — an unconditional swallow would exit 0 with report=None and no perf artifacts. On the session-failure path it still swallows, so a teardown error can't replace the run's in-flight exception. (Refines the earlier best-effort-drain change, which over-corrected.) - test_benchmark.py: a clean session with a failing drain propagates the error. - test_benchmark_final_snapshot.py: drop the duplicate `import ... as pipeline_mod` (kept the `from ... import` style) and patch via string targets — resolves the dual-import lint on the file. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../commands/benchmark/execute.py | 19 +++-- tests/unit/commands/test_benchmark.py | 69 +++++++++++++++++++ .../commands/test_benchmark_final_snapshot.py | 19 +++-- 3 files changed, 92 insertions(+), 15 deletions(-) diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index a3db0218c..a114966a9 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -926,13 +926,22 @@ def _on_phase_start(phase: PhaseConfig) -> None: # failure paths (BenchmarkSession.run publishes ENDED in its own # finally, so a failed run still has a terminal snapshot worth # draining). Nulls pipe.publisher so __aexit__ releases the ZMQ - # scope without killing the services. Best-effort: a drain error - # (publisher.close / wait_for_exit) must not replace the run's - # in-flight exception on the session-failure path. + # scope without killing the services. try: report = await pipe.drain_and_build_report() - except Exception as e: # noqa: BLE001 — drain best-effort; keep run error - logger.warning("Drain/report build error: %s", e) + except Exception as e: # noqa: BLE001 + # On a clean run a drain / report-build failure must be loud: + # silently returning report=None would exit 0 with no perf + # artifacts. On the session-failure path the run is already + # raising, so swallow it there rather than let a teardown + # error replace the in-flight exception. + if session_completed_normally: + raise + logger.warning( + "Drain/report build error suppressed (run already " + "failing): %s", + e, + ) finally: # Runs on every path, including a setup error before session.run # (which never reaches the session finally above). pbar.close() is diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index ccea2d2a3..004638f10 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -1438,6 +1438,75 @@ async def _launch_ok(service_configs, *, timeout): # outer finally must have shut it down (idempotent, called once here). mock_client.shutdown_async.assert_awaited_once() + @pytest.mark.unit + @pytest.mark.asyncio + async def test_clean_run_drain_failure_propagates(self, tmp_path): + """A clean session (session.run succeeds) with a failing drain must NOT + silently exit with report=None — the drain error propagates so the run + fails loudly instead of producing a success with no perf artifacts.""" + config = OfflineConfig(**_OFFLINE_KWARGS, settings=OfflineSettings()) + ctx = self._make_ctx(config, tmp_path) + + async def _launch_ok(service_configs, *, timeout): + return None + + mock_zmq = MagicMock() + mock_zmq.socket_dir = str(tmp_path / "sockets") + mock_client = MagicMock() + mock_client.shutdown_async = AsyncMock() + mock_session = MagicMock() + mock_session.run = AsyncMock(return_value=MagicMock()) # clean success + + loop = asyncio.get_event_loop() + with ( + patch( + "inference_endpoint.commands.benchmark.pipeline.ManagedZMQContext" + ) as MockZMQ, + patch( + "inference_endpoint.commands.benchmark.pipeline.EventPublisherService" + ) as MockPub, + patch( + "inference_endpoint.commands.benchmark.pipeline.MetricsSnapshotSubscriber" + ) as MockSub, + patch( + "inference_endpoint.commands.benchmark.pipeline.ServiceLauncher" + ) as MockLauncher, + patch("inference_endpoint.commands.benchmark.execute.tqdm"), + patch( + "inference_endpoint.commands.benchmark.execute._create_issuer", + new=AsyncMock(return_value=(MagicMock(), mock_client)), + ), + patch( + "inference_endpoint.commands.benchmark.execute._build_agentic_strategy", + return_value=None, + ), + patch( + "inference_endpoint.commands.benchmark.execute.BenchmarkSession", + return_value=mock_session, + ), + patch( + "inference_endpoint.commands.benchmark.execute._build_phases", + return_value=[], + ), + patch( + "inference_endpoint.commands.benchmark.pipeline." + "MetricsPipeline.drain_and_build_report", + new=AsyncMock(side_effect=RuntimeError("drain boom")), + ), + # The perf run reaches the SIGINT handler on the clean path; patch it + # out so the test doesn't touch real process signal state. + patch.object(loop, "add_signal_handler"), + patch.object(loop, "remove_signal_handler"), + ): + MockZMQ.scoped.return_value.__enter__ = MagicMock(return_value=mock_zmq) + MockZMQ.scoped.return_value.__exit__ = MagicMock(return_value=False) + MockPub.return_value.socket_name = "test_pub" + MockSub.return_value.start = MagicMock() + MockLauncher.return_value.launch = _launch_ok + + with pytest.raises(RuntimeError, match="drain boom"): + await _run_benchmark_async(ctx, loop) + class TestAccuracyOnlyDatasetLoading: """`--accuracy-only` must skip the performance dataset even when the config diff --git a/tests/unit/commands/test_benchmark_final_snapshot.py b/tests/unit/commands/test_benchmark_final_snapshot.py index 0bfaaae2c..64d5428c4 100644 --- a/tests/unit/commands/test_benchmark_final_snapshot.py +++ b/tests/unit/commands/test_benchmark_final_snapshot.py @@ -33,7 +33,6 @@ from types import SimpleNamespace from unittest.mock import MagicMock -import inference_endpoint.commands.benchmark.pipeline as pipeline_mod import pytest from inference_endpoint.async_utils.services.metrics_aggregator.snapshot import ( SessionState, @@ -46,6 +45,10 @@ from inference_endpoint.config.schema import LoadPatternType from inference_endpoint.metrics.report import Report +# Dotted path for monkeypatch string targets — keeps a single import style (the +# module is imported once via ``from ... import`` and patched by path here). +_PIPE = "inference_endpoint.commands.benchmark.pipeline" + def _snapshot_dict( *, @@ -200,8 +203,7 @@ def test_incomplete_report_logs_warning(self, monkeypatch, caplog): complete=False, state="complete", legacy_loadgen_window_duration_ns=None ) monkeypatch.setattr( - pipeline_mod, - "Report", + f"{_PIPE}.Report", SimpleNamespace(from_snapshot=lambda *a, **k: fake_report), ) with caplog.at_level("WARNING"): @@ -218,8 +220,7 @@ def test_legacy_loadgen_qps_warning(self, monkeypatch, caplog): legacy_loadgen_window_duration_ns=5_000_000_000, ) monkeypatch.setattr( - pipeline_mod, - "Report", + f"{_PIPE}.Report", SimpleNamespace(from_snapshot=lambda *a, **k: fake_report), ) with caplog.at_level("WARNING"): @@ -233,9 +234,7 @@ def test_malformed_snapshot_swallowed_to_none(self, monkeypatch, caplog): def _boom(*a, **k): raise ValueError("bad snapshot") - monkeypatch.setattr( - pipeline_mod, "Report", SimpleNamespace(from_snapshot=_boom) - ) + monkeypatch.setattr(f"{_PIPE}.Report", SimpleNamespace(from_snapshot=_boom)) with caplog.at_level("WARNING"): out = _build_report_from_snapshot(_snapshot_dict(), _fake_config()) assert out is None @@ -268,9 +267,9 @@ async def test_falls_back_to_subscriber_when_no_disk_snapshot( pipe.publisher = MagicMock(buffered_count=0, pending_count=0) pipe._launcher = MagicMock() pipe.subscriber = MagicMock(latest=object()) - monkeypatch.setattr(pipeline_mod, "snapshot_to_dict", lambda s: {"k": "v"}) + monkeypatch.setattr(f"{_PIPE}.snapshot_to_dict", lambda s: {"k": "v"}) monkeypatch.setattr( - pipeline_mod, "_build_report_from_snapshot", lambda d, c: "REPORT" + f"{_PIPE}._build_report_from_snapshot", lambda d, c: "REPORT" ) with caplog.at_level("WARNING"): report = await pipe.drain_and_build_report() From 9456604d40c76d17fd7c8c2f429a8a4f45de5e35 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Sun, 2 Aug 2026 18:32:12 -0500 Subject: [PATCH 10/13] fix(benchmark): harden report and pipeline teardown failures Fail clean benchmark sessions when metrics draining yields no usable report instead of exiting successfully without performance artifacts. Run service termination through the pipeline ExitStack so publisher, subscriber, and ZMQ cleanup still execute when teardown is interrupted. Add regressions for missing reports and repeated interrupts during both normal and partial-start teardown. --- .../commands/benchmark/execute.py | 4 ++ .../commands/benchmark/pipeline.py | 34 ++++++++------ tests/unit/commands/test_benchmark.py | 35 +++++++++++--- .../commands/test_benchmark_final_snapshot.py | 46 ++++++++++++++++++- 4 files changed, 98 insertions(+), 21 deletions(-) diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index a114966a9..81519c89f 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -929,6 +929,10 @@ def _on_phase_start(phase: PhaseConfig) -> None: # scope without killing the services. try: report = await pipe.drain_and_build_report() + if report is None: + raise ExecutionError( + "Benchmark completed without a usable metrics report" + ) except Exception as e: # noqa: BLE001 # On a clean run a drain / report-build failure must be loud: # silently returning report=None would exit 0 with no perf diff --git a/src/inference_endpoint/commands/benchmark/pipeline.py b/src/inference_endpoint/commands/benchmark/pipeline.py index 19cd69eee..8a7ccb10a 100644 --- a/src/inference_endpoint/commands/benchmark/pipeline.py +++ b/src/inference_endpoint/commands/benchmark/pipeline.py @@ -230,7 +230,7 @@ async def __aexit__( exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None, - ) -> None: + ) -> bool | None: """Release the pipeline. Kill the services iff the run never drained. ``drain_and_build_report`` nulls ``self.publisher`` on success; a still-set @@ -240,14 +240,14 @@ async def __aexit__( and the ZMQ scope — running every step even under ``BaseException``. """ if self._stack is None: - return + return None stack, self._stack = self._stack, None if self.publisher is not None and self._launcher is not None: - try: - self._launcher.kill_all() - except Exception as e: # noqa: BLE001 — teardown best-effort - logger.warning("Service kill_all error: %s", e) - stack.close() + # Register this last so it runs first. ExitStack still executes the + # publisher/subscriber/ZMQ callbacks if kill_all raises BaseException + # (for example, a second Ctrl-C during teardown). + stack.callback(self._kill_services) + return stack.__exit__(exc_type, exc, tb) async def start(self) -> None: """Bring up ZMQ + publisher + subscriber + service subprocesses. @@ -302,13 +302,12 @@ async def start(self) -> None: ], timeout=self._config.settings.service_ready_timeout_s, ) - except BaseException: + except BaseException as e: if self._launcher is not None: # launch may have spawned children - try: - self._launcher.kill_all() - except Exception as e: # noqa: BLE001 — teardown best-effort - logger.warning("Service kill_all error: %s", e) - stack.close() + # Run service termination first, while keeping the remaining + # publisher/subscriber/ZMQ callbacks BaseException-safe. + stack.callback(self._kill_services) + stack.__exit__(type(e), e, e.__traceback__) raise self._stack = stack @@ -357,6 +356,15 @@ async def drain_and_build_report(self) -> Report | None: self.publisher = None return report + def _kill_services(self) -> None: + """Best-effort service termination owned by the pipeline ExitStack.""" + if self._launcher is None: + return + try: + self._launcher.kill_all() + except Exception as e: # noqa: BLE001 — teardown best-effort + logger.warning("Service kill_all error: %s", e) + def _close_publisher(self) -> None: """Best-effort publisher close (ExitStack callback).""" if self.publisher is not None: diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index 004638f10..b35a016c3 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -88,7 +88,11 @@ Scorer, SWEBenchScorer, ) -from inference_endpoint.exceptions import InputValidationError, SetupError +from inference_endpoint.exceptions import ( + ExecutionError, + InputValidationError, + SetupError, +) from inference_endpoint.load_generator.sample_order import create_sample_order from inference_endpoint.load_generator.session import ( PhaseResult, @@ -1440,10 +1444,27 @@ async def _launch_ok(service_configs, *, timeout): @pytest.mark.unit @pytest.mark.asyncio - async def test_clean_run_drain_failure_propagates(self, tmp_path): - """A clean session (session.run succeeds) with a failing drain must NOT - silently exit with report=None — the drain error propagates so the run - fails loudly instead of producing a success with no perf artifacts.""" + @pytest.mark.parametrize( + ("drain_kwargs", "expected_error", "expected_match"), + [ + pytest.param( + {"side_effect": RuntimeError("drain boom")}, + RuntimeError, + "drain boom", + id="drain-raises", + ), + pytest.param( + {"return_value": None}, + ExecutionError, + "without a usable metrics report", + id="drain-returns-no-report", + ), + ], + ) + async def test_clean_run_drain_failure_propagates( + self, tmp_path, drain_kwargs, expected_error, expected_match + ): + """A clean session must fail if draining raises or yields no report.""" config = OfflineConfig(**_OFFLINE_KWARGS, settings=OfflineSettings()) ctx = self._make_ctx(config, tmp_path) @@ -1491,7 +1512,7 @@ async def _launch_ok(service_configs, *, timeout): patch( "inference_endpoint.commands.benchmark.pipeline." "MetricsPipeline.drain_and_build_report", - new=AsyncMock(side_effect=RuntimeError("drain boom")), + new=AsyncMock(**drain_kwargs), ), # The perf run reaches the SIGINT handler on the clean path; patch it # out so the test doesn't touch real process signal state. @@ -1504,7 +1525,7 @@ async def _launch_ok(service_configs, *, timeout): MockSub.return_value.start = MagicMock() MockLauncher.return_value.launch = _launch_ok - with pytest.raises(RuntimeError, match="drain boom"): + with pytest.raises(expected_error, match=expected_match): await _run_benchmark_async(ctx, loop) diff --git a/tests/unit/commands/test_benchmark_final_snapshot.py b/tests/unit/commands/test_benchmark_final_snapshot.py index 64d5428c4..f1f2459f7 100644 --- a/tests/unit/commands/test_benchmark_final_snapshot.py +++ b/tests/unit/commands/test_benchmark_final_snapshot.py @@ -31,7 +31,7 @@ import json from pathlib import Path from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from inference_endpoint.async_utils.services.metrics_aggregator.snapshot import ( @@ -318,3 +318,47 @@ async def test_never_drained_kills_once(self, tmp_path): pipe.subscriber = None await pipe.__aexit__(None, None, None) pipe._launcher.kill_all.assert_called_once() + + @pytest.mark.asyncio + async def test_kill_interrupt_still_closes_remaining_resources(self, tmp_path): + pipe = _make_pipe(tmp_path) + pipe._stack = contextlib.ExitStack() + remaining_cleanup = MagicMock() + pipe._stack.callback(remaining_cleanup) + pipe._launcher = MagicMock() + pipe._launcher.kill_all.side_effect = KeyboardInterrupt + pipe.publisher = MagicMock() + pipe.subscriber = None + + with pytest.raises(KeyboardInterrupt): + await pipe.__aexit__(None, None, None) + + remaining_cleanup.assert_called_once() + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_start_failure_kill_interrupt_still_closes_resources(tmp_path): + pipe = _make_pipe(tmp_path) + mock_zmq = MagicMock(socket_dir=str(tmp_path / "sockets")) + + with ( + patch(f"{_PIPE}.ManagedZMQContext") as mock_context, + patch(f"{_PIPE}.EventPublisherService") as mock_publisher, + patch(f"{_PIPE}.MetricsSnapshotSubscriber") as mock_subscriber, + patch(f"{_PIPE}.ServiceLauncher") as mock_launcher, + ): + mock_context.scoped.return_value.__enter__.return_value = mock_zmq + mock_context.scoped.return_value.__exit__.return_value = False + mock_publisher.return_value.socket_name = "test_pub" + mock_launcher.return_value.launch = AsyncMock( + side_effect=RuntimeError("launch failed") + ) + mock_launcher.return_value.kill_all.side_effect = KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + await pipe.start() + + mock_publisher.return_value.close.assert_called_once() + mock_subscriber.return_value.close.assert_called_once() + mock_context.scoped.return_value.__exit__.assert_called_once() From 42056f0509a2fb6f5e31ce477602d37ace4ca2c2 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:07:37 -0500 Subject: [PATCH 11/13] fix(benchmark): SIGTERM-before-SIGKILL teardown, publisher sentinel race, conditional tmpfs rmtree MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - launcher: add terminate_all() — sends SIGTERM first so the metrics aggregator can flush an INTERRUPTED final_snapshot.json via its signal handler, then escalates to SIGKILL after a 5s timeout per process - pipeline: swap-and-null self.publisher atomically before publisher.close() in drain_and_build_report so __aexit__ sees "drain initiated" even if a CancelledError fires before wait_for_exit returns; _kill_services now calls terminate_all instead of kill_all - execute: move shutil.rmtree inside the _salvage_tmpfs try block at both call sites so tmpfs is only deleted when the salvage copy succeeds; on failure, a warning is logged and tmpfs is retained for manual recovery Co-Authored-By: Claude Sonnet 4.6 (1M context) --- .../async_utils/services/launcher.py | 26 +++++++++++++++++++ .../commands/benchmark/execute.py | 14 +++++++--- .../commands/benchmark/pipeline.py | 25 +++++++++++------- tests/unit/commands/test_benchmark.py | 4 +-- .../commands/test_benchmark_final_snapshot.py | 8 +++--- 5 files changed, 58 insertions(+), 19 deletions(-) diff --git a/src/inference_endpoint/async_utils/services/launcher.py b/src/inference_endpoint/async_utils/services/launcher.py index 9e53c19c4..ccf8af161 100644 --- a/src/inference_endpoint/async_utils/services/launcher.py +++ b/src/inference_endpoint/async_utils/services/launcher.py @@ -145,6 +145,32 @@ async def launch( # re-raise the exception. raise + def terminate_all(self, timeout: float = 5.0) -> None: + """Terminate all managed subprocesses: SIGTERM then escalate to SIGKILL. + + Sends SIGTERM first so subprocesses can flush final state (e.g. the + metrics aggregator writes ``INTERRUPTED`` ``final_snapshot.json`` in its + SIGTERM handler). After ``timeout`` seconds any process that has not + exited is killed unconditionally with SIGKILL. + """ + for proc in self._procs: + if proc.poll() is None: + try: + proc.terminate() + except OSError: + pass + deadline = time.monotonic() + timeout + for proc in self._procs: + remaining = max(0.0, deadline - time.monotonic()) + try: + proc.wait(timeout=remaining) + except subprocess.TimeoutExpired: + try: + proc.kill() + proc.wait(timeout=5) + except OSError: + pass + def kill_all(self) -> None: """Kill all managed subprocesses.""" for proc in self._procs: diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index 81519c89f..e76203c17 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -968,9 +968,11 @@ def _on_phase_start(phase: PhaseConfig) -> None: if tmpfs_dir.exists(): try: _salvage_tmpfs(ctx.report_dir, tmpfs_dir) + shutil.rmtree(tmpfs_dir, ignore_errors=True) except Exception as e: # noqa: BLE001 — salvage best-effort; keep original exc - logger.warning("Failed to salvage tmpfs: %s", e) - shutil.rmtree(tmpfs_dir, ignore_errors=True) + logger.warning( + "Failed to salvage tmpfs: %s — tmpfs retained at %s", e, tmpfs_dir + ) raise return BenchmarkResult( @@ -1180,9 +1182,13 @@ def run_benchmark( if bench.tmpfs_dir.exists(): try: _salvage_tmpfs(ctx.report_dir, bench.tmpfs_dir) + shutil.rmtree(bench.tmpfs_dir, ignore_errors=True) except Exception as e: # noqa: BLE001 — salvage best-effort - logger.warning("Failed to salvage tmpfs: %s", e) - shutil.rmtree(bench.tmpfs_dir, ignore_errors=True) + logger.warning( + "Failed to salvage tmpfs: %s — tmpfs retained at %s", + e, + bench.tmpfs_dir, + ) logger.info(f"Partial results saved to {ctx.report_dir}") return ctx.report_dir diff --git a/src/inference_endpoint/commands/benchmark/pipeline.py b/src/inference_endpoint/commands/benchmark/pipeline.py index 8a7ccb10a..1fd445d15 100644 --- a/src/inference_endpoint/commands/benchmark/pipeline.py +++ b/src/inference_endpoint/commands/benchmark/pipeline.py @@ -244,7 +244,7 @@ async def __aexit__( stack, self._stack = self._stack, None if self.publisher is not None and self._launcher is not None: # Register this last so it runs first. ExitStack still executes the - # publisher/subscriber/ZMQ callbacks if kill_all raises BaseException + # publisher/subscriber/ZMQ callbacks if terminate_all raises BaseException # (for example, a second Ctrl-C during teardown). stack.callback(self._kill_services) return stack.__exit__(exc_type, exc, tb) @@ -327,7 +327,13 @@ async def drain_and_build_report(self) -> Report | None: self.publisher.buffered_count, self.publisher.pending_count, ) - self.publisher.close() + # Null the publisher before closing it so __aexit__ sees "drain initiated" + # regardless of whether the subsequent await completes or is cancelled + # (e.g. second Ctrl-C). If we close first and then CancelledError fires + # before the null assignment, __aexit__ would call terminate_all() on an + # aggregator that is already draining and writing final_snapshot.json. + publisher, self.publisher = self.publisher, None + publisher.close() logger.info("Waiting for services to finish processing...") await asyncio.to_thread(self._launcher.wait_for_exit, None) @@ -350,20 +356,21 @@ async def drain_and_build_report(self) -> Report | None: if snap_dict is not None else None ) - # Null the publisher so __aexit__ sees "drained cleanly" and releases the - # scope without killing the services. The subscriber and ZMQ scope are - # released by the ExitStack in __aexit__. - self.publisher = None return report def _kill_services(self) -> None: - """Best-effort service termination owned by the pipeline ExitStack.""" + """Best-effort service termination owned by the pipeline ExitStack. + + Sends SIGTERM first so the metrics aggregator can flush an INTERRUPTED + final_snapshot.json via its signal handler. Escalates to SIGKILL after + a short timeout for any process that does not exit cleanly. + """ if self._launcher is None: return try: - self._launcher.kill_all() + self._launcher.terminate_all() except Exception as e: # noqa: BLE001 — teardown best-effort - logger.warning("Service kill_all error: %s", e) + logger.warning("Service terminate_all error: %s", e) def _close_publisher(self) -> None: """Best-effort publisher close (ExitStack callback).""" diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index b35a016c3..c9be38cd5 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -1334,7 +1334,7 @@ async def _launch_ok(service_configs, *, timeout): # __aexit__ kills the services (drain never ran) → launcher.kill_all(); # called exactly once, and never for a clean run. - MockLauncher.return_value.kill_all.assert_called_once() + MockLauncher.return_value.terminate_all.assert_called_once() @pytest.mark.unit @pytest.mark.asyncio @@ -1382,7 +1382,7 @@ async def _launch_ok(service_configs, *, timeout): await _run_benchmark_async(ctx, loop) # The setup-error path never drains, so __aexit__ kills the services once. - MockLauncher.return_value.kill_all.assert_called_once() + MockLauncher.return_value.terminate_all.assert_called_once() @pytest.mark.unit @pytest.mark.asyncio diff --git a/tests/unit/commands/test_benchmark_final_snapshot.py b/tests/unit/commands/test_benchmark_final_snapshot.py index f1f2459f7..523369c60 100644 --- a/tests/unit/commands/test_benchmark_final_snapshot.py +++ b/tests/unit/commands/test_benchmark_final_snapshot.py @@ -307,7 +307,7 @@ async def test_clean_drain_does_not_kill(self, tmp_path): pipe.publisher = None # drain_and_build_report nulled it on a clean drain pipe.subscriber = None await pipe.__aexit__(None, None, None) - pipe._launcher.kill_all.assert_not_called() + pipe._launcher.terminate_all.assert_not_called() @pytest.mark.asyncio async def test_never_drained_kills_once(self, tmp_path): @@ -317,7 +317,7 @@ async def test_never_drained_kills_once(self, tmp_path): pipe.publisher = MagicMock() # still set ⇒ setup/session error before drain pipe.subscriber = None await pipe.__aexit__(None, None, None) - pipe._launcher.kill_all.assert_called_once() + pipe._launcher.terminate_all.assert_called_once() @pytest.mark.asyncio async def test_kill_interrupt_still_closes_remaining_resources(self, tmp_path): @@ -326,7 +326,7 @@ async def test_kill_interrupt_still_closes_remaining_resources(self, tmp_path): remaining_cleanup = MagicMock() pipe._stack.callback(remaining_cleanup) pipe._launcher = MagicMock() - pipe._launcher.kill_all.side_effect = KeyboardInterrupt + pipe._launcher.terminate_all.side_effect = KeyboardInterrupt pipe.publisher = MagicMock() pipe.subscriber = None @@ -354,7 +354,7 @@ async def test_start_failure_kill_interrupt_still_closes_resources(tmp_path): mock_launcher.return_value.launch = AsyncMock( side_effect=RuntimeError("launch failed") ) - mock_launcher.return_value.kill_all.side_effect = KeyboardInterrupt + mock_launcher.return_value.terminate_all.side_effect = KeyboardInterrupt with pytest.raises(KeyboardInterrupt): await pipe.start() From fb10049c0783354985b9bcbe38f87e0f4c9e5a18 Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:48:31 -0500 Subject: [PATCH 12/13] fix(launcher): add debug logging to empty OSError handlers in terminate_all Replace bare `except OSError: pass` with `logger.debug(...)` so race-condition process exits during terminate/kill are observable without raising. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- src/inference_endpoint/async_utils/services/launcher.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/inference_endpoint/async_utils/services/launcher.py b/src/inference_endpoint/async_utils/services/launcher.py index ccf8af161..1b9e9df45 100644 --- a/src/inference_endpoint/async_utils/services/launcher.py +++ b/src/inference_endpoint/async_utils/services/launcher.py @@ -157,8 +157,8 @@ def terminate_all(self, timeout: float = 5.0) -> None: if proc.poll() is None: try: proc.terminate() - except OSError: - pass + except OSError as e: # noqa: BLE001 — proc already exited between poll and terminate + logger.debug("terminate pid=%d: %s", proc.pid, e) deadline = time.monotonic() + timeout for proc in self._procs: remaining = max(0.0, deadline - time.monotonic()) @@ -168,8 +168,8 @@ def terminate_all(self, timeout: float = 5.0) -> None: try: proc.kill() proc.wait(timeout=5) - except OSError: - pass + except OSError as e: # noqa: BLE001 — proc already exited between kill and wait + logger.debug("kill pid=%d: %s", proc.pid, e) def kill_all(self) -> None: """Kill all managed subprocesses.""" From be6b9acdf730225416df6fe98b289053fc62177c Mon Sep 17 00:00:00 2001 From: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:58:16 -0500 Subject: [PATCH 13/13] fix(benchmark): update docstrings and harden terminate_all SIGKILL escalation - pipeline: update module and __aexit__ docstrings to reflect that self.publisher is nulled at drain initiation, not on success - launcher: catch subprocess.TimeoutExpired alongside OSError in the SIGKILL escalation block so a D-state stuck process logs a warning and cleanup continues for remaining processes - test_benchmark: update three prose references from kill_all to terminate_all to match the renamed method Co-Authored-By: Claude Sonnet 4.6 (1M context) --- .../async_utils/services/launcher.py | 7 +++++-- .../commands/benchmark/pipeline.py | 14 ++++++++------ tests/unit/commands/test_benchmark.py | 6 +++--- 3 files changed, 16 insertions(+), 11 deletions(-) diff --git a/src/inference_endpoint/async_utils/services/launcher.py b/src/inference_endpoint/async_utils/services/launcher.py index 1b9e9df45..f1e1dac15 100644 --- a/src/inference_endpoint/async_utils/services/launcher.py +++ b/src/inference_endpoint/async_utils/services/launcher.py @@ -168,8 +168,11 @@ def terminate_all(self, timeout: float = 5.0) -> None: try: proc.kill() proc.wait(timeout=5) - except OSError as e: # noqa: BLE001 — proc already exited between kill and wait - logger.debug("kill pid=%d: %s", proc.pid, e) + except (OSError, subprocess.TimeoutExpired) as e: # noqa: BLE001 + # OSError: proc already exited between kill and wait. + # TimeoutExpired: proc is stuck in kernel D-state; log and + # continue so remaining processes in _procs are still handled. + logger.warning("Could not kill pid=%d: %s", proc.pid, e) def kill_all(self) -> None: """Kill all managed subprocesses.""" diff --git a/src/inference_endpoint/commands/benchmark/pipeline.py b/src/inference_endpoint/commands/benchmark/pipeline.py index 1fd445d15..7778ac416 100644 --- a/src/inference_endpoint/commands/benchmark/pipeline.py +++ b/src/inference_endpoint/commands/benchmark/pipeline.py @@ -29,10 +29,11 @@ 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. It - nulls ``self.publisher`` on success — the signal ``__aexit__`` reads to choose - release-only vs. kill: a still-set publisher means the drain never ran (setup / - connect / session error before it), so the services are killed rather than left to - linger on the aggregator's unlimited drain-timeout. + nulls ``self.publisher`` at drain initiation (before ``publisher.close()``) — + the signal ``__aexit__`` reads to choose release-only vs. kill: a still-set + publisher means the drain was never initiated (setup / connect / session error + before it), so the services are killed rather than left to linger on the + aggregator's unlimited drain-timeout. """ from __future__ import annotations @@ -233,8 +234,9 @@ async def __aexit__( ) -> bool | None: """Release the pipeline. Kill the services iff the run never drained. - ``drain_and_build_report`` nulls ``self.publisher`` on success; a still-set - publisher means we never drained (setup / connect / session error before it), + ``drain_and_build_report`` nulls ``self.publisher`` at drain initiation; a + still-set publisher means the drain was never initiated (setup / connect / + session error before it), so the service subprocesses are killed rather than left on the aggregator's unlimited drain-timeout. The ``ExitStack`` then releases publisher, subscriber and the ZMQ scope — running every step even under ``BaseException``. diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index c9be38cd5..48c217856 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/tests/unit/commands/test_benchmark.py @@ -1292,7 +1292,7 @@ async def test_setup_failure_after_launch_kills_services(self, tmp_path): ENDED, and with drain-timeout defaulting to unlimited the services would otherwise linger. Here launch succeeds and _create_issuer then raises, so the run never drains — MetricsPipeline.__aexit__ sees the publisher still - set and kills the services (kill_all). + set and kills the services (terminate_all). """ config = OfflineConfig(**_OFFLINE_KWARGS, settings=OfflineSettings()) ctx = self._make_ctx(config, tmp_path) @@ -1332,7 +1332,7 @@ async def _launch_ok(service_configs, *, timeout): with pytest.raises(RuntimeError, match="setup boom"): await _run_benchmark_async(ctx, loop) - # __aexit__ kills the services (drain never ran) → launcher.kill_all(); + # __aexit__ kills the services (drain never ran) → launcher.terminate_all(); # called exactly once, and never for a clean run. MockLauncher.return_value.terminate_all.assert_called_once() @@ -1340,7 +1340,7 @@ async def _launch_ok(service_configs, *, timeout): @pytest.mark.asyncio async def test_setup_error_after_launch_aborts_exactly_once(self, tmp_path): """A SetupError after launch never drains, so MetricsPipeline.__aexit__ - kills the services exactly once — kill_all() runs a single time (there is + kills the services exactly once — terminate_all() runs a single time (there is one teardown path now, not an explicit abort plus a finally abort). """ config = OfflineConfig(**_OFFLINE_KWARGS, settings=OfflineSettings())