mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge e817cb2d22 into e88ba94da5
This commit is contained in:
yujonglee
GitHub
commit
2fa02021f2
17 changed files with 1523 additions and 4 deletions
74
.github/workflows/codspeed.yml
vendored
74
.github/workflows/codspeed.yml
vendored
|
|
@ -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"
|
||||
|
|
@ -31,7 +35,6 @@ on:
|
|||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
|
|
@ -39,6 +42,9 @@ concurrency:
|
|||
|
||||
jobs:
|
||||
benchmarks:
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
runs-on: ubuntu-24.04
|
||||
timeout-minutes: 60
|
||||
|
||||
|
|
@ -92,3 +98,69 @@ jobs:
|
|||
-p pytest_codspeed.plugin
|
||||
tests/benchmarks/
|
||||
--codspeed
|
||||
|
||||
e2e-walltime:
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
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 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 -c /dev/null --rootdir=.
|
||||
--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 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
|
||||
|
|
|
|||
|
|
@ -215,6 +215,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",
|
||||
|
|
|
|||
|
|
@ -58,11 +58,12 @@ 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 <strategy>|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
|
||||
- `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(
|
||||
|
|
|
|||
50
tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md
Normal file
50
tests/rust-python-harness/strategies/e2e_benchmark/AGENTS.md
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
# 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. 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 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
|
||||
|
||||
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 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`:
|
||||
|
||||
```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=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
|
||||
|
|
@ -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 AGENTS.md",
|
||||
),
|
||||
)
|
||||
161
tests/rust-python-harness/strategies/e2e_benchmark/codspeed.py
Normal file
161
tests/rust-python-harness/strategies/e2e_benchmark/codspeed.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
from typing import Final
|
||||
|
||||
PYTHON_SENTINEL: Final = "litellm-benchmark-python"
|
||||
161
tests/rust-python-harness/strategies/e2e_benchmark/execution.py
Normal file
161
tests/rust-python-harness/strategies/e2e_benchmark/execution.py
Normal file
|
|
@ -0,0 +1,161 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from collections.abc import Generator, Iterator
|
||||
from contextlib import contextmanager, suppress
|
||||
from pathlib import Path
|
||||
from time import monotonic, sleep
|
||||
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 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
|
||||
|
||||
|
||||
@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=log,
|
||||
stderr=log,
|
||||
text=True,
|
||||
)
|
||||
try:
|
||||
yield process
|
||||
except BaseException:
|
||||
process.terminate()
|
||||
raise
|
||||
finally:
|
||||
if process.stdin is not None:
|
||||
with suppress(BrokenPipeError):
|
||||
process.stdin.close()
|
||||
try:
|
||||
process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.wait(timeout=5)
|
||||
|
||||
|
||||
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(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(
|
||||
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())
|
||||
ready_file: Final = directory / "ready.json"
|
||||
timing_file: Final = directory / "timing.json"
|
||||
with (directory / "worker.log").open("w+") as log:
|
||||
try:
|
||||
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)
|
||||
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),
|
||||
),
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
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",
|
||||
min_time=options.min_time,
|
||||
)
|
||||
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,
|
||||
)
|
||||
70
tests/rust-python-harness/strategies/e2e_benchmark/models.py
Normal file
70
tests/rust-python-harness/strategies/e2e_benchmark/models.py
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import 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"]
|
||||
|
||||
|
||||
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=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)
|
||||
sample_interval_ms: float = Field(default=5, ge=1, allow_inf_nan=False)
|
||||
output: str | None = None
|
||||
|
||||
|
||||
class Invocation(BenchmarkModel):
|
||||
model: str
|
||||
document_url: str
|
||||
route: Route
|
||||
provider_url: str
|
||||
iterations: int
|
||||
warmup: int
|
||||
phase: Phase
|
||||
min_time: float = Field(default=0, ge=0, allow_inf_nan=False)
|
||||
|
||||
|
||||
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,70 @@
|
|||
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 .constants import PYTHON_SENTINEL
|
||||
from .models import Backend
|
||||
|
||||
|
||||
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(timeout=5)
|
||||
process.close()
|
||||
124
tests/rust-python-harness/strategies/e2e_benchmark/reporting.py
Normal file
124
tests/rust-python-harness/strategies/e2e_benchmark/reporting.py
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
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 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 | mean/stdev ms | CPU ms/call | calls/s | "
|
||||
"RSS baseline/peak/after MiB | pooled p50 ratio"
|
||||
)
|
||||
|
||||
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} | "
|
||||
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"
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
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, ...]:
|
||||
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 ())),
|
||||
),
|
||||
)
|
||||
156
tests/rust-python-harness/strategies/e2e_benchmark/runner.py
Normal file
156
tests/rust-python-harness/strategies/e2e_benchmark/runner.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import platform
|
||||
import subprocess
|
||||
from collections.abc import Sequence
|
||||
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
|
||||
from .models import Backend, BenchmarkModel, Measurement, Options, Profile, Route
|
||||
from .reporting import ARTIFACT_KIND, MEASUREMENTS, measurements, measurement_warnings
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .workloads import Workload
|
||||
|
||||
|
||||
class Report(BenchmarkModel):
|
||||
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:
|
||||
defaults: Final = Options()
|
||||
command: Final = click.Command(
|
||||
"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(
|
||||
("--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(
|
||||
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),
|
||||
warnings=measurement_warnings(measurements(tuple(run.results.values()))),
|
||||
)
|
||||
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
|
||||
|
|
@ -0,0 +1,374 @@
|
|||
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
|
||||
|
||||
import click
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
|
||||
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 Backend, Invocation, Measurement, Memory, Options, Ready, Route, Timing
|
||||
from .provider import provider_process
|
||||
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
|
||||
|
||||
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"startxref")[0])
|
||||
assert padded.endswith(b"\nstartxref\n9\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(click.BadParameter):
|
||||
parse_options(arguments)
|
||||
|
||||
|
||||
def test_unknown_options_are_not_silently_ignored() -> None:
|
||||
with pytest.raises(click.NoSuchOption, match="No such option"):
|
||||
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 == ()
|
||||
|
||||
|
||||
@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=r"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:
|
||||
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
|
||||
|
||||
|
||||
@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, "route": route}
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="backend mismatch"):
|
||||
execute_phase(request, "python", Options(iterations=3, warmup=1), REPO_ROOT)
|
||||
|
||||
|
||||
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: Final = main(
|
||||
(
|
||||
"run",
|
||||
"e2e_benchmark",
|
||||
"--surface",
|
||||
"sdk",
|
||||
"--function",
|
||||
"ocr",
|
||||
"--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}",
|
||||
)
|
||||
)
|
||||
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 tuple((value.repeat, value.backend) for value in report.measurements) == (
|
||||
(0, "python"),
|
||||
(0, "rust"),
|
||||
(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
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"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]
|
||||
) -> 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 == ()
|
||||
|
||||
|
||||
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()
|
||||
|
||||
|
||||
@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
|
||||
140
tests/rust-python-harness/strategies/e2e_benchmark/worker.py
Normal file
140
tests/rust-python-harness/strategies/e2e_benchmark/worker.py
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import gc
|
||||
import hashlib
|
||||
import json
|
||||
import platform
|
||||
import sys
|
||||
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
|
||||
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
|
||||
from .models import Invocation, Ready, Timing
|
||||
|
||||
|
||||
def file_sha256(path: Path) -> str:
|
||||
digest: Final = hashlib.sha256()
|
||||
with path.open("rb") as source:
|
||||
for chunk in iter(lambda: source.read(256 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def capture_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=file_sha256(Path(native_path)) if native_path else None,
|
||||
)
|
||||
|
||||
|
||||
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, directory: Path) -> None:
|
||||
gc.collect()
|
||||
_publish(ready, directory / "ready.json")
|
||||
if sys.stdin.readline().strip() != "go":
|
||||
raise RuntimeError("benchmark controller disconnected before measurement")
|
||||
|
||||
|
||||
def _finish(timing: Timing, directory: Path) -> None:
|
||||
gc.collect()
|
||||
_publish(timing, directory / "timing.json")
|
||||
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 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):
|
||||
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 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)
|
||||
|
||||
|
||||
async def measure_async(call: Callable[[], Awaitable[OCRResponse]], invocation: Invocation) -> Timing:
|
||||
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 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)
|
||||
|
||||
|
||||
async def _run_async(call: Callable[[], Awaitable[OCRResponse]], invocation: Invocation, directory: Path) -> None:
|
||||
for _ in range(invocation.warmup):
|
||||
await call()
|
||||
ready: Final = capture_ready(await call())
|
||||
_handshake(ready, directory)
|
||||
_finish(await measure_async(call, invocation), directory)
|
||||
|
||||
|
||||
def run_worker(invocation: Invocation, directory: Path) -> 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, directory))
|
||||
return
|
||||
sync_route: Final = cast(Callable[..., OCRResponse], litellm.ocr)
|
||||
call: Final = lambda: sync_route(**kwargs)
|
||||
for _ in range(invocation.warmup):
|
||||
call()
|
||||
ready: Final = capture_ready(call())
|
||||
_handshake(ready, directory)
|
||||
_finish(measure_sync(call, invocation), directory)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
case_file: Final = Path(sys.argv[1])
|
||||
run_worker(Invocation.model_validate_json(case_file.read_bytes()), case_file.parent)
|
||||
|
|
@ -0,0 +1,87 @@
|
|||
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"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
|
||||
|
||||
|
||||
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
|
|
@ -4531,6 +4531,7 @@ dev = [
|
|||
{ name = "opentelemetry-instrumentation-fastapi" },
|
||||
{ name = "opentelemetry-sdk" },
|
||||
{ name = "parameterized" },
|
||||
{ name = "psutil" },
|
||||
{ name = "psycopg" },
|
||||
{ name = "psycopg-binary" },
|
||||
{ name = "pytest" },
|
||||
|
|
@ -4720,6 +4721,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