mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
feat(harness): benchmark Python and Rust SDK resource usage
This commit is contained in:
parent
e6705510f8
commit
e73d21508d
15 changed files with 1012 additions and 2 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
Measures Python and Rust SDK latency, CPU time, and process memory against deterministic local provider replays, outside correctness-check overhead
|
||||
87
tests/rust-python-harness/strategies/e2e_benchmark/README.md
Normal file
87
tests/rust-python-harness/strategies/e2e_benchmark/README.md
Normal 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
|
||||
```
|
||||
|
|
@ -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",
|
||||
),
|
||||
)
|
||||
153
tests/rust-python-harness/strategies/e2e_benchmark/execution.py
Normal file
153
tests/rust-python-harness/strategies/e2e_benchmark/execution.py
Normal 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,
|
||||
)
|
||||
69
tests/rust-python-harness/strategies/e2e_benchmark/models.py
Normal file
69
tests/rust-python-harness/strategies/e2e_benchmark/models.py
Normal 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
|
||||
|
|
@ -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()
|
||||
|
|
@ -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 ())),
|
||||
),
|
||||
)
|
||||
133
tests/rust-python-harness/strategies/e2e_benchmark/runner.py
Normal file
133
tests/rust-python-harness/strategies/e2e_benchmark/runner.py
Normal 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
|
||||
|
|
@ -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
|
||||
120
tests/rust-python-harness/strategies/e2e_benchmark/worker.py
Normal file
120
tests/rust-python-harness/strategies/e2e_benchmark/worker.py
Normal 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()))
|
||||
|
|
@ -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
2
uv.lock
generated
|
|
@ -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" },
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue