feat(harness): benchmark Python and Rust SDK resource usage

This commit is contained in:
Yujong Lee 2026-09-05 11:51:01 -07:00
parent e6705510f8
commit e73d21508d
15 changed files with 1012 additions and 2 deletions

View file

@ -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",

View file

@ -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

View file

@ -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(

View file

@ -0,0 +1 @@
Measures Python and Rust SDK latency, CPU time, and process memory against deterministic local provider replays, outside correctness-check overhead

View file

@ -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
```

View file

@ -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",
),
)

View file

@ -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,
)

View file

@ -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

View file

@ -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()

View file

@ -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 ())),
),
)

View file

@ -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

View file

@ -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

View file

@ -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()))

View file

@ -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(),
)

2
uv.lock generated
View file

@ -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" },