From e9dbe047a2ed65e9b9f4e58cc4bbf25569791292 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 13:14:53 -0700 Subject: [PATCH] fix(harness): handle benchmark CLI validation and output errors --- .../strategies/e2e_benchmark/constants.py | 3 + .../strategies/e2e_benchmark/execution.py | 3 +- .../strategies/e2e_benchmark/models.py | 4 +- .../strategies/e2e_benchmark/provider.py | 3 +- .../strategies/e2e_benchmark/runner.py | 50 ++++++++---- .../e2e_benchmark/test_benchmark.py | 79 +++++++++++++++---- 6 files changed, 105 insertions(+), 37 deletions(-) create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/constants.py diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/constants.py b/tests/rust-python-harness/strategies/e2e_benchmark/constants.py new file mode 100644 index 00000000000..752cbb378d0 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/constants.py @@ -0,0 +1,3 @@ +from typing import Final + +PYTHON_SENTINEL: Final = "litellm-benchmark-python" diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/execution.py b/tests/rust-python-harness/strategies/e2e_benchmark/execution.py index 72e833ac2b7..d24d1107469 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/execution.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/execution.py @@ -12,8 +12,9 @@ from typing import TYPE_CHECKING, Final, TextIO import psutil +from .constants import PYTHON_SENTINEL from .models import Backend, BenchmarkModel, Invocation, Measurement, Memory, Options, Ready, Route, Timing -from .provider import PYTHON_SENTINEL, provider_process +from .provider import provider_process if TYPE_CHECKING: from .workloads import Workload diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/models.py b/tests/rust-python-harness/strategies/e2e_benchmark/models.py index 6f5201a1a47..78875b21fe9 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/models.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/models.py @@ -20,8 +20,8 @@ class Options(BenchmarkModel): repeats: int = Field(default=3, ge=1) profiles: tuple[Profile, ...] = ("small", "request_medium", "request_large", "response_medium", "response_large") routes: tuple[Route, ...] = ("ocr", "aocr") - timeout: float = Field(default=120, gt=0) - sample_interval_ms: float = Field(default=5, ge=1) + timeout: float = Field(default=120, gt=0, allow_inf_nan=False) + sample_interval_ms: float = Field(default=5, ge=1, allow_inf_nan=False) output: str | None = None diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/provider.py b/tests/rust-python-harness/strategies/e2e_benchmark/provider.py index efed21f8302..94eae039ee4 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/provider.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/provider.py @@ -7,10 +7,9 @@ from multiprocessing.connection import Connection from typing import ClassVar, Final from ...shared.parity.local_server import LocalHttpHandler, LocalHttpServer +from .constants import PYTHON_SENTINEL from .models import Backend -PYTHON_SENTINEL: Final = "litellm-benchmark-python" - class Provider(LocalHttpServer): def __init__(self, response: bytes, backend: Backend) -> None: diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/runner.py b/tests/rust-python-harness/strategies/e2e_benchmark/runner.py index feb7de2e980..e72c817368d 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/runner.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/runner.py @@ -1,6 +1,5 @@ from __future__ import annotations -import argparse import platform import subprocess from collections.abc import Sequence @@ -8,6 +7,9 @@ from pathlib import Path from time import monotonic from typing import TYPE_CHECKING, Final +import click +from pydantic import ValidationError + from ...shared.reporting.models import CaseResult, HarnessCase, HarnessRun, ResultArtifact, RunStatus from ...shared.reporting.strategy import ModuleCaseSpec, UpdateCallback from .execution import benchmark @@ -29,19 +31,32 @@ class Report(BenchmarkModel): def parse_options(arguments: Sequence[str]) -> Options: - parser: Final = argparse.ArgumentParser(prog="e2e_benchmark", exit_on_error=False) - parser.add_argument("--iterations", type=int, default=100) - parser.add_argument("--warmup", type=int, default=10) - parser.add_argument("--repeats", type=int, default=3) - parser.add_argument("--profile", dest="profiles", action="append", default=argparse.SUPPRESS) - parser.add_argument("--route", dest="routes", action="append", default=argparse.SUPPRESS) - parser.add_argument("--timeout", type=float, default=120) - parser.add_argument("--sample-interval-ms", type=float, default=5) - parser.add_argument("--output") - parsed, unknown = parser.parse_known_args(arguments) - if unknown: - raise ValueError(f"unknown benchmark arguments: {' '.join(unknown)}") - return Options.model_validate(vars(parsed)) + defaults: Final = Options() + command: Final = click.Command( + "e2e_benchmark", + params=[ + click.Option(("--iterations",), type=click.IntRange(min=1), default=defaults.iterations), + click.Option(("--warmup",), type=click.IntRange(min=1), default=defaults.warmup), + click.Option(("--repeats",), type=click.IntRange(min=1), default=defaults.repeats), + click.Option( + ("--profile", "profiles"), + type=click.Choice(defaults.profiles), + multiple=True, + default=defaults.profiles, + ), + click.Option( + ("--route", "routes"), type=click.Choice(defaults.routes), multiple=True, default=defaults.routes + ), + click.Option(("--timeout",), type=click.FloatRange(min=0, min_open=True), default=defaults.timeout), + click.Option(("--sample-interval-ms",), type=click.FloatRange(min=1), default=defaults.sample_interval_ms), + click.Option(("--output",)), + ], + ) + with command.make_context("e2e_benchmark", list(arguments)) as context: + try: + return Options.model_validate(context.params) + except ValidationError as error: + raise click.UsageError(str(error)) from error def _run_pair( @@ -128,6 +143,11 @@ def run_benchmark_cases( measurements=measurements(tuple(run.results.values())), failures=tuple(run.failures), ) - Path(options.output).write_text(report.model_dump_json(indent=2) + "\n") + try: + Path(options.output).write_text(report.model_dump_json(indent=2) + "\n") + except OSError as error: + raise click.ClickException( + f"cannot write benchmark report to {options.output}: {error.strerror}" + ) from error on_update(run) return int(bool(run.failures)), run diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py b/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py index dd6e01f0d7f..f6449f5d39c 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py @@ -8,20 +8,20 @@ from pathlib import Path from time import monotonic, sleep from typing import Final +import click import httpx import psutil import pytest -from pydantic import ValidationError from litellm.llms.base_llm.ocr.transformation import OCRResponse -from ...cli.catalog import load_catalog -from ...shared.reporting.models import RunStatus +from ...cli import main from .execution import execute_phase, sdk_process, wait_for_output +from .constants import PYTHON_SENTINEL from .models import Invocation, Options -from .provider import PYTHON_SENTINEL, provider_process +from .provider import provider_process from .reporting import percentile, render_measurements -from .runner import Report, parse_options, run_benchmark_cases +from .runner import Report, parse_options from .worker import measure_async, measure_sync from .workloads import JSON_OBJECT, JSON_PAGES, ocr_workload, padded_pdf @@ -71,12 +71,12 @@ def test_pdf_padding_preserves_existing_offsets_and_exact_size() -> None: @pytest.mark.parametrize("arguments", (("--iterations=0",), ("--warmup=0",), ("--route=chat",), ("--profile=unknown",))) def test_invalid_benchmark_options_fail_before_running(arguments: tuple[str, ...]) -> None: - with pytest.raises(ValidationError): + with pytest.raises(click.BadParameter): parse_options(arguments) def test_unknown_options_are_not_silently_ignored() -> None: - with pytest.raises(ValueError, match="unknown benchmark arguments"): + with pytest.raises(click.NoSuchOption, match="No such option"): parse_options(("--concurrency=8",)) @@ -170,18 +170,27 @@ def test_worker_errors_are_reported_instead_of_counted_as_fast_calls() -> None: execute_phase(request, "python", Options(iterations=3, warmup=1), REPO_ROOT) -def test_strategy_runs_both_backends_and_exports_measurements(tmp_path: Path) -> None: - strategy: Final = next(strategy for strategy in load_catalog() if strategy.id == "e2e_benchmark") - case: Final = next(case for case in strategy.cases if case.sdk_function == "ocr") +def test_cli_runs_both_backends_and_exports_measurements(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: output: Final = tmp_path / "measurements.json" - exit_code, run = run_benchmark_cases( - (case,), - REPO_ROOT, - lambda _: None, - ("--profile=small", "--route=aocr", "--iterations=3", "--warmup=1", "--repeats=1", f"--output={output}"), + exit_code: Final = main( + ( + "run", + "e2e_benchmark", + "--surface", + "sdk", + "--function", + "ocr", + "--benchmark-arg=--profile=small", + "--benchmark-arg=--route=aocr", + "--benchmark-arg=--iterations=3", + "--benchmark-arg=--warmup=1", + "--benchmark-arg=--repeats=1", + f"--benchmark-arg=--output={output}", + ) ) - assert exit_code == 0, run.failures - assert run.results[case.key].status is RunStatus.PASSED + captured: Final = capsys.readouterr() + assert exit_code == 0, captured.out + captured.err + assert "Result: PASSED" in captured.out report: Final = Report.model_validate_json(output.read_bytes()) assert {value.backend for value in report.measurements} == {"python", "rust"} assert len({value.ready.response_digest for value in report.measurements}) == 1 @@ -197,3 +206,39 @@ def test_strategy_runs_both_backends_and_exports_measurements(tmp_path: Path) -> assert "aocr/small | python" in table assert "aocr/small | rust" in table assert "CPU ms/call" in table + + +@pytest.mark.parametrize( + "argument", ("--iterations=0", "--warmup=invalid", "--route=chat", "--unknown=1", "--timeout=nan", "--timeout=inf") +) +def test_cli_rejects_invalid_benchmark_options_without_a_traceback( + argument: str, capsys: pytest.CaptureFixture[str] +) -> None: + assert main(("run", "e2e_benchmark", "--function", "ocr", f"--benchmark-arg={argument}")) == 2 + captured: Final = capsys.readouterr() + assert "Error:" in captured.err + assert "Traceback" not in captured.err + assert "sdk/ocr: running" not in captured.out + + +@pytest.mark.parametrize("destination", ("missing/report.json", ".")) +def test_cli_reports_output_errors_without_a_traceback( + tmp_path: Path, destination: str, capsys: pytest.CaptureFixture[str] +) -> None: + output: Final = tmp_path / destination + assert main(("run", "e2e_benchmark", "--function", "chat_completions", f"--benchmark-arg=--output={output}")) == 1 + captured: Final = capsys.readouterr() + assert "Error: cannot write benchmark report" in captured.err + assert str(output) in captured.err + assert "Traceback" not in captured.err + + +def test_cli_reports_unsupported_functions_without_measurements( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + output: Final = tmp_path / "unsupported.json" + assert main(("run", "e2e_benchmark", "--function", "chat_completions", f"--benchmark-arg=--output={output}")) == 0 + captured: Final = capsys.readouterr() + assert "not implemented" in captured.out + report: Final = Report.model_validate_json(output.read_bytes()) + assert report.measurements == ()