From e73d21508dabe3d5072a98f8b6bf041f45458cc9 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 11:51:01 -0700 Subject: [PATCH 01/10] feat(harness): benchmark Python and Rust SDK resource usage --- pyproject.toml | 1 + tests/rust-python-harness/AGENTS.md | 1 + tests/rust-python-harness/cli/test_cli.py | 8 +- .../strategies/e2e_benchmark/AGENTS.md | 1 + .../strategies/e2e_benchmark/README.md | 87 +++++++++ .../strategies/e2e_benchmark/__init__.py | 43 +++++ .../strategies/e2e_benchmark/execution.py | 153 ++++++++++++++++ .../strategies/e2e_benchmark/models.py | 69 ++++++++ .../strategies/e2e_benchmark/provider.py | 71 ++++++++ .../strategies/e2e_benchmark/reporting.py | 78 +++++++++ .../strategies/e2e_benchmark/runner.py | 133 ++++++++++++++ .../e2e_benchmark/test_benchmark.py | 165 ++++++++++++++++++ .../strategies/e2e_benchmark/worker.py | 120 +++++++++++++ .../strategies/e2e_benchmark/workloads.py | 82 +++++++++ uv.lock | 2 + 15 files changed, 1012 insertions(+), 2 deletions(-) create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/README.md create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/__init__.py create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/execution.py create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/models.py create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/provider.py create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/reporting.py create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/runner.py create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/worker.py create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/workloads.py diff --git a/pyproject.toml b/pyproject.toml index b889a3a0e60..48b97bbe4f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -216,6 +216,7 @@ dev = [ "pytest-timeout==2.4.0", "vcrpy==8.2.1", "pytest-recording==0.13.4", + "psutil==7.2.2", ] e2e-dev = [ "playwright==1.61.0", diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md index 017668d4289..700af8e41d2 100644 --- a/tests/rust-python-harness/AGENTS.md +++ b/tests/rust-python-harness/AGENTS.md @@ -63,6 +63,7 @@ tests/rust-python-harness/ - Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr` - `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases - `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses +- `e2e_benchmark/` measures SDK latency, CPU time, and RSS with size-varied local provider replays in isolated processes - `trace_parity/` compares mapped operations, call counts, and required execution ordering; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) - E2E and trace strategies load their registered module cases and run surface-specific execution from their folders - `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest diff --git a/tests/rust-python-harness/cli/test_cli.py b/tests/rust-python-harness/cli/test_cli.py index 226c89843d0..42e1c4af993 100644 --- a/tests/rust-python-harness/cli/test_cli.py +++ b/tests/rust-python-harness/cli/test_cli.py @@ -90,6 +90,7 @@ def test_should_load_surface_aware_and_function_only_strategies() -> None: assert [strategy.id for strategy in strategies] == [ "e2e_parity", + "e2e_benchmark", "trace_parity", "unit_tests_mapping", "unit_tests_parity", @@ -242,6 +243,7 @@ def _assert_unavailable_cell(strategy: Strategy, case: HarnessCase, section_titl def test_every_unavailable_case_finishes_and_explains_itself() -> None: section_titles: Final = { "e2e_parity": "End-to-end parity outcomes", + "e2e_benchmark": "End-to-end benchmark measurements", "trace_parity": "trace comparisons", "unit_tests_mapping": "Python/Rust unit-test mappings", "unit_tests_parity": "Python backend parity outcomes", @@ -262,6 +264,7 @@ def test_every_unavailable_case_finishes_and_explains_itself() -> None: ("strategy_id", "present", "absent"), ( ("e2e_parity", "--surface", "--pytest-arg"), + ("e2e_benchmark", "--benchmark-arg", "--pytest-arg"), ("trace_parity", "--surface", "--pytest-arg"), ("unit_tests_parity", "--pytest-arg", "--surface"), ("unit_tests_mapping", "--detail", "--surface"), @@ -290,6 +293,7 @@ def test_run_help_lists_all_and_every_strategy(capsys: pytest.CaptureFixture[str for command in ( "all", "e2e_parity", + "e2e_benchmark", "trace_parity", "unit_tests_mapping", "unit_tests_parity", @@ -394,9 +398,9 @@ def test_run_all_selects_every_declared_case_once(monkeypatch: pytest.MonkeyPatc monkeypatch.setattr(cli, "run_command", capture_run) assert main(["run", "all", "--function", "ocr"]) == 0 - assert len(selected) == 7 + assert len(selected) == 8 assert sum(case.surface is None for case in selected) == 3 - assert sum(case.surface is not None for case in selected) == 4 + assert sum(case.surface is not None for case in selected) == 5 def test_run_reports_not_implemented_surface_as_not_run( diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md b/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md new file mode 100644 index 00000000000..8aa3a4a38d5 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md @@ -0,0 +1 @@ +Measures Python and Rust SDK latency, CPU time, and process memory against deterministic local provider replays, outside correctness-check overhead diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/README.md b/tests/rust-python-harness/strategies/e2e_benchmark/README.md new file mode 100644 index 00000000000..f1fe6e96c84 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/README.md @@ -0,0 +1,87 @@ +# End-to-end SDK benchmark + +Compare `LITELLM_RUST=0` and `LITELLM_RUST=1` against a separate local HTTP provider process, with no real provider calls, credentials, or Docker required + +The initial workload covers synchronous and asynchronous Mistral OCR using the existing `e2e_parity` recording. Other SDK functions are explicitly unimplemented. This measures SDK calls including loopback HTTP transport and response construction. It does not measure gateway overhead, streaming, or concurrent load + +## Run + +Build the Rust extension in release mode first. An editable `uv sync` normally builds the development profile, which is unsuitable for a Python/Rust performance comparison + +```sh +uv sync --frozen --python 3.12 +VIRTUAL_ENV="$PWD/.venv" uvx --from maturin==1.15.0 maturin develop --release +uv run --no-sync python -m tests.rust-python-harness run e2e_benchmark \ + --surface sdk --function ocr \ + --benchmark-arg=--output=/tmp/e2e-benchmark.json +``` + +Keep `--no-sync` on the benchmark command so it uses the extension you just built. Use an otherwise idle machine and run the same command on both revisions when evaluating a change + +For a short smoke run: + +```sh +uv run --no-sync python -m tests.rust-python-harness run e2e_benchmark \ + --function ocr \ + --benchmark-arg=--profile=small \ + --benchmark-arg=--route=ocr \ + --benchmark-arg=--iterations=10 \ + --benchmark-arg=--warmup=2 \ + --benchmark-arg=--repeats=1 \ + --benchmark-arg=--output=/tmp/e2e-benchmark-smoke.json +``` + +`run all` also runs this strategy with its defaults. No CI integration is added + +## Workloads + +The seed cassette stays under `e2e_parity/sdk/ocr/fixtures/data`. The benchmark derives synthetic size variants in memory; it never edits or re-records the correctness fixtures + +| Profile | Inline PDF bytes | Response pages | +| --- | ---: | ---: | +| small | 32 KiB | 1 | +| request_medium | 256 KiB | 1 | +| request_large | 2 MiB | 1 | +| response_medium | 32 KiB | 16 | +| response_large | 32 KiB | 128 | + +Request variants add PDF comment padding before the EOF marker, preserving existing object offsets. The SDK sends base64 plus JSON framing, so wire request sizes exceed the document sizes above. Response variants repeat recorded pages with contiguous indexes and adjusted usage. They exercise realistic response structure, but their page count intentionally varies independently of the input PDF's content + +## Measurements + +Each backend, route, size, and repeat gets a fresh SDK process for timing and another for memory. Python and Rust execute sequentially, with their order reversed on alternating repeats. The local provider serves preloaded bytes without parsing or capturing request JSON. Its CPU and RSS are outside the SDK measurements + +Workers warm up their clients and run an untimed response check before measuring. Python and Rust response digests must match. Every provider request also checks the existing parity harness's User-Agent convention: a Rust run using Python's HTTP path fails instead of reporting a comparison between two Python runs. Missing native extensions, SDK exceptions, timeouts, and incomplete samples fail the run + +Latency starts immediately before the SDK call and ends when its result has been returned and discarded. Async calls are awaited on a persistent event loop. CPU is process CPU time during the timed batch, including Python and native threads. Fixture loading, process startup, warmup, preflight serialization, and report generation are excluded. Default SDK behavior is retained, so deferred background work can extend beyond a call's return; these metrics describe the measurement window, not the eventual cost of every callback + +The memory controller uses `psutil` to sample only the SDK worker's RSS during a separate run, avoiding polling overhead in latency results. Baseline RSS is taken after warmup and garbage collection. Peak is the highest sampled RSS, including the baseline and final sample. After RSS is measured after the workload and another garbage collection, with input/client state still resident. RSS includes native allocations and shared resident pages, so it is not equivalent to Python heap size or uniquely owned memory. Sampling can miss brief peaks; these values are not an exact allocator high-water mark or proof of a leak + +The terminal reports pooled p50/p95/p99 latency, CPU milliseconds per call, sequential calls per second, baseline/peak/after RSS, and speedup (`Python p50 / backend p50`). Throughput is at concurrency one, not saturation capacity. Short runs cannot estimate tail latency reliably. The JSON retains each repeat's raw latency samples, CPU and memory measurements, input/response sizes, seed hash, Python version, native extension hash, settings, platform, Git revision, and whether the working tree has changes + +## Options + +Pass each option through `--benchmark-arg=...` + +| Option | Default | Meaning | +| --- | --- | --- | +| `--iterations=N` | 100 | Measured calls per worker | +| `--warmup=N` | 10 | Warmup calls, followed by one preflight call | +| `--repeats=N` | 3 | Fresh paired runs per workload | +| `--profile=NAME` | All five | Select a size profile; repeat for several | +| `--route=ocr` or `--route=aocr` | Both | Select SDK entrypoint; repeat for both | +| `--timeout=SECONDS` | 120 | Worker readiness and measurement deadline | +| `--sample-interval-ms=N` | 5 | Memory sampling interval, at least 1 ms | +| `--output=PATH` | None | Write a JSON report, including partial results on worker failure | + +Run the strategy tests and existing harness checks with: + +```sh +uv run --no-sync pytest -o consider_namespace_packages=true \ + tests/rust-python-harness/strategies/e2e_benchmark \ + tests/rust-python-harness/shared tests/rust-python-harness/cli \ + tests/rust-python-harness/strategies/unit_tests_mapping \ + tests/rust-python-harness/strategies/unit_tests_parity \ + tests/rust-python-harness/strategies/unit_tests_rust \ + tests/test_rust_python_harness.py -q +``` diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/__init__.py b/tests/rust-python-harness/strategies/e2e_benchmark/__init__.py new file mode 100644 index 00000000000..692d3d485af --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/__init__.py @@ -0,0 +1,43 @@ +from pathlib import Path +from typing import Final + +from ...shared.reporting.models import SDK_FUNCTIONS, Coverage +from ...shared.reporting.strategy import ( + CaseDefinition, + ModuleCaseSpec, + NotImplementedCaseSpec, + RunnerArgumentDefinition, + StrategyDefinition, +) +from .reporting import render_benchmark_results +from .runner import run_benchmark_cases + +STRATEGY: Final = StrategyDefinition( + id="e2e_benchmark", + order=15, + label="End-to-end benchmark", + description="Compare Python/Rust SDK latency, CPU time, and RSS using local provider replays.", + directory=Path(__file__).parent, + runnable_spec=ModuleCaseSpec, + cases=tuple( + CaseDefinition( + function, + ModuleCaseSpec( + coverage=Coverage.PARTIAL, + module="tests.rust-python-harness.strategies.e2e_benchmark.workloads", + note="Sync/async Mistral OCR with scaled recorded fixtures; concurrency is one.", + ) + if function == "ocr" + else NotImplementedCaseSpec(reason="No benchmark workload is implemented for this SDK function yet."), + surface="sdk", + ) + for function in SDK_FUNCTIONS + ), + run=run_benchmark_cases, + render=render_benchmark_results, + surfaces=("sdk",), + runner_argument=RunnerArgumentDefinition( + option="--benchmark-arg", + help="benchmark option, e.g. --benchmark-arg=--iterations=100; see the strategy README", + ), +) diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/execution.py b/tests/rust-python-harness/strategies/e2e_benchmark/execution.py new file mode 100644 index 00000000000..43fee4e67d3 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/execution.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import os +import subprocess +import sys +import tempfile +from collections.abc import Generator, Iterator +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from time import monotonic, sleep +from typing import TYPE_CHECKING, Final, TextIO, cast + +import psutil + +from .models import PREFIX, Backend, BenchmarkModel, Invocation, Measurement, Memory, Options, Ready, Route, Timing +from .provider import PYTHON_SENTINEL, provider_process + +if TYPE_CHECKING: + from .workloads import Workload + +WORKER_MODULE: Final = "tests.rust-python-harness.strategies.e2e_benchmark.worker" + + +class ProcessMemory(BenchmarkModel): + rss: int + + +def rss_bytes(process: psutil.Process) -> int: + return ProcessMemory.model_validate(process.memory_info(), from_attributes=True).rss + + +def _read_message(stream: TextIO) -> str: + for line in stream: + if line.startswith(PREFIX): + return line.removeprefix(PREFIX) + raise RuntimeError("SDK worker exited without returning a measurement") + + +@contextmanager +def sdk_process(case_file: Path, backend: Backend, repo_root: Path, log: TextIO) -> Generator[subprocess.Popen[str]]: + process: Final = subprocess.Popen( + (sys.executable, "-m", WORKER_MODULE, str(case_file)), + cwd=repo_root, + env={ + **os.environ, + "LITELLM_RUST": "1" if backend == "rust" else "0", + "LITELLM_USER_AGENT": PYTHON_SENTINEL, + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "NO_PROXY": "127.0.0.1,localhost", + "no_proxy": "127.0.0.1,localhost", + "PYTHONPATH": str(repo_root), + }, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=log, + text=True, + ) + try: + yield process + finally: + if process.stdin is not None: + process.stdin.close() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + if process.stdout is not None: + process.stdout.close() + + +def sample_rss(process: psutil.Process, completed: Future[str], interval: float, timeout: float) -> Iterator[int]: + deadline: Final = monotonic() + timeout + while not completed.done(): + if monotonic() >= deadline: + raise TimeoutError("memory measurement timed out") + yield rss_bytes(process) + sleep(interval) + + +def execute_phase( + invocation: Invocation, backend: Backend, options: Options, repo_root: Path +) -> tuple[Ready, Timing, Memory]: + with tempfile.TemporaryDirectory(prefix="litellm-benchmark-") as raw_directory: + directory: Final = Path(raw_directory) + case_file: Final = directory / "invocation.json" + case_file.write_text(invocation.model_dump_json()) + with (directory / "worker.log").open("w+") as log: + try: + with ThreadPoolExecutor(max_workers=1) as reader: + with sdk_process(case_file, backend, repo_root, log) as child: + assert child.stdout is not None and child.stdin is not None + stdout: Final = cast(TextIO, child.stdout) + ready: Final = Ready.model_validate_json( + reader.submit(_read_message, stdout).result(timeout=options.timeout) + ) + process: Final = psutil.Process(child.pid) + baseline: Final = rss_bytes(process) if invocation.phase == "memory" else 0 + child.stdin.write("go\n") + child.stdin.flush() + result: Final = reader.submit(_read_message, stdout) + samples: Final = ( + tuple(sample_rss(process, result, options.sample_interval_ms / 1000, options.timeout)) + if invocation.phase == "memory" + else () + ) + timing: Final = Timing.model_validate_json(result.result(timeout=options.timeout)) + retained: Final = rss_bytes(process) if invocation.phase == "memory" else 0 + memory: Final = Memory( + baseline_rss_bytes=baseline, + sampled_peak_rss_bytes=max((baseline, retained, *samples)), + retained_rss_bytes=retained, + samples=len(samples), + ) + return ready, timing, memory + except (RuntimeError, OSError, ValueError, TimeoutError) as error: + log.seek(0) + raise RuntimeError(f"{backend}/{invocation.phase}: {error}\n{log.read()[-6000:]}") from error + + +def benchmark( + workload: Workload, route: Route, backend: Backend, repeat: int, options: Options, repo_root: Path +) -> Measurement: + with provider_process(workload.response, backend) as url: + invocation: Final = Invocation( + model=workload.model, + document_url=workload.document_url, + route=route, + provider_url=url, + iterations=options.iterations, + warmup=options.warmup, + phase="timing", + ) + ready, timing, _ = execute_phase(invocation, backend, options, repo_root) + memory_ready, _, memory = execute_phase( + invocation.model_copy(update={"phase": "memory"}), backend, options, repo_root + ) + if ready != memory_ready: + raise ValueError("timing and memory workers returned different preflight results") + return Measurement( + backend=backend, + repeat=repeat, + profile=workload.profile, + route=route, + document_bytes=workload.document_bytes, + response_bytes=len(workload.response), + response_pages=workload.response_pages, + fixture_sha256=workload.fixture_sha256, + ready=ready, + timing=timing, + memory=memory, + ) diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/models.py b/tests/rust-python-harness/strategies/e2e_benchmark/models.py new file mode 100644 index 00000000000..12dbd27f320 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/models.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, Field + +Backend = Literal["python", "rust"] +Route = Literal["ocr", "aocr"] +Phase = Literal["timing", "memory"] +Profile = Literal["small", "request_medium", "request_large", "response_medium", "response_large"] +PREFIX: Final = "LITELLM_BENCHMARK " + + +class BenchmarkModel(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class Options(BenchmarkModel): + iterations: int = Field(default=100, ge=1) + warmup: int = Field(default=10, ge=1) + 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) + output: str | None = None + + +class Invocation(BenchmarkModel): + model: str + document_url: str + route: Route + provider_url: str + iterations: int + warmup: int + phase: Phase + + +class Ready(BenchmarkModel): + response_digest: str + python_version: str + native_sha256: str | None + + +class Timing(BenchmarkModel): + latency_ms: tuple[float, ...] + cpu_ms: float + elapsed_ms: float + + +class Memory(BenchmarkModel): + baseline_rss_bytes: int + sampled_peak_rss_bytes: int + retained_rss_bytes: int + samples: int + + +class Measurement(BenchmarkModel): + backend: Backend + repeat: int + profile: Profile + route: Route + document_bytes: int + response_bytes: int + response_pages: int + fixture_sha256: str + ready: Ready + timing: Timing + memory: Memory diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/provider.py b/tests/rust-python-harness/strategies/e2e_benchmark/provider.py new file mode 100644 index 00000000000..f8c0b1607c1 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/provider.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +import multiprocessing +from collections.abc import Generator +from contextlib import contextmanager +from multiprocessing.connection import Connection +from typing import ClassVar, Final + +from ...shared.parity.local_server import LocalHttpHandler, LocalHttpServer +from .models import Backend + +PYTHON_SENTINEL: Final = "litellm-benchmark-python" + + +class Provider(LocalHttpServer): + def __init__(self, response: bytes, backend: Backend) -> None: + super().__init__(("127.0.0.1", 0), Handler) + self.response: Final = response + self.backend: Final = backend + + +class Handler(LocalHttpHandler): + disable_nagle_algorithm: ClassVar[bool] = True + + def do_POST(self) -> None: + provider: Final = self.server + assert isinstance(provider, Provider) + self.rfile.read(int(self.headers.get("content-length", "0"))) + python_http: Final = self.headers.get("user-agent") == PYTHON_SENTINEL + if python_http != (provider.backend == "python"): + self.send_error(409, "SDK backend mismatch: Rust may have fallen back to Python") + return + if self.path != "/v1/ocr": + self.send_error(404, "unexpected benchmark endpoint") + return + self.send_response_only(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(provider.response))) + self.end_headers() + self.wfile.write(provider.response) + + +def _serve(response: bytes, backend: Backend, pipe: Connection) -> None: + with Provider(response, backend) as provider: + pipe.send_bytes(provider.url.encode()) + pipe.close() + provider.serve_forever() + + +@contextmanager +def provider_process(response: bytes, backend: Backend) -> Generator[str]: + context: Final = multiprocessing.get_context("spawn") + receive, send = context.Pipe(duplex=False) + process: Final = context.Process(target=_serve, args=(response, backend, send)) + process.start() + send.close() + try: + if not receive.poll(30): + raise TimeoutError("benchmark provider did not start within 30 seconds") + url: Final = receive.recv_bytes().decode() + if not url.startswith("http://127.0.0.1:"): + raise ValueError("benchmark provider returned an invalid local address") + yield url + finally: + receive.close() + process.terminate() + process.join(timeout=5) + if process.is_alive(): + process.kill() + process.join() + process.close() diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/reporting.py b/tests/rust-python-harness/strategies/e2e_benchmark/reporting.py new file mode 100644 index 00000000000..a996d216032 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/reporting.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import math +import statistics +from collections.abc import Sequence +from typing import Final + +from pydantic import TypeAdapter + +from ...shared.reporting.models import CaseResult +from ...shared.reporting.rendering import ReportSection, render_case_outcome +from .models import Measurement + +MEASUREMENTS: Final = TypeAdapter(tuple[Measurement, ...]) +ARTIFACT_KIND: Final = "e2e_benchmark" + + +def percentile(samples: Sequence[float], quantile: float) -> float: + if not samples or not 0 < quantile <= 1: + raise ValueError("percentile requires samples and a quantile in (0, 1]") + return sorted(samples)[math.ceil(len(samples) * quantile) - 1] + + +def measurements(results: Sequence[CaseResult]) -> tuple[Measurement, ...]: + return tuple( + measurement + for result in results + for artifacts in result.artifacts.values() + for artifact in artifacts + if artifact.kind == ARTIFACT_KIND + for measurement in MEASUREMENTS.validate_json(artifact.body) + ) + + +def render_measurements(values: tuple[Measurement, ...]) -> str: + keys: Final = tuple(dict.fromkeys((value.route, value.profile) for value in values)) + header: Final = ( + "route/profile | backend | p50/p95/p99 ms | CPU ms/call | calls/s | RSS baseline/peak/after MiB | speedup" + ) + + def row(group: tuple[Measurement, ...], baseline: float) -> str: + samples: Final = tuple(sample for value in group for sample in value.timing.latency_ms) + median: Final = statistics.median(samples) + cpu: Final = sum(value.timing.cpu_ms for value in group) / len(samples) + rps: Final = len(samples) * 1000 / sum(value.timing.elapsed_ms for value in group) + rss: Final = ( + statistics.median(value.memory.baseline_rss_bytes for value in group), + max(value.memory.sampled_peak_rss_bytes for value in group), + statistics.median(value.memory.retained_rss_bytes for value in group), + ) + return ( + f"{group[0].route}/{group[0].profile} | {group[0].backend} | " + f"{median:.3f}/{percentile(samples, 0.95):.3f}/{percentile(samples, 0.99):.3f} | {cpu:.3f} | {rps:.1f} | " + f"{'/'.join(f'{value / 2**20:.1f}' for value in rss)} | {baseline / median:.2f}x" + ) + + def rows(route: str, profile: str) -> tuple[str, ...]: + python: Final = tuple( + value for value in values if (value.route, value.profile, value.backend) == (route, profile, "python") + ) + rust: Final = tuple( + value for value in values if (value.route, value.profile, value.backend) == (route, profile, "rust") + ) + baseline: Final = statistics.median(sample for value in python for sample in value.timing.latency_ms) + return row(python, baseline), row(rust, baseline) + + return "\n".join((header, *(line for route, profile in keys for line in rows(route, profile)))) + + +def render_benchmark_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: + values: Final = measurements(results) + blocks: Final = tuple(render_case_outcome(result) for result in results) + return ( + ReportSection( + "End-to-end benchmark measurements", + (*blocks, *((render_measurements(values),) if values else ())), + ), + ) diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/runner.py b/tests/rust-python-harness/strategies/e2e_benchmark/runner.py new file mode 100644 index 00000000000..feb7de2e980 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/runner.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import argparse +import platform +import subprocess +from collections.abc import Sequence +from pathlib import Path +from time import monotonic +from typing import TYPE_CHECKING, Final + +from ...shared.reporting.models import CaseResult, HarnessCase, HarnessRun, ResultArtifact, RunStatus +from ...shared.reporting.strategy import ModuleCaseSpec, UpdateCallback +from .execution import benchmark +from .models import Backend, BenchmarkModel, Measurement, Options, Profile, Route +from .reporting import ARTIFACT_KIND, MEASUREMENTS, measurements + +if TYPE_CHECKING: + from .workloads import Workload + + +class Report(BenchmarkModel): + schema_version: int = 1 + revision: str + working_tree_dirty: bool + platform: str + options: Options + measurements: tuple[Measurement, ...] + failures: tuple[tuple[str, str], ...] + + +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)) + + +def _run_pair( + workload: Workload, route: Route, repeat: int, options: Options, repo_root: Path +) -> tuple[Measurement, ...]: + order: Final[tuple[Backend, Backend]] = ("python", "rust") if repeat % 2 == 0 else ("rust", "python") + pair: Final = tuple(benchmark(workload, route, backend, repeat, options, repo_root) for backend in order) + if pair[0].ready.response_digest != pair[1].ready.response_digest: + raise ValueError("Python and Rust preflight SDK responses differ; run e2e_parity before comparing performance") + if any(len(value.timing.latency_ms) != options.iterations for value in pair): + raise ValueError("SDK worker returned an incomplete measurement") + return pair + + +def _run_job( + result: CaseResult, + run: HarnessRun, + options: Options, + repo_root: Path, + job: tuple[Profile, Route, int, str], +) -> bool: + from .workloads import ocr_workload + + profile, route, repeat, nodeid = job + start: Final = monotonic() + try: + pair: Final = _run_pair(ocr_workload(profile), route, repeat, options, repo_root) + except Exception as error: + result.record(nodeid, RunStatus.ERROR, monotonic() - start) + run.failures.append((nodeid, f"{type(error).__name__}: {error}")) + return False + result.record( + nodeid, + RunStatus.PASSED, + monotonic() - start, + artifacts=(ResultArtifact(ARTIFACT_KIND, MEASUREMENTS.dump_json(pair).decode()),), + ) + return True + + +def _run_case(result: CaseResult, run: HarnessRun, options: Options, repo_root: Path, update: UpdateCallback) -> None: + jobs: Final[tuple[tuple[Profile, Route, int, str], ...]] = tuple( + (profile, route, repeat, f"benchmark:{route}:{profile}:{repeat}") + for profile in dict.fromkeys(options.profiles) + for route in dict.fromkeys(options.routes) + for repeat in range(options.repeats) + ) + result.collected.update(nodeid for _, _, _, nodeid in jobs) + for job in jobs: + result.status = RunStatus.RUNNING + run.current_nodeid = job[3] + update(run) + if not _run_job(result, run, options, repo_root, job): + update(run) + return + update(run) + + +def run_benchmark_cases( + cases: Sequence[HarnessCase], + repo_root: Path, + on_update: UpdateCallback, + runner_args: Sequence[str] = (), +) -> tuple[int, HarnessRun]: + options: Final = parse_options(runner_args) + run: Final = HarnessRun.from_cases(cases) + for result in run.results.values(): + if isinstance(result.case.spec, ModuleCaseSpec): + _run_case(result, run, options, repo_root, on_update) + run.finished_at = monotonic() + if options.output: + revision: Final = subprocess.run( + ("git", "rev-parse", "HEAD"), cwd=repo_root, capture_output=True, text=True, check=True + ).stdout.strip() + report: Final = Report( + revision=revision, + working_tree_dirty=bool( + subprocess.run( + ("git", "status", "--porcelain"), cwd=repo_root, capture_output=True, text=True, check=True + ).stdout.strip() + ), + platform=platform.platform(), + options=options, + measurements=measurements(tuple(run.results.values())), + failures=tuple(run.failures), + ) + Path(options.output).write_text(report.model_dump_json(indent=2) + "\n") + 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 new file mode 100644 index 00000000000..e927847acd9 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import asyncio +import base64 +from concurrent.futures import Future +from pathlib import Path +from time import sleep +from typing import Final + +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 .execution import execute_phase, sample_rss +from .models import Invocation, Options +from .provider import PYTHON_SENTINEL, provider_process +from .reporting import percentile, render_measurements +from .runner import Report, parse_options, run_benchmark_cases +from .worker import measure_async, measure_sync +from .workloads import JSON_OBJECT, JSON_PAGES, ocr_workload, padded_pdf + +REPO_ROOT: Final = Path(__file__).resolve().parents[4] + + +def invocation(*, phase: str = "timing") -> Invocation: + return Invocation.model_validate( + { + "model": "mistral/mistral-ocr-latest", + "document_url": "data:application/pdf;base64,AA==", + "route": "ocr", + "provider_url": "http://127.0.0.1:1", + "iterations": 3, + "warmup": 1, + "phase": phase, + } + ) + + +def test_request_and_response_sizes_vary_independently() -> None: + small: Final = ocr_workload("small") + request: Final = ocr_workload("request_large") + response: Final = ocr_workload("response_large") + small_body: Final = JSON_OBJECT.validate_json(small.response) + request_body: Final = JSON_OBJECT.validate_json(request.response) + response_body: Final = JSON_OBJECT.validate_json(response.response) + + assert small.document_bytes == 32 * 1024 + assert request.document_bytes == 2 * 1024 * 1024 + assert base64.b64decode(request.document_url.split(",", 1)[1]).startswith(b"%PDF-") + assert small_body["pages"] == request_body["pages"] + assert response.document_url == small.document_url + assert len(response.response) > 100 * len(small.response) + assert tuple(page["index"] for page in JSON_PAGES.validate_python(response_body["pages"])) == tuple(range(128)) + assert JSON_OBJECT.validate_python(response_body["usage_info"])["pages_processed"] == 128 + assert small.fixture_sha256 == request.fixture_sha256 == response.fixture_sha256 + + +def test_pdf_padding_preserves_existing_offsets_and_exact_size() -> None: + seed: Final = b"%PDF-1.7\n1 0 obj\n<<>>\nendobj\nstartxref\n9\n%%EOF\n" + padded: Final = padded_pdf(seed, 1024) + assert len(padded) == 1024 + assert padded.startswith(seed.split(b"%%EOF")[0]) + assert padded.endswith(b"\n%%EOF\n") + + +@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): + parse_options(arguments) + + +def test_unknown_options_are_not_silently_ignored() -> None: + with pytest.raises(ValueError, match="unknown benchmark arguments"): + parse_options(("--concurrency=8",)) + + +def test_percentiles_use_nearest_rank_without_dropping_the_tail() -> None: + assert percentile(tuple(range(1, 101)), 0.95) == 95 + assert percentile((4, 1, 3, 2), 0.99) == 4 + with pytest.raises(ValueError, match="requires samples"): + percentile((), 0.95) + + +def test_sync_timing_excludes_waiting_from_cpu_time() -> None: + def call() -> OCRResponse: + sleep(0.02) + return OCRResponse(model="benchmark", pages=[]) + + result: Final = measure_sync(call, invocation()) + assert len(result.latency_ms) == 3 + assert min(result.latency_ms) >= 20 + assert result.elapsed_ms >= sum(result.latency_ms) + assert result.cpu_ms < result.elapsed_ms / 2 + + +def test_async_timing_awaits_the_sdk_operation() -> None: + async def call() -> OCRResponse: + await asyncio.sleep(0.02) + return OCRResponse(model="benchmark", pages=[]) + + result: Final = asyncio.run(measure_async(call, invocation())) + assert len(result.latency_ms) == 3 + assert min(result.latency_ms) >= 20 + assert result.cpu_ms < result.elapsed_ms / 2 + + +def test_memory_pass_does_not_accumulate_latency_samples() -> None: + result: Final = measure_sync(lambda: OCRResponse(model="benchmark", pages=[]), invocation(phase="memory")) + assert result.latency_ms == () + + +def test_memory_monitor_has_a_deadline() -> None: + pending: Final[Future[str]] = Future() + with pytest.raises(TimeoutError, match="memory measurement timed out"): + tuple(sample_rss(psutil.Process(), pending, interval=0.001, timeout=0.01)) + + +def test_replay_rejects_python_fallback_during_rust_measurement() -> None: + workload: Final = ocr_workload("small") + with provider_process(workload.response, "rust") as url: + response: Final = httpx.post(url + "/v1/ocr", content=b"{}", headers={"user-agent": PYTHON_SENTINEL}) + assert response.status_code == 409 + assert "backend mismatch" in response.text + + +def test_worker_errors_are_reported_instead_of_counted_as_fast_calls() -> None: + workload: Final = ocr_workload("small") + with provider_process(workload.response, "rust") as url: + request: Final = invocation().model_copy(update={"provider_url": url, "document_url": workload.document_url}) + with pytest.raises(RuntimeError, match="backend mismatch"): + 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") + 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}"), + ) + assert exit_code == 0, run.failures + assert run.results[case.key].status is RunStatus.PASSED + 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 + for value in report.measurements: + assert len(value.timing.latency_ms) == 3 + assert value.timing.cpu_ms > 0 + assert min(value.timing.latency_ms) > 0 + assert value.memory.baseline_rss_bytes > 0 + assert value.memory.sampled_peak_rss_bytes >= value.memory.baseline_rss_bytes + assert value.memory.sampled_peak_rss_bytes >= value.memory.retained_rss_bytes > 0 + assert (value.ready.native_sha256 is not None) == (value.backend == "rust") + table: Final = render_measurements(report.measurements) + assert "aocr/small | python" in table + assert "aocr/small | rust" in table + assert "CPU ms/call" in table diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/worker.py b/tests/rust-python-harness/strategies/e2e_benchmark/worker.py new file mode 100644 index 00000000000..2de690239a9 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/worker.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import asyncio +import gc +import hashlib +import json +import platform +import sys +from collections.abc import Awaitable, Callable +from pathlib import Path +from time import perf_counter_ns, process_time_ns +from typing import Final, cast + +from litellm.llms.base_llm.ocr.transformation import OCRResponse + +from .models import PREFIX, Invocation, Ready, Timing + + +def _ready(response: OCRResponse) -> Ready: + from litellm.rust_bridge import get_native_bridge + from litellm.rust_bridge.configuration import rust_enabled + + bridge: Final = get_native_bridge() if rust_enabled() else None + if rust_enabled() and bridge is None: + raise RuntimeError("native Rust bridge is unavailable; build it with maturin develop --release") + native_path: Final = bridge.__file__ if bridge is not None else None + return Ready( + response_digest=hashlib.sha256(json.dumps(response.model_dump(), sort_keys=True).encode()).hexdigest(), + python_version=platform.python_version(), + native_sha256=hashlib.sha256(Path(native_path).read_bytes()).hexdigest() if native_path else None, + ) + + +def _emit(value: Ready | Timing) -> None: + print(PREFIX + value.model_dump_json(), flush=True) + + +def _handshake(ready: Ready) -> None: + gc.collect() + _emit(ready) + if sys.stdin.readline().strip() != "go": + raise RuntimeError("benchmark controller disconnected before measurement") + + +def _finish(timing: Timing) -> None: + gc.collect() + _emit(timing) + sys.stdin.readline() + + +def _sync_sample(call: Callable[[], OCRResponse]) -> float: + start: Final = perf_counter_ns() + call() + return (perf_counter_ns() - start) / 1e6 + + +async def _async_sample(call: Callable[[], Awaitable[OCRResponse]]) -> float: + start: Final = perf_counter_ns() + await call() + return (perf_counter_ns() - start) / 1e6 + + +def measure_sync(call: Callable[[], OCRResponse], invocation: Invocation) -> Timing: + cpu_start: Final = process_time_ns() + wall_start: Final = perf_counter_ns() + if invocation.phase == "memory": + for _ in range(invocation.iterations): + call() + return Timing(latency_ms=(), cpu_ms=0, elapsed_ms=0) + samples: Final = tuple(_sync_sample(call) for _ in range(invocation.iterations)) + elapsed: Final = perf_counter_ns() - wall_start + return Timing(latency_ms=samples, cpu_ms=(process_time_ns() - cpu_start) / 1e6, elapsed_ms=elapsed / 1e6) + + +async def measure_async(call: Callable[[], Awaitable[OCRResponse]], invocation: Invocation) -> Timing: + cpu_start: Final = process_time_ns() + wall_start: Final = perf_counter_ns() + if invocation.phase == "memory": + for _ in range(invocation.iterations): + await call() + return Timing(latency_ms=(), cpu_ms=0, elapsed_ms=0) + samples: Final = tuple([await _async_sample(call) for _ in range(invocation.iterations)]) + elapsed: Final = perf_counter_ns() - wall_start + return Timing(latency_ms=samples, cpu_ms=(process_time_ns() - cpu_start) / 1e6, elapsed_ms=elapsed / 1e6) + + +async def _run_async(call: Callable[[], Awaitable[OCRResponse]], invocation: Invocation) -> None: + for _ in range(invocation.warmup): + await call() + ready: Final = _ready(await call()) + _handshake(ready) + _finish(await measure_async(call, invocation)) + + +def run_worker(invocation: Invocation) -> None: + import litellm + + kwargs: Final = { + "model": invocation.model, + "document": {"type": "document_url", "document_url": invocation.document_url}, + "api_key": "benchmark-local-only", + "api_base": invocation.provider_url, + "timeout": 10, + "num_retries": 0, + } + if invocation.route == "aocr": + async_route: Final = cast(Callable[..., Awaitable[OCRResponse]], litellm.aocr) + asyncio.run(_run_async(lambda: async_route(**kwargs), invocation)) + return + sync_route: Final = cast(Callable[..., OCRResponse], litellm.ocr) + call: Final = lambda: sync_route(**kwargs) + for _ in range(invocation.warmup): + call() + ready: Final = _ready(call()) + _handshake(ready) + _finish(measure_sync(call, invocation)) + + +if __name__ == "__main__": + run_worker(Invocation.model_validate_json(Path(sys.argv[1]).read_bytes())) diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/workloads.py b/tests/rust-python-harness/strategies/e2e_benchmark/workloads.py new file mode 100644 index 00000000000..d78e8699791 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/workloads.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import base64 +import hashlib +import json +from dataclasses import dataclass +from typing import Final + +from pydantic import JsonValue, TypeAdapter + +from ...shared.parity.fixtures.store import read_fixture +from ...shared.parity.recorded_http import RecordedHttpResponse +from ..e2e_parity.sdk.ocr.fixtures.config import DEFAULT_FIXTURE_DIRECTORY +from ..e2e_parity.sdk.ocr.fixtures.models import OcrParityCase +from .models import Profile + +SEED: Final = ( + DEFAULT_FIXTURE_DIRECTORY / "mistral-ocr/7727f65058eebe0c68c2a9be97c4777f9a19a7c5a860f5953037e19690bc1154.yaml" +) +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +JSON_PAGES: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]) + + +@dataclass(frozen=True, slots=True) +class Workload: + profile: Profile + model: str + document_url: str + document_bytes: int + response: bytes + response_pages: int + fixture_sha256: str + + +def profile_sizes(profile: Profile) -> tuple[int, int]: + match profile: + case "small": + return 32 * 1024, 1 + case "request_medium": + return 256 * 1024, 1 + case "request_large": + return 2 * 1024 * 1024, 1 + case "response_medium": + return 32 * 1024, 16 + case "response_large": + return 32 * 1024, 128 + + +def padded_pdf(document: bytes, size: int) -> bytes: + prefix, marker, suffix = document.rpartition(b"%%EOF") + if not marker or not document.startswith(b"%PDF-") or size < len(document) + 3: + raise ValueError("expected a PDF seed smaller than the requested document size") + return prefix + b"%" + b"x" * (size - len(document) - 2) + b"\n" + marker + suffix + + +def ocr_workload(profile: Profile) -> Workload: + seed: Final = read_fixture(SEED, OcrParityCase) + document: Final = seed.litellm_input.document + response: Final = seed.provider_responses[0] + if document.type != "document_url" or not isinstance(response, RecordedHttpResponse): + raise ValueError("OCR benchmark seed must contain an inline PDF and a non-streaming response") + if not document.document_url.startswith("data:application/pdf;base64,") or response.status_code != 200: + raise ValueError("OCR benchmark seed must be a successful inline PDF recording") + document_size, page_count = profile_sizes(profile) + pdf: Final = padded_pdf(base64.b64decode(document.document_url.split(",", 1)[1], validate=True), document_size) + body: Final = JSON_OBJECT.validate_json(response.body_bytes()) + pages: Final = JSON_PAGES.validate_python(body["pages"]) + usage: Final = JSON_OBJECT.validate_python(body["usage_info"]) + scaled: Final = { + **body, + "pages": tuple({**pages[index % len(pages)], "index": index} for index in range(page_count)), + "usage_info": {**usage, "pages_processed": page_count, "doc_size_bytes": len(pdf)}, + } + return Workload( + profile=profile, + model=seed.litellm_input.model, + document_url="data:application/pdf;base64," + base64.b64encode(pdf).decode("ascii"), + document_bytes=len(pdf), + response=json.dumps(scaled, separators=(",", ":"), ensure_ascii=False).encode(), + response_pages=page_count, + fixture_sha256=hashlib.sha256(SEED.read_bytes()).hexdigest(), + ) diff --git a/uv.lock b/uv.lock index 89205cd9527..06a3c0ebcd1 100644 --- a/uv.lock +++ b/uv.lock @@ -4534,6 +4534,7 @@ dev = [ { name = "opentelemetry-instrumentation-fastapi" }, { name = "opentelemetry-sdk" }, { name = "parameterized" }, + { name = "psutil" }, { name = "psycopg" }, { name = "psycopg-binary" }, { name = "pytest" }, @@ -4723,6 +4724,7 @@ dev = [ { name = "opentelemetry-instrumentation-fastapi", specifier = "==0.49b0" }, { name = "opentelemetry-sdk", specifier = "==1.28.0" }, { name = "parameterized", specifier = "==0.9.0" }, + { name = "psutil", specifier = "==7.2.2" }, { name = "psycopg", specifier = "==3.3.3" }, { name = "psycopg-binary", specifier = "==3.3.3" }, { name = "pytest", specifier = "==9.0.3" }, From af1c8eb5ed32e87e17c04702b5e2d1b6b2351621 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 11:55:29 -0700 Subject: [PATCH 02/10] fix(harness): preserve PDF trailers in benchmark fixtures --- .../strategies/e2e_benchmark/README.md | 2 +- .../strategies/e2e_benchmark/test_benchmark.py | 4 ++-- .../strategies/e2e_benchmark/workloads.py | 9 +++++++-- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/README.md b/tests/rust-python-harness/strategies/e2e_benchmark/README.md index f1fe6e96c84..58b17a65579 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/README.md +++ b/tests/rust-python-harness/strategies/e2e_benchmark/README.md @@ -45,7 +45,7 @@ The seed cassette stays under `e2e_parity/sdk/ocr/fixtures/data`. The benchmark | response_medium | 32 KiB | 16 | | response_large | 32 KiB | 128 | -Request variants add PDF comment padding before the EOF marker, preserving existing object offsets. The SDK sends base64 plus JSON framing, so wire request sizes exceed the document sizes above. Response variants repeat recorded pages with contiguous indexes and adjusted usage. They exercise realistic response structure, but their page count intentionally varies independently of the input PDF's content +Request variants add PDF comment padding before the final `startxref` marker, preserving existing object offsets and the EOF trailer. The SDK sends base64 plus JSON framing, so wire request sizes exceed the document sizes above. Response variants repeat recorded pages with contiguous indexes and adjusted usage. They exercise realistic response structure, but their page count intentionally varies independently of the input PDF's content ## Measurements 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 e927847acd9..367cea1c410 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py @@ -64,8 +64,8 @@ def test_pdf_padding_preserves_existing_offsets_and_exact_size() -> None: seed: Final = b"%PDF-1.7\n1 0 obj\n<<>>\nendobj\nstartxref\n9\n%%EOF\n" padded: Final = padded_pdf(seed, 1024) assert len(padded) == 1024 - assert padded.startswith(seed.split(b"%%EOF")[0]) - assert padded.endswith(b"\n%%EOF\n") + assert padded.startswith(seed.split(b"startxref")[0]) + assert padded.endswith(b"\nstartxref\n9\n%%EOF\n") @pytest.mark.parametrize("arguments", (("--iterations=0",), ("--warmup=0",), ("--route=chat",), ("--profile=unknown",))) diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/workloads.py b/tests/rust-python-harness/strategies/e2e_benchmark/workloads.py index d78e8699791..8b984b1ce85 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/workloads.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/workloads.py @@ -47,8 +47,13 @@ def profile_sizes(profile: Profile) -> tuple[int, int]: def padded_pdf(document: bytes, size: int) -> bytes: - prefix, marker, suffix = document.rpartition(b"%%EOF") - if not marker or not document.startswith(b"%PDF-") or size < len(document) + 3: + prefix, marker, suffix = document.rpartition(b"startxref") + if ( + not marker + or not suffix.rstrip().endswith(b"%%EOF") + or not document.startswith(b"%PDF-") + or size < len(document) + 3 + ): raise ValueError("expected a PDF seed smaller than the requested document size") return prefix + b"%" + b"x" * (size - len(document) - 2) + b"\n" + marker + suffix From de28a65e960d5595cbc01568cea1273e0ca52c50 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 13:07:01 -0700 Subject: [PATCH 03/10] refactor(harness): simplify benchmark worker coordination --- .../strategies/e2e_benchmark/README.md | 2 +- .../strategies/e2e_benchmark/execution.py | 92 ++++++++++--------- .../strategies/e2e_benchmark/models.py | 3 +- .../strategies/e2e_benchmark/provider.py | 2 +- .../e2e_benchmark/test_benchmark.py | 48 ++++++++-- .../strategies/e2e_benchmark/worker.py | 33 ++++--- 6 files changed, 111 insertions(+), 69 deletions(-) diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/README.md b/tests/rust-python-harness/strategies/e2e_benchmark/README.md index 58b17a65579..22b103e37a3 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/README.md +++ b/tests/rust-python-harness/strategies/e2e_benchmark/README.md @@ -51,7 +51,7 @@ Request variants add PDF comment padding before the final `startxref` marker, pr Each backend, route, size, and repeat gets a fresh SDK process for timing and another for memory. Python and Rust execute sequentially, with their order reversed on alternating repeats. The local provider serves preloaded bytes without parsing or capturing request JSON. Its CPU and RSS are outside the SDK measurements -Workers warm up their clients and run an untimed response check before measuring. Python and Rust response digests must match. Every provider request also checks the existing parity harness's User-Agent convention: a Rust run using Python's HTTP path fails instead of reporting a comparison between two Python runs. Missing native extensions, SDK exceptions, timeouts, and incomplete samples fail the run +Workers warm up their clients and run an untimed response check before measuring. Python and Rust response digests must match. Every provider request also checks the existing parity harness's User-Agent convention: a Rust run using Python's HTTP path fails instead of reporting a comparison between two Python runs. Missing native extensions, SDK exceptions, timeouts, and incomplete samples fail the run. Workers publish readiness and results atomically in temporary JSON files; stdout and stderr go to a diagnostic log. The controller waits for timing workers to exit and samples RSS only for memory workers. A timeout or interruption terminates and reaps the SDK worker, with bounded shutdown waits for both SDK and provider processes Latency starts immediately before the SDK call and ends when its result has been returned and discarded. Async calls are awaited on a persistent event loop. CPU is process CPU time during the timed batch, including Python and native threads. Fixture loading, process startup, warmup, preflight serialization, and report generation are excluded. Default SDK behavior is retained, so deferred background work can extend beyond a call's return; these metrics describe the measurement window, not the eventual cost of every callback diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/execution.py b/tests/rust-python-harness/strategies/e2e_benchmark/execution.py index 43fee4e67d3..72e833ac2b7 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/execution.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/execution.py @@ -5,15 +5,14 @@ import subprocess import sys import tempfile from collections.abc import Generator, Iterator -from concurrent.futures import Future, ThreadPoolExecutor -from contextlib import contextmanager +from contextlib import contextmanager, suppress from pathlib import Path from time import monotonic, sleep -from typing import TYPE_CHECKING, Final, TextIO, cast +from typing import TYPE_CHECKING, Final, TextIO import psutil -from .models import PREFIX, Backend, BenchmarkModel, Invocation, Measurement, Memory, Options, Ready, Route, Timing +from .models import Backend, BenchmarkModel, Invocation, Measurement, Memory, Options, Ready, Route, Timing from .provider import PYTHON_SENTINEL, provider_process if TYPE_CHECKING: @@ -30,13 +29,6 @@ def rss_bytes(process: psutil.Process) -> int: return ProcessMemory.model_validate(process.memory_info(), from_attributes=True).rss -def _read_message(stream: TextIO) -> str: - for line in stream: - if line.startswith(PREFIX): - return line.removeprefix(PREFIX) - raise RuntimeError("SDK worker exited without returning a measurement") - - @contextmanager def sdk_process(case_file: Path, backend: Backend, repo_root: Path, log: TextIO) -> Generator[subprocess.Popen[str]]: process: Final = subprocess.Popen( @@ -52,31 +44,40 @@ def sdk_process(case_file: Path, backend: Backend, repo_root: Path, log: TextIO) "PYTHONPATH": str(repo_root), }, stdin=subprocess.PIPE, - stdout=subprocess.PIPE, + stdout=log, stderr=log, text=True, ) try: yield process + except BaseException: + process.terminate() + raise finally: if process.stdin is not None: - process.stdin.close() + with suppress(BrokenPipeError): + process.stdin.close() try: process.wait(timeout=5) except subprocess.TimeoutExpired: process.kill() process.wait(timeout=5) - if process.stdout is not None: - process.stdout.close() -def sample_rss(process: psutil.Process, completed: Future[str], interval: float, timeout: float) -> Iterator[int]: - deadline: Final = monotonic() + timeout - while not completed.done(): +def wait_for_output( + output: Path, child: subprocess.Popen[str], options: Options, *, sample_memory: bool = False +) -> Iterator[int]: + deadline: Final = monotonic() + options.timeout + process: Final = psutil.Process(child.pid) if sample_memory else None + interval: Final = options.sample_interval_ms / 1000 if sample_memory else 0.01 + while not output.exists(): + if child.poll() is not None: + raise RuntimeError(f"SDK worker exited with code {child.returncode} before writing {output.name}") if monotonic() >= deadline: - raise TimeoutError("memory measurement timed out") - yield rss_bytes(process) - sleep(interval) + raise TimeoutError(f"SDK worker timed out waiting for {output.name}") + if process is not None: + yield rss_bytes(process) + sleep(min(interval, max(0, deadline - monotonic()))) def execute_phase( @@ -86,35 +87,40 @@ def execute_phase( directory: Final = Path(raw_directory) case_file: Final = directory / "invocation.json" case_file.write_text(invocation.model_dump_json()) + ready_file: Final = directory / "ready.json" + timing_file: Final = directory / "timing.json" with (directory / "worker.log").open("w+") as log: try: - with ThreadPoolExecutor(max_workers=1) as reader: - with sdk_process(case_file, backend, repo_root, log) as child: - assert child.stdout is not None and child.stdin is not None - stdout: Final = cast(TextIO, child.stdout) - ready: Final = Ready.model_validate_json( - reader.submit(_read_message, stdout).result(timeout=options.timeout) + with sdk_process(case_file, backend, repo_root, log) as child: + tuple(wait_for_output(ready_file, child, options)) + ready: Final = Ready.model_validate_json(ready_file.read_bytes()) + if invocation.phase == "timing": + child.communicate(input="go\n", timeout=options.timeout) + if child.returncode != 0: + raise RuntimeError(f"SDK worker exited with code {child.returncode}") + return ( + ready, + Timing.model_validate_json(timing_file.read_bytes()), + Memory(baseline_rss_bytes=0, sampled_peak_rss_bytes=0, retained_rss_bytes=0, samples=0), ) - process: Final = psutil.Process(child.pid) - baseline: Final = rss_bytes(process) if invocation.phase == "memory" else 0 - child.stdin.write("go\n") - child.stdin.flush() - result: Final = reader.submit(_read_message, stdout) - samples: Final = ( - tuple(sample_rss(process, result, options.sample_interval_ms / 1000, options.timeout)) - if invocation.phase == "memory" - else () - ) - timing: Final = Timing.model_validate_json(result.result(timeout=options.timeout)) - retained: Final = rss_bytes(process) if invocation.phase == "memory" else 0 - memory: Final = Memory( + process: Final = psutil.Process(child.pid) + baseline: Final = rss_bytes(process) + assert child.stdin is not None + child.stdin.write("go\n") + child.stdin.flush() + samples: Final = tuple(wait_for_output(timing_file, child, options, sample_memory=True)) + retained: Final = rss_bytes(process) + return ( + ready, + Timing.model_validate_json(timing_file.read_bytes()), + Memory( baseline_rss_bytes=baseline, sampled_peak_rss_bytes=max((baseline, retained, *samples)), retained_rss_bytes=retained, samples=len(samples), - ) - return ready, timing, memory - except (RuntimeError, OSError, ValueError, TimeoutError) as error: + ), + ) + except (RuntimeError, OSError, ValueError, TimeoutError, subprocess.TimeoutExpired, psutil.Error) as error: log.seek(0) raise RuntimeError(f"{backend}/{invocation.phase}: {error}\n{log.read()[-6000:]}") from error diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/models.py b/tests/rust-python-harness/strategies/e2e_benchmark/models.py index 12dbd27f320..6f5201a1a47 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/models.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/models.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Final, Literal +from typing import Literal from pydantic import BaseModel, ConfigDict, Field @@ -8,7 +8,6 @@ Backend = Literal["python", "rust"] Route = Literal["ocr", "aocr"] Phase = Literal["timing", "memory"] Profile = Literal["small", "request_medium", "request_large", "response_medium", "response_large"] -PREFIX: Final = "LITELLM_BENCHMARK " class BenchmarkModel(BaseModel): diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/provider.py b/tests/rust-python-harness/strategies/e2e_benchmark/provider.py index f8c0b1607c1..efed21f8302 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/provider.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/provider.py @@ -67,5 +67,5 @@ def provider_process(response: bytes, backend: Backend) -> Generator[str]: process.join(timeout=5) if process.is_alive(): process.kill() - process.join() + process.join(timeout=5) process.close() 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 367cea1c410..dd6e01f0d7f 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py @@ -2,9 +2,10 @@ from __future__ import annotations import asyncio import base64 -from concurrent.futures import Future +import subprocess +import sys from pathlib import Path -from time import sleep +from time import monotonic, sleep from typing import Final import httpx @@ -16,7 +17,7 @@ from litellm.llms.base_llm.ocr.transformation import OCRResponse from ...cli.catalog import load_catalog from ...shared.reporting.models import RunStatus -from .execution import execute_phase, sample_rss +from .execution import execute_phase, sdk_process, wait_for_output from .models import Invocation, Options from .provider import PYTHON_SENTINEL, provider_process from .reporting import percentile, render_measurements @@ -114,10 +115,43 @@ def test_memory_pass_does_not_accumulate_latency_samples() -> None: assert result.latency_ms == () -def test_memory_monitor_has_a_deadline() -> None: - pending: Final[Future[str]] = Future() - with pytest.raises(TimeoutError, match="memory measurement timed out"): - tuple(sample_rss(psutil.Process(), pending, interval=0.001, timeout=0.01)) +@pytest.mark.parametrize("sample_memory", (False, True)) +def test_worker_deadline_terminates_and_reaps_the_process(tmp_path: Path, sample_memory: bool) -> None: + case_file: Final = tmp_path / "invocation.json" + case_file.write_text(invocation().model_dump_json()) + existing_children: Final = frozenset(process.pid for process in psutil.Process().children()) + start: Final = monotonic() + with (tmp_path / "worker.log").open("w+") as log: + with pytest.raises(TimeoutError, match="timed out waiting for missing.json"): + with sdk_process(case_file, "python", REPO_ROOT, log) as child: + tuple( + wait_for_output( + tmp_path / "missing.json", + child, + Options(timeout=0.02, sample_interval_ms=10000), + sample_memory=sample_memory, + ) + ) + assert frozenset(process.pid for process in psutil.Process().children()) <= existing_children + assert monotonic() - start < 5 + + +def test_worker_cleanup_handles_buffered_input_after_early_exit(tmp_path: Path) -> None: + case_file: Final = tmp_path / "invocation.json" + case_file.write_text(invocation().model_dump_json()) + with (tmp_path / "worker.log").open("w+") as log: + with sdk_process(case_file, "python", REPO_ROOT, log) as child: + child.terminate() + child.wait(timeout=5) + assert child.stdin is not None + child.stdin.write("go\n") + assert child.returncode is not None + + +def test_worker_exit_is_detected_without_waiting_for_the_deadline(tmp_path: Path) -> None: + with subprocess.Popen((sys.executable, "-c", "raise SystemExit(7)"), cwd=REPO_ROOT, text=True) as child: + with pytest.raises(RuntimeError, match="exited with code 7"): + tuple(wait_for_output(tmp_path / "missing.json", child, Options(timeout=30))) def test_replay_rejects_python_fallback_during_rust_measurement() -> None: diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/worker.py b/tests/rust-python-harness/strategies/e2e_benchmark/worker.py index 2de690239a9..a3514f4d62d 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/worker.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/worker.py @@ -13,7 +13,7 @@ from typing import Final, cast from litellm.llms.base_llm.ocr.transformation import OCRResponse -from .models import PREFIX, Invocation, Ready, Timing +from .models import Invocation, Ready, Timing def _ready(response: OCRResponse) -> Ready: @@ -31,20 +31,22 @@ def _ready(response: OCRResponse) -> Ready: ) -def _emit(value: Ready | Timing) -> None: - print(PREFIX + value.model_dump_json(), flush=True) +def _publish(value: Ready | Timing, destination: Path) -> None: + temporary: Final = destination.with_suffix(".tmp") + temporary.write_text(value.model_dump_json()) + temporary.replace(destination) -def _handshake(ready: Ready) -> None: +def _handshake(ready: Ready, directory: Path) -> None: gc.collect() - _emit(ready) + _publish(ready, directory / "ready.json") if sys.stdin.readline().strip() != "go": raise RuntimeError("benchmark controller disconnected before measurement") -def _finish(timing: Timing) -> None: +def _finish(timing: Timing, directory: Path) -> None: gc.collect() - _emit(timing) + _publish(timing, directory / "timing.json") sys.stdin.readline() @@ -84,15 +86,15 @@ async def measure_async(call: Callable[[], Awaitable[OCRResponse]], invocation: return Timing(latency_ms=samples, cpu_ms=(process_time_ns() - cpu_start) / 1e6, elapsed_ms=elapsed / 1e6) -async def _run_async(call: Callable[[], Awaitable[OCRResponse]], invocation: Invocation) -> None: +async def _run_async(call: Callable[[], Awaitable[OCRResponse]], invocation: Invocation, directory: Path) -> None: for _ in range(invocation.warmup): await call() ready: Final = _ready(await call()) - _handshake(ready) - _finish(await measure_async(call, invocation)) + _handshake(ready, directory) + _finish(await measure_async(call, invocation), directory) -def run_worker(invocation: Invocation) -> None: +def run_worker(invocation: Invocation, directory: Path) -> None: import litellm kwargs: Final = { @@ -105,16 +107,17 @@ def run_worker(invocation: Invocation) -> None: } if invocation.route == "aocr": async_route: Final = cast(Callable[..., Awaitable[OCRResponse]], litellm.aocr) - asyncio.run(_run_async(lambda: async_route(**kwargs), invocation)) + asyncio.run(_run_async(lambda: async_route(**kwargs), invocation, directory)) return sync_route: Final = cast(Callable[..., OCRResponse], litellm.ocr) call: Final = lambda: sync_route(**kwargs) for _ in range(invocation.warmup): call() ready: Final = _ready(call()) - _handshake(ready) - _finish(measure_sync(call, invocation)) + _handshake(ready, directory) + _finish(measure_sync(call, invocation), directory) if __name__ == "__main__": - run_worker(Invocation.model_validate_json(Path(sys.argv[1]).read_bytes())) + case_file: Final = Path(sys.argv[1]) + run_worker(Invocation.model_validate_json(case_file.read_bytes()), case_file.parent) From e9dbe047a2ed65e9b9f4e58cc4bbf25569791292 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 13:14:53 -0700 Subject: [PATCH 04/10] 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 == () From 75ae4d591adc56387740dbe0619793cd635f8a13 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 13:22:44 -0700 Subject: [PATCH 05/10] fix(harness): bound benchmark provenance memory and document measurement limits --- .../strategies/e2e_benchmark/AGENTS.md | 34 +++++++- .../strategies/e2e_benchmark/README.md | 87 ------------------- .../strategies/e2e_benchmark/__init__.py | 2 +- .../e2e_benchmark/test_benchmark.py | 36 ++++++-- .../strategies/e2e_benchmark/worker.py | 15 ++-- 5 files changed, 74 insertions(+), 100 deletions(-) delete mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/README.md diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md b/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md index 8aa3a4a38d5..34d3cfefd4f 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md +++ b/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md @@ -1 +1,33 @@ -Measures Python and Rust SDK latency, CPU time, and process memory against deterministic local provider replays, outside correctness-check overhead +# What this is + +Measures Python and Rust SDK latency, process CPU time, and sampled RSS against local provider replays. Initial coverage is sync/async Mistral OCR at concurrency one. Other SDK functions remain explicitly unimplemented. Run locally, with no CI integration + +# How it works + +Derive five profiles from the existing `e2e_parity` recording without editing it. `small` uses a 32 KiB PDF and one response page. `request_medium` and `request_large` increase only the PDF to 256 KiB and 2 MiB. `response_medium` and `response_large` increase only the response to 16 and 128 pages. Insert PDF comment padding before the original final `startxref` and EOF trailer, preserving object offsets. Response pages are synthetic repetitions, independent of actual PDF content + +Generate fixtures in the controller, outside SDK worker imports. Each backend, route, profile, and repeat gets fresh timing and memory workers against a separate local HTTP provider process. Run Python and Rust sequentially, reversing their order on alternating repeats. The provider drains request bytes and serves preloaded responses without JSON parsing or capture. Its CPU and RSS are excluded + +Warmup, preflight response checks, bounded-memory native hashing, and garbage collection precede readiness. Require matching Python/Rust preflight digests and matching timing/memory worker digests. Every request checks the parity harness's User-Agent convention to reject Python fallback during a Rust run. Use `e2e_parity` for complete request and response semantics + +Time each SDK call until its returned result is discarded. Async calls share a persistent event loop. Process CPU and batch elapsed time include sample collection overhead and work on Python/native threads. Startup, fixture loading, warmup, preflight, final garbage collection, and reporting are excluded. Deferred callbacks can outlive the measured batch + +Sample only the SDK worker's RSS in the separate memory pass. Baseline follows warmup and garbage collection. Peak includes baseline, periodic samples, and the final sample. After RSS follows another garbage collection with input/client state resident. RSS includes native allocations and shared pages. Sampling can miss short peaks, and retained RSS does not prove a leak + +Publish readiness/results atomically in temporary JSON files, reserving stdout/stderr for diagnostics. Bound readiness, measurement, and shutdown waits. Terminate and reap workers on failure or interruption. Reject missing extensions, backend mismatches, exceptions, and incomplete samples. Preserve completed pairs in partial reports when a worker fails + +Report pooled p50/p95/p99 latency, CPU milliseconds per call, sequential calls per second, baseline/peak/after RSS, and Python p50 divided by backend p50. JSON retains raw per-repeat samples, fixture/extension hashes, Python version, options, platform, Git revision, and working-tree state. A pass confirms valid measurements without imposing performance thresholds. Use an idle host, inspect repeat variation, and treat short-run tail estimates cautiously. Loopback HTTP, allocator behavior, and deferred work affect results. Streaming, gateway overhead, live provider latency, and concurrent throughput are outside scope + +Editable `uv sync` builds a development extension. Build release explicitly and retain `--no-sync`: + +```sh +uv sync --frozen --python 3.12 +VIRTUAL_ENV="$PWD/.venv" uvx --from maturin==1.15.0 maturin develop --release +uv run --no-sync python -m tests.rust-python-harness run e2e_benchmark \ + --surface sdk --function ocr \ + --benchmark-arg=--output=/tmp/e2e-benchmark.json +``` + +Forward each option through `--benchmark-arg=...`. Defaults: `--iterations=100`, `--warmup=10`, `--repeats=3`, `--timeout=120` seconds, `--sample-interval-ms=5`. Repeat `--profile=NAME` or `--route=ocr|aocr` to select subsets. `--output=PATH` exports JSON. Invalid values and unwritable destinations produce handled CLI errors. `run all` includes this strategy. A smoke run can select `small`, `ocr`, 10 iterations, 2 warmups, and 1 repeat + +Run focused checks with `uv run --no-sync pytest -o consider_namespace_packages=true tests/rust-python-harness/strategies/e2e_benchmark tests/rust-python-harness/cli -q` diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/README.md b/tests/rust-python-harness/strategies/e2e_benchmark/README.md deleted file mode 100644 index 22b103e37a3..00000000000 --- a/tests/rust-python-harness/strategies/e2e_benchmark/README.md +++ /dev/null @@ -1,87 +0,0 @@ -# End-to-end SDK benchmark - -Compare `LITELLM_RUST=0` and `LITELLM_RUST=1` against a separate local HTTP provider process, with no real provider calls, credentials, or Docker required - -The initial workload covers synchronous and asynchronous Mistral OCR using the existing `e2e_parity` recording. Other SDK functions are explicitly unimplemented. This measures SDK calls including loopback HTTP transport and response construction. It does not measure gateway overhead, streaming, or concurrent load - -## Run - -Build the Rust extension in release mode first. An editable `uv sync` normally builds the development profile, which is unsuitable for a Python/Rust performance comparison - -```sh -uv sync --frozen --python 3.12 -VIRTUAL_ENV="$PWD/.venv" uvx --from maturin==1.15.0 maturin develop --release -uv run --no-sync python -m tests.rust-python-harness run e2e_benchmark \ - --surface sdk --function ocr \ - --benchmark-arg=--output=/tmp/e2e-benchmark.json -``` - -Keep `--no-sync` on the benchmark command so it uses the extension you just built. Use an otherwise idle machine and run the same command on both revisions when evaluating a change - -For a short smoke run: - -```sh -uv run --no-sync python -m tests.rust-python-harness run e2e_benchmark \ - --function ocr \ - --benchmark-arg=--profile=small \ - --benchmark-arg=--route=ocr \ - --benchmark-arg=--iterations=10 \ - --benchmark-arg=--warmup=2 \ - --benchmark-arg=--repeats=1 \ - --benchmark-arg=--output=/tmp/e2e-benchmark-smoke.json -``` - -`run all` also runs this strategy with its defaults. No CI integration is added - -## Workloads - -The seed cassette stays under `e2e_parity/sdk/ocr/fixtures/data`. The benchmark derives synthetic size variants in memory; it never edits or re-records the correctness fixtures - -| Profile | Inline PDF bytes | Response pages | -| --- | ---: | ---: | -| small | 32 KiB | 1 | -| request_medium | 256 KiB | 1 | -| request_large | 2 MiB | 1 | -| response_medium | 32 KiB | 16 | -| response_large | 32 KiB | 128 | - -Request variants add PDF comment padding before the final `startxref` marker, preserving existing object offsets and the EOF trailer. The SDK sends base64 plus JSON framing, so wire request sizes exceed the document sizes above. Response variants repeat recorded pages with contiguous indexes and adjusted usage. They exercise realistic response structure, but their page count intentionally varies independently of the input PDF's content - -## Measurements - -Each backend, route, size, and repeat gets a fresh SDK process for timing and another for memory. Python and Rust execute sequentially, with their order reversed on alternating repeats. The local provider serves preloaded bytes without parsing or capturing request JSON. Its CPU and RSS are outside the SDK measurements - -Workers warm up their clients and run an untimed response check before measuring. Python and Rust response digests must match. Every provider request also checks the existing parity harness's User-Agent convention: a Rust run using Python's HTTP path fails instead of reporting a comparison between two Python runs. Missing native extensions, SDK exceptions, timeouts, and incomplete samples fail the run. Workers publish readiness and results atomically in temporary JSON files; stdout and stderr go to a diagnostic log. The controller waits for timing workers to exit and samples RSS only for memory workers. A timeout or interruption terminates and reaps the SDK worker, with bounded shutdown waits for both SDK and provider processes - -Latency starts immediately before the SDK call and ends when its result has been returned and discarded. Async calls are awaited on a persistent event loop. CPU is process CPU time during the timed batch, including Python and native threads. Fixture loading, process startup, warmup, preflight serialization, and report generation are excluded. Default SDK behavior is retained, so deferred background work can extend beyond a call's return; these metrics describe the measurement window, not the eventual cost of every callback - -The memory controller uses `psutil` to sample only the SDK worker's RSS during a separate run, avoiding polling overhead in latency results. Baseline RSS is taken after warmup and garbage collection. Peak is the highest sampled RSS, including the baseline and final sample. After RSS is measured after the workload and another garbage collection, with input/client state still resident. RSS includes native allocations and shared resident pages, so it is not equivalent to Python heap size or uniquely owned memory. Sampling can miss brief peaks; these values are not an exact allocator high-water mark or proof of a leak - -The terminal reports pooled p50/p95/p99 latency, CPU milliseconds per call, sequential calls per second, baseline/peak/after RSS, and speedup (`Python p50 / backend p50`). Throughput is at concurrency one, not saturation capacity. Short runs cannot estimate tail latency reliably. The JSON retains each repeat's raw latency samples, CPU and memory measurements, input/response sizes, seed hash, Python version, native extension hash, settings, platform, Git revision, and whether the working tree has changes - -## Options - -Pass each option through `--benchmark-arg=...` - -| Option | Default | Meaning | -| --- | --- | --- | -| `--iterations=N` | 100 | Measured calls per worker | -| `--warmup=N` | 10 | Warmup calls, followed by one preflight call | -| `--repeats=N` | 3 | Fresh paired runs per workload | -| `--profile=NAME` | All five | Select a size profile; repeat for several | -| `--route=ocr` or `--route=aocr` | Both | Select SDK entrypoint; repeat for both | -| `--timeout=SECONDS` | 120 | Worker readiness and measurement deadline | -| `--sample-interval-ms=N` | 5 | Memory sampling interval, at least 1 ms | -| `--output=PATH` | None | Write a JSON report, including partial results on worker failure | - -Run the strategy tests and existing harness checks with: - -```sh -uv run --no-sync pytest -o consider_namespace_packages=true \ - tests/rust-python-harness/strategies/e2e_benchmark \ - tests/rust-python-harness/shared tests/rust-python-harness/cli \ - tests/rust-python-harness/strategies/unit_tests_mapping \ - tests/rust-python-harness/strategies/unit_tests_parity \ - tests/rust-python-harness/strategies/unit_tests_rust \ - tests/test_rust_python_harness.py -q -``` diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/__init__.py b/tests/rust-python-harness/strategies/e2e_benchmark/__init__.py index 692d3d485af..0160ab2c602 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/__init__.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/__init__.py @@ -38,6 +38,6 @@ STRATEGY: Final = StrategyDefinition( surfaces=("sdk",), runner_argument=RunnerArgumentDefinition( option="--benchmark-arg", - help="benchmark option, e.g. --benchmark-arg=--iterations=100; see the strategy README", + help="benchmark option, e.g. --benchmark-arg=--iterations=100; see the strategy AGENTS.md", ), ) 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 f6449f5d39c..696156dddec 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py @@ -2,8 +2,10 @@ from __future__ import annotations import asyncio import base64 +import hashlib import subprocess import sys +import tracemalloc from pathlib import Path from time import monotonic, sleep from typing import Final @@ -18,11 +20,11 @@ from litellm.llms.base_llm.ocr.transformation import OCRResponse 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 .models import Invocation, Options, Route from .provider import provider_process from .reporting import percentile, render_measurements from .runner import Report, parse_options -from .worker import measure_async, measure_sync +from .worker import file_sha256, measure_async, measure_sync from .workloads import JSON_OBJECT, JSON_PAGES, ocr_workload, padded_pdf REPO_ROOT: Final = Path(__file__).resolve().parents[4] @@ -162,10 +164,13 @@ def test_replay_rejects_python_fallback_during_rust_measurement() -> None: assert "backend mismatch" in response.text -def test_worker_errors_are_reported_instead_of_counted_as_fast_calls() -> None: +@pytest.mark.parametrize("route", ("ocr", "aocr")) +def test_worker_errors_are_reported_instead_of_counted_as_fast_calls(route: Route) -> None: workload: Final = ocr_workload("small") with provider_process(workload.response, "rust") as url: - request: Final = invocation().model_copy(update={"provider_url": url, "document_url": workload.document_url}) + request: Final = invocation().model_copy( + update={"provider_url": url, "document_url": workload.document_url, "route": route} + ) with pytest.raises(RuntimeError, match="backend mismatch"): execute_phase(request, "python", Options(iterations=3, warmup=1), REPO_ROOT) @@ -184,7 +189,7 @@ def test_cli_runs_both_backends_and_exports_measurements(tmp_path: Path, capsys: "--benchmark-arg=--route=aocr", "--benchmark-arg=--iterations=3", "--benchmark-arg=--warmup=1", - "--benchmark-arg=--repeats=1", + "--benchmark-arg=--repeats=2", f"--benchmark-arg=--output={output}", ) ) @@ -192,7 +197,12 @@ def test_cli_runs_both_backends_and_exports_measurements(tmp_path: Path, capsys: 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 tuple((value.repeat, value.backend) for value in report.measurements) == ( + (0, "python"), + (0, "rust"), + (1, "rust"), + (1, "python"), + ) assert len({value.ready.response_digest for value in report.measurements}) == 1 for value in report.measurements: assert len(value.timing.latency_ms) == 3 @@ -242,3 +252,17 @@ def test_cli_reports_unsupported_functions_without_measurements( assert "not implemented" in captured.out report: Final = Report.model_validate_json(output.read_bytes()) assert report.measurements == () + + +def test_native_provenance_hash_uses_bounded_memory(tmp_path: Path) -> None: + native: Final = tmp_path / "native.so" + payload: Final = b"native-extension-data" * (1024 * 1024) + native.write_bytes(payload) + expected: Final = hashlib.sha256(payload).hexdigest() + tracemalloc.start() + try: + assert file_sha256(native) == expected + _, peak = tracemalloc.get_traced_memory() + assert peak < 1024 * 1024 + finally: + tracemalloc.stop() diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/worker.py b/tests/rust-python-harness/strategies/e2e_benchmark/worker.py index a3514f4d62d..89c5d1cd896 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/worker.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/worker.py @@ -16,6 +16,11 @@ from litellm.llms.base_llm.ocr.transformation import OCRResponse from .models import Invocation, Ready, Timing +def file_sha256(path: Path) -> str: + with path.open("rb") as source: + return hashlib.file_digest(source, "sha256").hexdigest() + + def _ready(response: OCRResponse) -> Ready: from litellm.rust_bridge import get_native_bridge from litellm.rust_bridge.configuration import rust_enabled @@ -27,7 +32,7 @@ def _ready(response: OCRResponse) -> Ready: return Ready( response_digest=hashlib.sha256(json.dumps(response.model_dump(), sort_keys=True).encode()).hexdigest(), python_version=platform.python_version(), - native_sha256=hashlib.sha256(Path(native_path).read_bytes()).hexdigest() if native_path else None, + native_sha256=file_sha256(Path(native_path)) if native_path else None, ) @@ -63,24 +68,24 @@ async def _async_sample(call: Callable[[], Awaitable[OCRResponse]]) -> float: def measure_sync(call: Callable[[], OCRResponse], invocation: Invocation) -> Timing: - cpu_start: Final = process_time_ns() - wall_start: Final = perf_counter_ns() if invocation.phase == "memory": for _ in range(invocation.iterations): call() return Timing(latency_ms=(), cpu_ms=0, elapsed_ms=0) + cpu_start: Final = process_time_ns() + wall_start: Final = perf_counter_ns() samples: Final = tuple(_sync_sample(call) for _ in range(invocation.iterations)) elapsed: Final = perf_counter_ns() - wall_start return Timing(latency_ms=samples, cpu_ms=(process_time_ns() - cpu_start) / 1e6, elapsed_ms=elapsed / 1e6) async def measure_async(call: Callable[[], Awaitable[OCRResponse]], invocation: Invocation) -> Timing: - cpu_start: Final = process_time_ns() - wall_start: Final = perf_counter_ns() if invocation.phase == "memory": for _ in range(invocation.iterations): await call() return Timing(latency_ms=(), cpu_ms=0, elapsed_ms=0) + cpu_start: Final = process_time_ns() + wall_start: Final = perf_counter_ns() samples: Final = tuple([await _async_sample(call) for _ in range(invocation.iterations)]) elapsed: Final = perf_counter_ns() - wall_start return Timing(latency_ms=samples, cpu_ms=(process_time_ns() - cpu_start) / 1e6, elapsed_ms=elapsed / 1e6) From 878f0da9aba282e3898ac5954f70e6fd3b42a328 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 13:28:08 -0700 Subject: [PATCH 06/10] fix(harness): support bounded provenance hashing on Python 3.10 --- tests/rust-python-harness/strategies/e2e_benchmark/worker.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/worker.py b/tests/rust-python-harness/strategies/e2e_benchmark/worker.py index 89c5d1cd896..201a2c6339e 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/worker.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/worker.py @@ -17,8 +17,11 @@ from .models import Invocation, Ready, Timing def file_sha256(path: Path) -> str: + digest: Final = hashlib.sha256() with path.open("rb") as source: - return hashlib.file_digest(source, "sha256").hexdigest() + for chunk in iter(lambda: source.read(256 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() def _ready(response: OCRResponse) -> Ready: From 34e2ef97f0aa8cca558e3098c51c96ef960d2ccd Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 17:21:44 -0700 Subject: [PATCH 07/10] feat: add CodSpeed walltime benchmarks and reliability diagnostics --- .github/workflows/codspeed.yml | 61 +++++++ tests/rust-python-harness/AGENTS.md | 2 +- .../strategies/e2e_benchmark/AGENTS.md | 25 ++- .../strategies/e2e_benchmark/codspeed.py | 161 ++++++++++++++++++ .../strategies/e2e_benchmark/execution.py | 1 + .../strategies/e2e_benchmark/models.py | 4 +- .../strategies/e2e_benchmark/reporting.py | 52 +++++- .../strategies/e2e_benchmark/runner.py | 9 +- .../e2e_benchmark/test_benchmark.py | 112 +++++++++++- .../strategies/e2e_benchmark/worker.py | 21 ++- 10 files changed, 427 insertions(+), 21 deletions(-) create mode 100644 tests/rust-python-harness/strategies/e2e_benchmark/codspeed.py diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 7e013b7bb0b..dd76609e26e 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -8,6 +8,8 @@ on: paths: - "litellm/**" - "tests/benchmarks/**" + - "tests/rust-python-harness/**" + - "litellm-rust/**" - "pyproject.toml" - "uv.lock" - ".github/workflows/codspeed.yml" @@ -20,6 +22,8 @@ on: paths: - "litellm/**" - "tests/benchmarks/**" + - "tests/rust-python-harness/**" + - "litellm-rust/**" - "pyproject.toml" - "uv.lock" - ".github/workflows/codspeed.yml" @@ -92,3 +96,60 @@ jobs: -p pytest_codspeed.plugin tests/benchmarks/ --codspeed + + e2e-walltime: + runs-on: codspeed-macro + timeout-minutes: 30 + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + litellm-rust/target + key: ${{ runner.os }}-${{ runner.arch }}-e2e-release-${{ hashFiles('litellm-rust/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-${{ runner.arch }}-e2e-release- + + - name: Build environment + run: >- + env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 + uv run --frozen --no-default-groups + --with pytest==8.3.5 + --with pytest-codspeed==5.0.3 + --with psutil==7.2.2 + --with "mcp>=1.26.0,<2.0" + --with "a2a-sdk>=1.1.0,<2.0" + pytest -p pytest_codspeed.plugin + --import-mode=importlib -o consider_namespace_packages=true + tests/rust-python-harness/strategies/e2e_benchmark/codspeed.py + --collect-only -q + + - name: Build release extension + run: VIRTUAL_ENV="$PWD/.venv" uvx --from maturin==1.15.0 maturin develop --release + + - name: Run OCR walltime benchmarks + uses: CodSpeedHQ/action@1c8ae4843586d3ba879736b7f6b7b0c990757fab # v4.12.1 + with: + mode: walltime + run: >- + uv run --no-sync --no-default-groups + --with pytest==8.3.5 + --with pytest-codspeed==5.0.3 + --with psutil==7.2.2 + --with "mcp>=1.26.0,<2.0" + --with "a2a-sdk>=1.1.0,<2.0" + python -m tests.rust-python-harness.strategies.e2e_benchmark.codspeed diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md index 700af8e41d2..31b55d1f6eb 100644 --- a/tests/rust-python-harness/AGENTS.md +++ b/tests/rust-python-harness/AGENTS.md @@ -58,7 +58,7 @@ tests/rust-python-harness/ - A strategy is a folder under `strategies/` with a one-line `AGENTS.md` and an `__init__.py` exporting exactly one `STRATEGY: StrategyDefinition`; its id must equal the folder name - `shared/reporting/strategy.py` is the contract: runnable module/suite specs, not-implemented/skipped specs, the runner protocol, and `StrategyDefinition` - Every `STRATEGY` explicitly classifies every SDK function; surface-aware strategies declare their surfaces and classify the complete surface-by-function matrix -- Run locally only; no CI integration +- Run locally; `e2e_benchmark` also has a CodSpeed walltime job - `python -m tests.rust-python-harness run |all` runs the selected strategy; `--function` is common, while each strategy exposes only its supported options - Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr` - `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md b/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md index 34d3cfefd4f..d8cbf38ca7d 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md +++ b/tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md @@ -1,12 +1,12 @@ # What this is -Measures Python and Rust SDK latency, process CPU time, and sampled RSS against local provider replays. Initial coverage is sync/async Mistral OCR at concurrency one. Other SDK functions remain explicitly unimplemented. Run locally, with no CI integration +Measures Python and Rust SDK latency, process CPU time, and sampled RSS against local provider replays. Initial coverage is sync/async Mistral OCR at concurrency one. Other SDK functions remain explicitly unimplemented. Keep implementation in this strategy folder. The local CLI reports latency, CPU, and RSS; the CodSpeed adapter uploads walltime measurements from the existing repository workflow # How it works Derive five profiles from the existing `e2e_parity` recording without editing it. `small` uses a 32 KiB PDF and one response page. `request_medium` and `request_large` increase only the PDF to 256 KiB and 2 MiB. `response_medium` and `response_large` increase only the response to 16 and 128 pages. Insert PDF comment padding before the original final `startxref` and EOF trailer, preserving object offsets. Response pages are synthetic repetitions, independent of actual PDF content -Generate fixtures in the controller, outside SDK worker imports. Each backend, route, profile, and repeat gets fresh timing and memory workers against a separate local HTTP provider process. Run Python and Rust sequentially, reversing their order on alternating repeats. The provider drains request bytes and serves preloaded responses without JSON parsing or capture. Its CPU and RSS are excluded +Generate fixtures in the controller, outside SDK worker imports. Each backend, route, profile, and repeat gets fresh timing and memory workers against a separate local HTTP provider process. Run Python and Rust sequentially, reversing their order on alternating repeats. The default four repeats balance backend order. Timing workers run at least the requested iteration count and minimum duration; the separate memory workers use the same fixed iteration count for both backends The provider drains request bytes and serves preloaded responses without JSON parsing or capture. Its CPU and RSS are excluded Warmup, preflight response checks, bounded-memory native hashing, and garbage collection precede readiness. Require matching Python/Rust preflight digests and matching timing/memory worker digests. Every request checks the parity harness's User-Agent convention to reject Python fallback during a Rust run. Use `e2e_parity` for complete request and response semantics @@ -16,7 +16,7 @@ Sample only the SDK worker's RSS in the separate memory pass. Baseline follows w Publish readiness/results atomically in temporary JSON files, reserving stdout/stderr for diagnostics. Bound readiness, measurement, and shutdown waits. Terminate and reap workers on failure or interruption. Reject missing extensions, backend mismatches, exceptions, and incomplete samples. Preserve completed pairs in partial reports when a worker fails -Report pooled p50/p95/p99 latency, CPU milliseconds per call, sequential calls per second, baseline/peak/after RSS, and Python p50 divided by backend p50. JSON retains raw per-repeat samples, fixture/extension hashes, Python version, options, platform, Git revision, and working-tree state. A pass confirms valid measurements without imposing performance thresholds. Use an idle host, inspect repeat variation, and treat short-run tail estimates cautiously. Loopback HTTP, allocator behavior, and deferred work affect results. Streaming, gateway overhead, live provider latency, and concurrent throughput are outside scope +Report pooled p50/p95/p99 and mean/standard-deviation latency, per-repeat sample counts, elapsed time, medians and throughput, CPU milliseconds per call, sequential calls per second, baseline/peak/after RSS, and Python p50 divided by backend p50. JSON retains raw per-repeat samples, fixture/extension hashes, Python version, options, platform, Git revision, and working-tree state. JSON schema version 2 also exports diagnostic warnings. A pass confirms completed measurements without establishing statistical significance or imposing performance thresholds. Warn for single repeats, unbalanced order, batches shorter than one second, or a repeat-median range exceeding 10% of its median. These are diagnostic heuristics; absence of warnings is not proof of stability. Pooled statistics weight each call equally, so duration-based sampling can weight faster repeats more heavily. Inspect independent repeat statistics before interpreting pooled ratios. Calls per second follows batch mean time, not median latency. Keep slow calls in the data Use an idle host, inspect repeat variation, and treat short-run tail estimates cautiously. Loopback HTTP, allocator behavior, and deferred work affect results. Streaming, gateway overhead, live provider latency, and concurrent throughput are outside scope Editable `uv sync` builds a development extension. Build release explicitly and retain `--no-sync`: @@ -28,6 +28,23 @@ uv run --no-sync python -m tests.rust-python-harness run e2e_benchmark \ --benchmark-arg=--output=/tmp/e2e-benchmark.json ``` -Forward each option through `--benchmark-arg=...`. Defaults: `--iterations=100`, `--warmup=10`, `--repeats=3`, `--timeout=120` seconds, `--sample-interval-ms=5`. Repeat `--profile=NAME` or `--route=ocr|aocr` to select subsets. `--output=PATH` exports JSON. Invalid values and unwritable destinations produce handled CLI errors. `run all` includes this strategy. A smoke run can select `small`, `ocr`, 10 iterations, 2 warmups, and 1 repeat +Forward each option through `--benchmark-arg=...`. Defaults: `--iterations=100`, `--warmup=10`, `--repeats=4`, `--min-time=1` second, `--timeout=120` seconds, `--sample-interval-ms=5`. Repeat `--profile=NAME` or `--route=ocr|aocr` to select subsets. `--output=PATH` exports JSON. Invalid values and unwritable destinations produce handled CLI errors. `run all` includes this strategy. A smoke run can select `small`, `ocr`, 10 iterations, 2 warmups, 1 repeat, and `--min-time=0` Run focused checks with `uv run --no-sync pytest -o consider_namespace_packages=true tests/rust-python-harness/strategies/e2e_benchmark tests/rust-python-harness/cli -q` + +The CodSpeed job in `.github/workflows/codspeed.yml` uses `mode: walltime` on `codspeed-macro`, with release compilation outside instrumentation. It leaves the existing CPU simulation job intact. Macro runners need explicit access to this public repository through the organization runner group. The ARM64 release cache includes runner architecture + +Run the CodSpeed adapter in place: + +```sh +uv pip install --python .venv/bin/python pytest-codspeed==5.0.3 +uv run --no-sync python -m tests.rust-python-harness.strategies.e2e_benchmark.codspeed +``` + +Use `--profile=small`, `--route=ocr|aocr`, and `--max-time=5` seconds to select work. Each backend/profile/route runs once in a fresh pytest process with a stable CodSpeed benchmark ID. Setup, provider startup, preflight validation, and teardown are outside the benchmark fixture. Compare Python/Rust response digests after both processes finish. The same provider rejects backend fallback. Timed warmup and round-count selection are supplied by CodSpeed, which stores history and profiling data in CI. Repeated rounds in one process do not replace independent process repeats; use the local CLI for that audit + +Set CodSpeed `min_time=0` so each round measures one full SDK call. This also avoids the pinned plugin's terminal display dividing normalized timings by iterations a second time for multi-call rounds. Async measurements await completion on a persistent loop, including `run_until_complete` entry/exit per call; local CLI async samples run inside the loop, so compare backends within each instrument rather than equating their absolute timings. Neither path overrides the logging executor. CodSpeed's best-time statistic is not the CLI's median or throughput. Memory and process CPU remain separate local measurements + +Interpret results per route and profile. A synchronous small-request improvement does not establish an async improvement or monotonic scaling. Neither path isolates PyO3 overhead. Sampled peak RSS reductions do not imply equivalent retained-memory reductions. Use a separate bridge microbenchmark and allocation profiling to investigate causes + +Methodology references: [Switowski](https://switowski.com/blog/how-to-benchmark-python-code/) on repeatability and setup boundaries, [CodSpeed](https://codspeed.io/docs/instruments/walltime) on I/O-inclusive measurement, [UCL](https://github-pages.arc.ucl.ac.uk/python-tooling/pages/benchmarking-profiling.html) on benchmarking versus profiling, and [Bencher](https://bencher.dev/learn/benchmarking/python/pytest-benchmark/) on distributions and mean-based throughput diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/codspeed.py b/tests/rust-python-harness/strategies/e2e_benchmark/codspeed.py new file mode 100644 index 00000000000..ffe36a4441a --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_benchmark/codspeed.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import asyncio +import gc +import math +import os +import signal +import subprocess +import sys +import tempfile +from collections.abc import Awaitable, Callable +from pathlib import Path +from typing import TYPE_CHECKING, Final, cast + +import click +import pytest + +from litellm.llms.base_llm.ocr.transformation import OCRResponse + +from .constants import PYTHON_SENTINEL +from .models import Backend, Options, Profile, Ready, Route +from .provider import provider_process +from .worker import capture_ready +from .workloads import ocr_workload + +if TYPE_CHECKING: + from pytest_codspeed.plugin import BenchmarkFixture + +CASES: Final = tuple( + (backend, route, profile) + for profile in Options().profiles + for route in Options().routes + for backend in ("python", "rust") +) + + +@pytest.mark.parametrize(("backend", "route", "profile"), CASES, ids=tuple("-".join(case) for case in CASES)) +@pytest.mark.benchmark(min_time=0) +def test_sdk(benchmark: BenchmarkFixture, backend: Backend, route: Route, profile: Profile) -> None: + import litellm + + assert os.environ.get("LITELLM_RUST") == ("1" if backend == "rust" else "0"), "use the codspeed module CLI" + workload: Final = ocr_workload(profile) + with provider_process(workload.response, backend) as url: + kwargs: Final = { + "model": workload.model, + "document": {"type": "document_url", "document_url": workload.document_url}, + "api_key": "benchmark-local-only", + "api_base": url, + "timeout": 10, + "num_retries": 0, + } + if route == "ocr": + sync_call: Final = cast(Callable[..., OCRResponse], litellm.ocr) + ready: Final = capture_ready(sync_call(**kwargs)) + + def call_sync() -> None: + sync_call(**kwargs) + + gc.collect() + benchmark(call_sync) + Path(os.environ["LITELLM_BENCHMARK_READY"]).write_text(ready.model_dump_json()) + return + async_call: Final = cast(Callable[..., Awaitable[OCRResponse]], litellm.aocr) + loop: Final = asyncio.new_event_loop() + try: + async_ready: Final = capture_ready(loop.run_until_complete(async_call(**kwargs))) + + async def call_async() -> None: + await async_call(**kwargs) + + def call_on_loop() -> None: + loop.run_until_complete(call_async()) + + gc.collect() + benchmark(call_on_loop) + Path(os.environ["LITELLM_BENCHMARK_READY"]).write_text(async_ready.model_dump_json()) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + +def run_case(backend: Backend, route: Route, profile: Profile, ready_file: Path, max_time: float) -> Ready: + root: Final = Path(__file__).resolve().parents[4] + nodeid: Final = f"{Path(__file__).relative_to(root)}::test_sdk[{backend}-{route}-{profile}]" + with subprocess.Popen( + ( + sys.executable, + "-m", + "pytest", + "-p", + "pytest_codspeed.plugin", + "-c", + os.devnull, + f"--rootdir={root}", + "--import-mode=importlib", + "-o", + "consider_namespace_packages=true", + "--codspeed", + "--codspeed-mode=walltime", + "--codspeed-warmup-time=1", + f"--codspeed-max-time={max_time}", + nodeid, + "-q", + ), + cwd=root, + env={ + **os.environ, + "PYTEST_DISABLE_PLUGIN_AUTOLOAD": "1", + "LITELLM_RUST": "1" if backend == "rust" else "0", + "LITELLM_USER_AGENT": PYTHON_SENTINEL, + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_BENCHMARK_READY": str(ready_file), + "NO_PROXY": "127.0.0.1,localhost", + "no_proxy": "127.0.0.1,localhost", + }, + start_new_session=True, + ) as child: + try: + code: Final = child.wait(timeout=max_time + 120) + if code: + raise click.ClickException(f"CodSpeed benchmark failed: {backend}/{route}/{profile}, exit {code}") + except subprocess.TimeoutExpired as error: + raise click.ClickException(f"CodSpeed worker timed out: {backend}/{route}/{profile}") from error + finally: + try: + os.killpg(child.pid, signal.SIGTERM) + except ProcessLookupError: + pass + try: + child.wait(timeout=5) + except subprocess.TimeoutExpired: + os.killpg(child.pid, signal.SIGKILL) + child.wait() + return Ready.model_validate_json(ready_file.read_bytes()) + + +def run_pair(route: Route, profile: Profile, directory: Path, max_time: float) -> None: + pair: Final = tuple( + run_case(backend, route, profile, directory / f"{backend}.json", max_time) for backend in ("python", "rust") + ) + if pair[0].response_digest != pair[1].response_digest: + raise click.ClickException(f"Python/Rust responses differ for {route}/{profile}; run e2e_parity") + + +@click.command() +@click.option("--profile", "profiles", multiple=True, type=click.Choice(Options().profiles), default=Options().profiles) +@click.option("--route", "routes", multiple=True, type=click.Choice(Options().routes), default=Options().routes) +@click.option("--max-time", type=click.FloatRange(min=0.1, max=60), default=5.0, show_default=True) +def main(profiles: tuple[Profile, ...], routes: tuple[Route, ...], max_time: float) -> None: + """Run isolated SDK workers with CodSpeed walltime calibration and profiling.""" + if not math.isfinite(max_time): + raise click.BadParameter("must be finite", param_hint="--max-time") + with tempfile.TemporaryDirectory(prefix="litellm-codspeed-") as directory: + for profile in dict.fromkeys(profiles): + for route in dict.fromkeys(routes): + run_pair(route, profile, Path(directory), max_time) + + +if __name__ == "__main__": + main() diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/execution.py b/tests/rust-python-harness/strategies/e2e_benchmark/execution.py index d24d1107469..5f312e6f322 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/execution.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/execution.py @@ -138,6 +138,7 @@ def benchmark( iterations=options.iterations, warmup=options.warmup, phase="timing", + min_time=options.min_time, ) ready, timing, _ = execute_phase(invocation, backend, options, repo_root) memory_ready, _, memory = execute_phase( diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/models.py b/tests/rust-python-harness/strategies/e2e_benchmark/models.py index 78875b21fe9..2cd28e8ced1 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/models.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/models.py @@ -17,7 +17,8 @@ class BenchmarkModel(BaseModel): class Options(BenchmarkModel): iterations: int = Field(default=100, ge=1) warmup: int = Field(default=10, ge=1) - repeats: int = Field(default=3, ge=1) + repeats: int = Field(default=4, ge=1) + min_time: float = Field(default=1, ge=0, allow_inf_nan=False) profiles: tuple[Profile, ...] = ("small", "request_medium", "request_large", "response_medium", "response_large") routes: tuple[Route, ...] = ("ocr", "aocr") timeout: float = Field(default=120, gt=0, allow_inf_nan=False) @@ -33,6 +34,7 @@ class Invocation(BenchmarkModel): iterations: int warmup: int phase: Phase + min_time: float = Field(default=0, ge=0, allow_inf_nan=False) class Ready(BenchmarkModel): diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/reporting.py b/tests/rust-python-harness/strategies/e2e_benchmark/reporting.py index a996d216032..af457ca8937 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/reporting.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/reporting.py @@ -32,10 +32,38 @@ def measurements(results: Sequence[CaseResult]) -> tuple[Measurement, ...]: ) +def measurement_warnings(values: tuple[Measurement, ...]) -> tuple[str, ...]: + keys: Final = tuple(dict.fromkeys((value.route, value.profile, value.backend) for value in values)) + + def warnings(route: str, profile: str, backend: str) -> tuple[str, ...]: + group: Final = tuple( + value for value in values if (value.route, value.profile, value.backend) == (route, profile, backend) + ) + medians: Final = tuple(statistics.median(value.timing.latency_ms) for value in group) + prefix: Final = f"{route}/{profile}/{backend}" + return ( + *((f"{prefix}: one process repeat cannot establish repeatability",) if len(group) == 1 else ()), + *((f"{prefix}: backend order is unbalanced",) if len(group) % 2 else ()), + *( + (f"{prefix}: batch shorter than 1 second; increase --min-time",) + if any(value.timing.elapsed_ms < 1000 for value in group) + else () + ), + *( + (f"{prefix}: repeat p50 range exceeds 10% of its median; investigate variability",) + if max(medians) - min(medians) > statistics.median(medians) * 0.1 + else () + ), + ) + + return tuple(warning for route, profile, backend in keys for warning in warnings(route, profile, backend)) + + def render_measurements(values: tuple[Measurement, ...]) -> str: keys: Final = tuple(dict.fromkeys((value.route, value.profile) for value in values)) header: Final = ( - "route/profile | backend | p50/p95/p99 ms | CPU ms/call | calls/s | RSS baseline/peak/after MiB | speedup" + "route/profile | backend | p50/p95/p99 ms | mean/stdev ms | CPU ms/call | calls/s | " + "RSS baseline/peak/after MiB | pooled p50 ratio" ) def row(group: tuple[Measurement, ...], baseline: float) -> str: @@ -50,7 +78,8 @@ def render_measurements(values: tuple[Measurement, ...]) -> str: ) return ( f"{group[0].route}/{group[0].profile} | {group[0].backend} | " - f"{median:.3f}/{percentile(samples, 0.95):.3f}/{percentile(samples, 0.99):.3f} | {cpu:.3f} | {rps:.1f} | " + f"{median:.3f}/{percentile(samples, 0.95):.3f}/{percentile(samples, 0.99):.3f} | " + f"{statistics.mean(samples):.3f}/{statistics.pstdev(samples):.3f} | {cpu:.3f} | {rps:.1f} | " f"{'/'.join(f'{value / 2**20:.1f}' for value in rss)} | {baseline / median:.2f}x" ) @@ -64,7 +93,24 @@ def render_measurements(values: tuple[Measurement, ...]) -> str: baseline: Final = statistics.median(sample for value in python for sample in value.timing.latency_ms) return row(python, baseline), row(rust, baseline) - return "\n".join((header, *(line for route, profile in keys for line in rows(route, profile)))) + repeats: Final = tuple( + f"{value.route}/{value.profile} | {value.backend} | {value.repeat + 1} | " + f"{len(value.timing.latency_ms)} | {value.timing.elapsed_ms:.1f} | " + f"{statistics.median(value.timing.latency_ms):.3f} | {statistics.mean(value.timing.latency_ms):.3f} | " + f"{len(value.timing.latency_ms) * 1000 / value.timing.elapsed_ms:.1f}" + for value in values + ) + return "\n".join( + ( + header, + *(line for route, profile in keys for line in rows(route, profile)), + "Per-repeat measurements (independent processes):", + "route/profile | backend | repeat | calls | batch ms | p50 ms | mean ms | calls/s", + *repeats, + "Completed measurements do not establish statistical significance or isolate PyO3 overhead", + *(f"WARNING: {warning}" for warning in measurement_warnings(values)), + ) + ) def render_benchmark_results(results: Sequence[CaseResult]) -> tuple[ReportSection, ...]: diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/runner.py b/tests/rust-python-harness/strategies/e2e_benchmark/runner.py index e72c817368d..4390c701742 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/runner.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/runner.py @@ -14,20 +14,21 @@ from ...shared.reporting.models import CaseResult, HarnessCase, HarnessRun, Resu from ...shared.reporting.strategy import ModuleCaseSpec, UpdateCallback from .execution import benchmark from .models import Backend, BenchmarkModel, Measurement, Options, Profile, Route -from .reporting import ARTIFACT_KIND, MEASUREMENTS, measurements +from .reporting import ARTIFACT_KIND, MEASUREMENTS, measurements, measurement_warnings if TYPE_CHECKING: from .workloads import Workload class Report(BenchmarkModel): - schema_version: int = 1 + schema_version: int = 2 revision: str working_tree_dirty: bool platform: str options: Options measurements: tuple[Measurement, ...] failures: tuple[tuple[str, str], ...] + warnings: tuple[str, ...] = () def parse_options(arguments: Sequence[str]) -> Options: @@ -36,6 +37,7 @@ def parse_options(arguments: Sequence[str]) -> Options: "e2e_benchmark", params=[ click.Option(("--iterations",), type=click.IntRange(min=1), default=defaults.iterations), + click.Option(("--min-time",), type=click.FloatRange(min=0), default=defaults.min_time), click.Option(("--warmup",), type=click.IntRange(min=1), default=defaults.warmup), click.Option(("--repeats",), type=click.IntRange(min=1), default=defaults.repeats), click.Option( @@ -66,7 +68,7 @@ def _run_pair( pair: Final = tuple(benchmark(workload, route, backend, repeat, options, repo_root) for backend in order) if pair[0].ready.response_digest != pair[1].ready.response_digest: raise ValueError("Python and Rust preflight SDK responses differ; run e2e_parity before comparing performance") - if any(len(value.timing.latency_ms) != options.iterations for value in pair): + if any(len(value.timing.latency_ms) < options.iterations for value in pair): raise ValueError("SDK worker returned an incomplete measurement") return pair @@ -142,6 +144,7 @@ def run_benchmark_cases( options=options, measurements=measurements(tuple(run.results.values())), failures=tuple(run.failures), + warnings=measurement_warnings(measurements(tuple(run.results.values()))), ) try: Path(options.output).write_text(report.model_dump_json(indent=2) + "\n") 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 696156dddec..e62ceaefd95 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py @@ -20,9 +20,9 @@ from litellm.llms.base_llm.ocr.transformation import OCRResponse from ...cli import main from .execution import execute_phase, sdk_process, wait_for_output from .constants import PYTHON_SENTINEL -from .models import Invocation, Options, Route +from .models import Backend, Invocation, Measurement, Memory, Options, Ready, Route, Timing from .provider import provider_process -from .reporting import percentile, render_measurements +from .reporting import measurement_warnings, percentile, render_measurements from .runner import Report, parse_options from .worker import file_sha256, measure_async, measure_sync from .workloads import JSON_OBJECT, JSON_PAGES, ocr_workload, padded_pdf @@ -188,6 +188,7 @@ def test_cli_runs_both_backends_and_exports_measurements(tmp_path: Path, capsys: "--benchmark-arg=--profile=small", "--benchmark-arg=--route=aocr", "--benchmark-arg=--iterations=3", + "--benchmark-arg=--min-time=0", "--benchmark-arg=--warmup=1", "--benchmark-arg=--repeats=2", f"--benchmark-arg=--output={output}", @@ -203,6 +204,9 @@ def test_cli_runs_both_backends_and_exports_measurements(tmp_path: Path, capsys: (1, "rust"), (1, "python"), ) + assert report.schema_version == 2 + assert report.warnings == measurement_warnings(report.measurements) + assert any("batch shorter than 1 second" in warning for warning in report.warnings) assert len({value.ready.response_digest for value in report.measurements}) == 1 for value in report.measurements: assert len(value.timing.latency_ms) == 3 @@ -219,7 +223,18 @@ def test_cli_runs_both_backends_and_exports_measurements(tmp_path: Path, capsys: @pytest.mark.parametrize( - "argument", ("--iterations=0", "--warmup=invalid", "--route=chat", "--unknown=1", "--timeout=nan", "--timeout=inf") + "argument", + ( + "--iterations=0", + "--warmup=invalid", + "--route=chat", + "--unknown=1", + "--timeout=nan", + "--timeout=inf", + "--min-time=-1", + "--min-time=nan", + "--min-time=inf", + ), ) def test_cli_rejects_invalid_benchmark_options_without_a_traceback( argument: str, capsys: pytest.CaptureFixture[str] @@ -266,3 +281,94 @@ def test_native_provenance_hash_uses_bounded_memory(tmp_path: Path) -> None: assert peak < 1024 * 1024 finally: tracemalloc.stop() + + +@pytest.mark.parametrize("route", ("ocr", "aocr")) +def test_duration_sampling_meets_time_and_iteration_minimums(route: Route) -> None: + request: Final = invocation().model_copy(update={"min_time": 0.05}) + + def sync_call() -> OCRResponse: + sleep(0.005) + return OCRResponse(model="benchmark", pages=[]) + + async def async_call() -> OCRResponse: + await asyncio.sleep(0.005) + return OCRResponse(model="benchmark", pages=[]) + + result: Final = ( + measure_sync(sync_call, request) if route == "ocr" else asyncio.run(measure_async(async_call, request)) + ) + assert len(result.latency_ms) > request.iterations + assert result.elapsed_ms >= 50 + count_limited: Final = request.model_copy(update={"min_time": 0.001}) + short: Final = ( + measure_sync(sync_call, count_limited) + if route == "ocr" + else asyncio.run(measure_async(async_call, count_limited)) + ) + assert len(short.latency_ms) == request.iterations + + +def measurement(backend: Backend, repeat: int, samples: tuple[float, ...]) -> Measurement: + return Measurement( + backend=backend, + repeat=repeat, + profile="small", + route="ocr", + document_bytes=32768, + response_bytes=100, + response_pages=1, + fixture_sha256="fixture", + ready=Ready(response_digest="response", python_version="3.12", native_sha256=None), + timing=Timing(latency_ms=samples, elapsed_ms=sum(samples), cpu_ms=1), + memory=Memory(baseline_rss_bytes=100, sampled_peak_rss_bytes=110, retained_rss_bytes=105, samples=2), + ) + + +def test_report_exposes_tail_effects_and_insufficient_repetition() -> None: + pair: Final = ( + measurement("python", 0, (1, 1, 1, 97)), + measurement("rust", 0, (1, 1, 1, 1)), + ) + output: Final = render_measurements(pair) + assert "25.000/41.569" in output + assert "40.0" in output + assert "1000.0" in output + assert "Per-repeat measurements" in output + assert "one process repeat cannot establish repeatability" in output + assert "backend order is unbalanced" in output + assert "batch shorter than 1 second" in output + + +def test_warnings_detect_between_process_variation_without_discarding_samples() -> None: + stable: Final = tuple(measurement("python", repeat, (500, 500)) for repeat in range(4)) + assert measurement_warnings(stable) == () + varied: Final = (*stable[:3], measurement("python", 3, (1000, 1000))) + assert measurement_warnings(varied) == ( + "ocr/small/python: repeat p50 range exceeds 10% of its median; investigate variability", + ) + + +def test_codspeed_cli_runs_both_backends_and_entrypoints() -> None: + pytest.importorskip("pytest_codspeed") + result: Final = subprocess.run( + ( + sys.executable, + "-m", + "tests.rust-python-harness.strategies.e2e_benchmark.codspeed", + "--profile", + "small", + "--max-time", + "0.1", + ), + cwd=REPO_ROOT, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stdout + result.stderr + for backend in ("python", "rust"): + for route in ("ocr", "aocr"): + assert f"test_sdk[{backend}-{route}-small]" in result.stdout + assert result.stdout.count("1 benchmarked") == 4 + assert "was never awaited" not in result.stderr + result.stdout diff --git a/tests/rust-python-harness/strategies/e2e_benchmark/worker.py b/tests/rust-python-harness/strategies/e2e_benchmark/worker.py index 201a2c6339e..1b1fd1a5a1c 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/worker.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/worker.py @@ -6,7 +6,8 @@ import hashlib import json import platform import sys -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Iterator +from itertools import count from pathlib import Path from time import perf_counter_ns, process_time_ns from typing import Final, cast @@ -24,7 +25,7 @@ def file_sha256(path: Path) -> str: return digest.hexdigest() -def _ready(response: OCRResponse) -> Ready: +def capture_ready(response: OCRResponse) -> Ready: from litellm.rust_bridge import get_native_bridge from litellm.rust_bridge.configuration import rust_enabled @@ -70,6 +71,14 @@ async def _async_sample(call: Callable[[], Awaitable[OCRResponse]]) -> float: return (perf_counter_ns() - start) / 1e6 +def sample_indices(invocation: Invocation, start_ns: int) -> Iterator[int]: + deadline: Final = start_ns + invocation.min_time * 1e9 + for index in count(): + if index >= invocation.iterations and perf_counter_ns() >= deadline: + return + yield index + + def measure_sync(call: Callable[[], OCRResponse], invocation: Invocation) -> Timing: if invocation.phase == "memory": for _ in range(invocation.iterations): @@ -77,7 +86,7 @@ def measure_sync(call: Callable[[], OCRResponse], invocation: Invocation) -> Tim return Timing(latency_ms=(), cpu_ms=0, elapsed_ms=0) cpu_start: Final = process_time_ns() wall_start: Final = perf_counter_ns() - samples: Final = tuple(_sync_sample(call) for _ in range(invocation.iterations)) + samples: Final = tuple(_sync_sample(call) for _ in sample_indices(invocation, wall_start)) elapsed: Final = perf_counter_ns() - wall_start return Timing(latency_ms=samples, cpu_ms=(process_time_ns() - cpu_start) / 1e6, elapsed_ms=elapsed / 1e6) @@ -89,7 +98,7 @@ async def measure_async(call: Callable[[], Awaitable[OCRResponse]], invocation: return Timing(latency_ms=(), cpu_ms=0, elapsed_ms=0) cpu_start: Final = process_time_ns() wall_start: Final = perf_counter_ns() - samples: Final = tuple([await _async_sample(call) for _ in range(invocation.iterations)]) + samples: Final = tuple([await _async_sample(call) for _ in sample_indices(invocation, wall_start)]) elapsed: Final = perf_counter_ns() - wall_start return Timing(latency_ms=samples, cpu_ms=(process_time_ns() - cpu_start) / 1e6, elapsed_ms=elapsed / 1e6) @@ -97,7 +106,7 @@ async def measure_async(call: Callable[[], Awaitable[OCRResponse]], invocation: async def _run_async(call: Callable[[], Awaitable[OCRResponse]], invocation: Invocation, directory: Path) -> None: for _ in range(invocation.warmup): await call() - ready: Final = _ready(await call()) + ready: Final = capture_ready(await call()) _handshake(ready, directory) _finish(await measure_async(call, invocation), directory) @@ -121,7 +130,7 @@ def run_worker(invocation: Invocation, directory: Path) -> None: call: Final = lambda: sync_route(**kwargs) for _ in range(invocation.warmup): call() - ready: Final = _ready(call()) + ready: Final = capture_ready(call()) _handshake(ready, directory) _finish(measure_sync(call, invocation), directory) From abf159911ff2ae6f9d656e930c3b741c9e96d691 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 17:24:05 -0700 Subject: [PATCH 08/10] fix(ci): scope CodSpeed OIDC permissions to each job --- .github/workflows/codspeed.yml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index dd76609e26e..5e0d7c88531 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -35,7 +35,6 @@ on: permissions: contents: read - id-token: write concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} @@ -43,6 +42,9 @@ concurrency: jobs: benchmarks: + permissions: + contents: read + id-token: write runs-on: ubuntu-24.04 timeout-minutes: 60 @@ -98,6 +100,9 @@ jobs: --codspeed e2e-walltime: + permissions: + contents: read + id-token: write runs-on: codspeed-macro timeout-minutes: 30 From 1006cad1195462ff9a37b67560917dade5437c88 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 17:28:51 -0700 Subject: [PATCH 09/10] test: escape literal filename in benchmark timeout assertion --- .../strategies/e2e_benchmark/test_benchmark.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 e62ceaefd95..84a3b0b541f 100644 --- a/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py +++ b/tests/rust-python-harness/strategies/e2e_benchmark/test_benchmark.py @@ -124,7 +124,7 @@ def test_worker_deadline_terminates_and_reaps_the_process(tmp_path: Path, sample existing_children: Final = frozenset(process.pid for process in psutil.Process().children()) start: Final = monotonic() with (tmp_path / "worker.log").open("w+") as log: - with pytest.raises(TimeoutError, match="timed out waiting for missing.json"): + with pytest.raises(TimeoutError, match=r"timed out waiting for missing\.json"): with sdk_process(case_file, "python", REPO_ROOT, log) as child: tuple( wait_for_output( From e817cb2d222320ebb276401926c63f3af8126e00 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 17:30:41 -0700 Subject: [PATCH 10/10] fix(ci): install OCR fixture dependencies for CodSpeed --- .github/workflows/codspeed.yml | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 5e0d7c88531..7fe6a1edf9f 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -136,9 +136,12 @@ jobs: --with pytest==8.3.5 --with pytest-codspeed==5.0.3 --with psutil==7.2.2 + --with vcrpy==8.2.1 + --with hypothesis==6.165.10 + --with reportlab==5.0.1 --with "mcp>=1.26.0,<2.0" --with "a2a-sdk>=1.1.0,<2.0" - pytest -p pytest_codspeed.plugin + pytest -p pytest_codspeed.plugin -c /dev/null --rootdir=. --import-mode=importlib -o consider_namespace_packages=true tests/rust-python-harness/strategies/e2e_benchmark/codspeed.py --collect-only -q @@ -155,6 +158,9 @@ jobs: --with pytest==8.3.5 --with pytest-codspeed==5.0.3 --with psutil==7.2.2 + --with vcrpy==8.2.1 + --with hypothesis==6.165.10 + --with reportlab==5.0.1 --with "mcp>=1.26.0,<2.0" --with "a2a-sdk>=1.1.0,<2.0" python -m tests.rust-python-harness.strategies.e2e_benchmark.codspeed