chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,145 @@
|
||||
# MLflow AI Gateway Benchmark
|
||||
|
||||
Measures the **proxy overhead** of the MLflow tracking-server-backed AI Gateway under
|
||||
concurrent load. A fake OpenAI server simulates the upstream provider at a fixed latency,
|
||||
so results reflect pure MLflow processing time rather than provider variance.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Python 3.10+ with [`uv`](https://docs.astral.sh/uv/) — all scripts must be run via `uv run`, which handles dependency installation automatically via inline script metadata
|
||||
- Docker (required for `--database postgres` and `multi` mode)
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
cd dev/benchmarks/gateway
|
||||
|
||||
# 4 instances behind nginx (default, requires Docker)
|
||||
uv run run.py
|
||||
|
||||
# Single instance, SQLite (no Docker needed)
|
||||
uv run run.py --instances 1
|
||||
|
||||
# Single instance, PostgreSQL
|
||||
uv run run.py --instances 1 --database postgres
|
||||
|
||||
# Scale up
|
||||
uv run run.py --instances 8 --workers 8
|
||||
|
||||
# Benchmark an existing endpoint directly (skips all setup)
|
||||
uv run run.py --url http://your-server/gateway/my-endpoint/mlflow/invocations
|
||||
|
||||
# Basic-auth enabled (starts MLflow with --app-name=basic-auth,
|
||||
# sends Authorization: Basic on every request)
|
||||
uv run run.py --instances 1 --auth
|
||||
|
||||
```
|
||||
|
||||
## What is measured
|
||||
|
||||
Latency is measured **client-side** using `time.perf_counter()` around each `aiohttp` request.
|
||||
Each sample covers the full round-trip: client serialization → loopback → full server processing → response deserialization. Only HTTP 200 responses count toward latency stats; errors are tracked separately.
|
||||
|
||||
Connection pooling and HTTP keep-alive are enabled, so TCP handshake cost is amortized after the warmup phase.
|
||||
|
||||
### What is NOT measured
|
||||
|
||||
| Factor | In this benchmark | In production |
|
||||
| ------------------ | ---------------------------------------------- | --------------------------- |
|
||||
| Network latency | ~0 ms (loopback) | 1–100 ms per hop |
|
||||
| TLS/SSL | None (plain HTTP) | ~5–20 ms per new connection |
|
||||
| Provider inference | Fixed fake delay (`--fake-delay-ms`) | Variable (50 ms – 60 s+) |
|
||||
| Authentication | Off by default; basic-auth opt-in via `--auth` | Token validation, RBAC |
|
||||
|
||||
## What MLflow does per request
|
||||
|
||||
Each invocation through the tracking-server gateway runs these steps:
|
||||
|
||||
```
|
||||
1. Config resolution (DB-backed, cached after first hit)
|
||||
2. Secret decryption (cached, 60 s TTL)
|
||||
3. Provider instantiation
|
||||
4. Tracing (if usage_tracking=True)
|
||||
5. HTTP call to LLM API
|
||||
```
|
||||
|
||||
Steps 1 (config resolution) and 4 (tracing) have historically been the dominant bottlenecks.
|
||||
Config caching (enabled by default) eliminates most of step 1's cost. Tracing overhead
|
||||
depends on the span processor in use.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Single instance (`--instances 1`)
|
||||
|
||||
```
|
||||
benchmark.py ──aiohttp──▶ MLflow server (:5731) ──▶ fake_server.py (:9137)
|
||||
│
|
||||
SQLite or PostgreSQL
|
||||
```
|
||||
|
||||
### Multi-instance (`--instances N`, default)
|
||||
|
||||
```
|
||||
benchmark.py ──aiohttp──▶ nginx LB (:5731) ──round-robin──▶ MLflow :5800
|
||||
MLflow :5801
|
||||
MLflow :580N
|
||||
│
|
||||
fake_server.py (:9137)
|
||||
PostgreSQL (Docker)
|
||||
```
|
||||
|
||||
MLflow instances are started **sequentially** (instance 0 first) to let it initialize the
|
||||
DB schema before the others join. All instances share one PostgreSQL database.
|
||||
|
||||
## Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
| ----------------------------- | -------------- | ----------------------------------------------------------------------- |
|
||||
| `--url URL` | — | Benchmark this URL directly, skip all setup |
|
||||
| `--instances N` | 4 | MLflow instances. Use 1 for single-instance (no nginx, optional SQLite) |
|
||||
| `--workers N` | 4 | MLflow worker processes per instance |
|
||||
| `--database sqlite\|postgres` | `sqlite` | Database to use — only applies when `--instances 1` |
|
||||
| `--no-usage-tracking` | — | Disable usage tracking (tracing) on the endpoint |
|
||||
| `--port N` | 5731 | Port to benchmark (MLflow port for single, nginx LB port for multi) |
|
||||
| `--base-port N` | 5800 | First MLflow instance port in multi mode (rest are +1, +2, …) |
|
||||
| `--fake-server-port N` | 9137 | Fake OpenAI server port |
|
||||
| `--requests N` | 2000 | Requests per run |
|
||||
| `--max-concurrent N` | 50 | Max concurrent requests |
|
||||
| `--runs N` | 3 | Number of benchmark runs |
|
||||
| `--fake-delay-ms N` | 50 | Simulated provider latency in ms |
|
||||
| `--min-rps N` | — | Fail (exit 1) if average throughput falls below N req/s |
|
||||
| `--max-p50-ms N` | — | Fail (exit 1) if average P50 latency exceeds N ms (CI threshold) |
|
||||
| `--max-p99-ms N` | — | Fail (exit 1) if average P99 latency exceeds N ms (CI threshold) |
|
||||
| `--auth` | off | Start MLflow with `--app-name=basic-auth`; send Basic auth on requests |
|
||||
| `--auth-username USER` | `admin` | Basic-auth username (matches `mlflow/server/auth/basic_auth.ini`) |
|
||||
| `--auth-password PASS` | `password1234` | Basic-auth password (matches `mlflow/server/auth/basic_auth.ini`) |
|
||||
|
||||
All flags can also be set via environment variables (same name, uppercased):
|
||||
`INSTANCES`, `WORKERS_PER_INSTANCE`, `REQUESTS`, `MAX_CONCURRENT`, `RUNS`,
|
||||
`FAKE_RESPONSE_DELAY_MS`, `MLFLOW_PORT`, `BASE_PORT`, `FAKE_SERVER_PORT`,
|
||||
`AUTH`, `AUTH_USERNAME`, `AUTH_PASSWORD`.
|
||||
|
||||
To avoid conflicts with a local PostgreSQL instance, override the port via `GATEWAY_BENCH_POSTGRES_PORT` (default: 5432).
|
||||
|
||||
## Known limitations
|
||||
|
||||
- **Loopback only** — all processes run on the same machine. Results don't include real
|
||||
network latency between client, gateway, and provider.
|
||||
- **No TLS** — MLflow is started with `--disable-security-middleware`. Production deployments
|
||||
add TLS termination overhead.
|
||||
- **Fixed provider latency** — `fake_server.py` always responds in exactly `--fake-delay-ms`.
|
||||
Real providers have high variance (P99 often 5–10× P50).
|
||||
- **Basic-auth is opt-in, no RBAC** — `--auth` enables `basic-auth` with the default
|
||||
admin user (full permissions), which measures the cost of HTTP Basic authentication
|
||||
and user lookup but not fine-grained RBAC checks against non-admin users.
|
||||
- **Single machine resource contention** — with multiple instances, all MLflow instances, nginx,
|
||||
PostgreSQL, and the benchmark client share CPU/memory. On a server with dedicated resources
|
||||
per instance, throughput will be higher.
|
||||
|
||||
## Files
|
||||
|
||||
| File | Purpose |
|
||||
| ---------------- | ------------------------------------------------------------------------------ |
|
||||
| `run.py` | Main entry point — orchestrates servers, Docker, endpoint setup, and benchmark |
|
||||
| `benchmark.py` | Async HTTP benchmark client (standalone or imported by `run.py`) |
|
||||
| `fake_server.py` | Fake OpenAI-compatible server for controlled latency simulation |
|
||||
@@ -0,0 +1,362 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.10"
|
||||
# dependencies = ["aiohttp>=3.13.3,<4", "rich>=14.3.3,<15"]
|
||||
# ///
|
||||
"""Async HTTP benchmark client for the MLflow AI Gateway.
|
||||
|
||||
Can be imported by run.py or used standalone:
|
||||
uv run benchmark.py --url http://127.0.0.1:5731/gateway/benchmark-chat/mlflow/invocations
|
||||
uv run benchmark.py --url http://... --requests 5000 --max-concurrent 100 --runs 3
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import math
|
||||
import statistics
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
from rich.console import Console # type: ignore[import-not-found]
|
||||
from rich.progress import ( # type: ignore[import-not-found]
|
||||
BarColumn,
|
||||
MofNCompleteColumn,
|
||||
Progress,
|
||||
SpinnerColumn,
|
||||
TaskID,
|
||||
TextColumn,
|
||||
TimeElapsedColumn,
|
||||
)
|
||||
from rich.table import Table # type: ignore[import-not-found]
|
||||
|
||||
console = Console()
|
||||
|
||||
_BODY = {
|
||||
"messages": [{"role": "user", "content": "benchmark request"}],
|
||||
"temperature": 0.0,
|
||||
"max_tokens": 50,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class RunResult:
|
||||
latencies_ms: list[float] = field(default_factory=list)
|
||||
failures: dict[str, int] = field(default_factory=dict)
|
||||
wall_time: float = 0.0
|
||||
|
||||
@property
|
||||
def n_success(self) -> int:
|
||||
return len(self.latencies_ms)
|
||||
|
||||
@property
|
||||
def n_failures(self) -> int:
|
||||
return sum(self.failures.values())
|
||||
|
||||
@property
|
||||
def throughput(self) -> float:
|
||||
return self.n_success / self.wall_time if self.wall_time > 0 else 0.0
|
||||
|
||||
def percentile(self, p: float) -> float:
|
||||
if not self.latencies_ms:
|
||||
return 0.0
|
||||
s = sorted(self.latencies_ms)
|
||||
idx = max(0, math.ceil(p / 100 * len(s)) - 1)
|
||||
return s[idx]
|
||||
|
||||
|
||||
async def _send(
|
||||
session: aiohttp.ClientSession,
|
||||
url: str,
|
||||
sem: asyncio.Semaphore,
|
||||
auth: aiohttp.BasicAuth | None = None,
|
||||
) -> tuple[float, str | None]:
|
||||
async with sem:
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
async with session.post(url, json=_BODY, auth=auth) as resp:
|
||||
await resp.read()
|
||||
ms = (time.perf_counter() - t0) * 1000
|
||||
if resp.status == 200:
|
||||
return ms, None
|
||||
return ms, f"HTTP {resp.status}"
|
||||
except Exception as e:
|
||||
return (time.perf_counter() - t0) * 1000, type(e).__name__
|
||||
|
||||
|
||||
async def _run_once(
|
||||
url: str,
|
||||
n: int,
|
||||
max_concurrent: int,
|
||||
progress: Progress,
|
||||
task_id: TaskID,
|
||||
auth: aiohttp.BasicAuth | None = None,
|
||||
) -> RunResult:
|
||||
sem = asyncio.Semaphore(max_concurrent)
|
||||
connector = aiohttp.TCPConnector(
|
||||
limit=max(max_concurrent * 2, 200),
|
||||
limit_per_host=max(max_concurrent, 200),
|
||||
force_close=False,
|
||||
enable_cleanup_closed=True,
|
||||
)
|
||||
result = RunResult()
|
||||
total_time = 0.0
|
||||
max_time = 0.0
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
t0 = time.perf_counter()
|
||||
for coro in asyncio.as_completed([_send(session, url, sem, auth) for _ in range(n)]):
|
||||
ms, error = await coro
|
||||
if error:
|
||||
result.failures[error] = result.failures.get(error, 0) + 1
|
||||
else:
|
||||
result.latencies_ms.append(ms)
|
||||
total_time += ms
|
||||
if ms > max_time:
|
||||
max_time = ms
|
||||
|
||||
n_ok = result.n_success
|
||||
n_fail = result.n_failures
|
||||
mean = total_time / n_ok if n_ok else 0.0
|
||||
fail_part = f"[red]✗{n_fail}[/red] " if n_fail else ""
|
||||
live = f"{fail_part}✓{n_ok} mean={mean:.0f}ms max={max_time:.0f}ms"
|
||||
progress.update(task_id, advance=1, live=live)
|
||||
|
||||
result.wall_time = time.perf_counter() - t0
|
||||
return result
|
||||
|
||||
|
||||
async def _warmup(
|
||||
url: str, n: int, max_concurrent: int, auth: aiohttp.BasicAuth | None = None
|
||||
) -> None:
|
||||
sem = asyncio.Semaphore(max_concurrent)
|
||||
connector = aiohttp.TCPConnector(limit=max(max_concurrent * 2, 200))
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
await asyncio.gather(*[_send(session, url, sem, auth) for _ in range(n)])
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
url: str,
|
||||
n_requests: int = 2000,
|
||||
max_concurrent: int = 50,
|
||||
runs: int = 3,
|
||||
auth: aiohttp.BasicAuth | None = None,
|
||||
) -> list[RunResult]:
|
||||
warmup_n = min(max(50, max_concurrent), n_requests)
|
||||
console.print(f" [dim]Warming up ({warmup_n} requests)...[/dim]")
|
||||
asyncio.run(_warmup(url, warmup_n, max_concurrent, auth))
|
||||
|
||||
results = []
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
MofNCompleteColumn(),
|
||||
TimeElapsedColumn(),
|
||||
TextColumn(" {task.fields[live]}"),
|
||||
console=console,
|
||||
) as progress:
|
||||
for i in range(runs):
|
||||
task_id = progress.add_task(f" Run {i + 1}/{runs}", total=n_requests, live="")
|
||||
results.append(
|
||||
asyncio.run(_run_once(url, n_requests, max_concurrent, progress, task_id, auth))
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def results_to_dict(results: list[RunResult]) -> dict[str, Any]:
|
||||
runs = [
|
||||
{
|
||||
"n_success": r.n_success,
|
||||
"n_failures": r.n_failures,
|
||||
"failures": r.failures,
|
||||
"wall_time_s": r.wall_time,
|
||||
"mean_ms": statistics.mean(r.latencies_ms) if r.latencies_ms else 0.0,
|
||||
"p50_ms": r.percentile(50),
|
||||
"p95_ms": r.percentile(95),
|
||||
"p99_ms": r.percentile(99),
|
||||
"max_ms": max(r.latencies_ms) if r.latencies_ms else 0.0,
|
||||
"rps": r.throughput,
|
||||
}
|
||||
for r in results
|
||||
]
|
||||
summary: dict[str, Any] = (
|
||||
{
|
||||
"avg_mean_ms": statistics.mean(
|
||||
statistics.mean(r.latencies_ms) if r.latencies_ms else 0.0 for r in results
|
||||
),
|
||||
"avg_p50_ms": statistics.mean(r.percentile(50) for r in results),
|
||||
"avg_p99_ms": statistics.mean(r.percentile(99) for r in results),
|
||||
"avg_rps": statistics.mean(r.throughput for r in results),
|
||||
}
|
||||
if results
|
||||
else {}
|
||||
)
|
||||
return {"runs": runs, "summary": summary}
|
||||
|
||||
|
||||
def print_results(results: list[RunResult]) -> None:
|
||||
table = Table(show_header=True, header_style="bold cyan", box=None, padding=(0, 2))
|
||||
table.add_column("Run", style="dim", width=5)
|
||||
table.add_column("Mean ms", justify="right")
|
||||
table.add_column("P50 ms", justify="right")
|
||||
table.add_column("P95 ms", justify="right")
|
||||
table.add_column("P99 ms", justify="right")
|
||||
table.add_column("Max ms", justify="right")
|
||||
table.add_column("Req/s", justify="right")
|
||||
table.add_column("Failures", justify="right")
|
||||
|
||||
means = []
|
||||
p50s = []
|
||||
p95s = []
|
||||
p99s = []
|
||||
maxes = []
|
||||
throughputs = []
|
||||
for i, r in enumerate(results):
|
||||
mean = statistics.mean(r.latencies_ms) if r.latencies_ms else 0.0
|
||||
p50 = r.percentile(50)
|
||||
p95 = r.percentile(95)
|
||||
p99 = r.percentile(99)
|
||||
mx = max(r.latencies_ms) if r.latencies_ms else 0.0
|
||||
means.append(mean)
|
||||
p50s.append(p50)
|
||||
p95s.append(p95)
|
||||
p99s.append(p99)
|
||||
maxes.append(mx)
|
||||
throughputs.append(r.throughput)
|
||||
fail_str = f"[red]{r.n_failures}[/red]" if r.n_failures else "0"
|
||||
table.add_row(
|
||||
str(i + 1),
|
||||
f"{mean:.1f}",
|
||||
f"{p50:.1f}",
|
||||
f"{p95:.1f}",
|
||||
f"{p99:.1f}",
|
||||
f"{mx:.1f}",
|
||||
f"{r.throughput:.0f}",
|
||||
fail_str,
|
||||
)
|
||||
|
||||
if len(results) > 1:
|
||||
table.add_section()
|
||||
table.add_row(
|
||||
"[bold]avg[/bold]",
|
||||
f"[bold]{statistics.mean(means):.1f}[/bold]",
|
||||
f"[bold]{statistics.mean(p50s):.1f}[/bold]",
|
||||
f"[bold]{statistics.mean(p95s):.1f}[/bold]",
|
||||
f"[bold]{statistics.mean(p99s):.1f}[/bold]",
|
||||
f"[bold]{statistics.mean(maxes):.1f}[/bold]",
|
||||
f"[bold]{statistics.mean(throughputs):.0f}[/bold]",
|
||||
"",
|
||||
)
|
||||
|
||||
console.print()
|
||||
console.print(table)
|
||||
|
||||
combined: dict[str, int] = {}
|
||||
for r in results:
|
||||
for k, v in r.failures.items():
|
||||
combined[k] = combined.get(k, 0) + v
|
||||
if combined:
|
||||
console.print()
|
||||
console.print("[red]Failure breakdown:[/red]")
|
||||
for reason, count in sorted(combined.items(), key=lambda x: -x[1]):
|
||||
console.print(f" {reason}: {count}")
|
||||
|
||||
|
||||
def check_thresholds(
|
||||
results: list[RunResult],
|
||||
min_rps: float | None = None,
|
||||
max_p50_ms: float | None = None,
|
||||
max_p99_ms: float | None = None,
|
||||
) -> bool:
|
||||
"""Check results against performance thresholds. Returns True if all pass."""
|
||||
avg_rps = statistics.mean(r.throughput for r in results)
|
||||
avg_p50 = statistics.mean(r.percentile(50) for r in results)
|
||||
avg_p99 = statistics.mean(r.percentile(99) for r in results)
|
||||
passed = True
|
||||
|
||||
if min_rps is not None and avg_rps < min_rps:
|
||||
console.print(
|
||||
f"\n[red]THRESHOLD FAILED:[/red] avg throughput {avg_rps:.0f} req/s"
|
||||
f" < minimum {min_rps:.0f} req/s"
|
||||
)
|
||||
passed = False
|
||||
|
||||
if max_p50_ms is not None and avg_p50 > max_p50_ms:
|
||||
console.print(
|
||||
f"\n[red]THRESHOLD FAILED:[/red] avg P50 {avg_p50:.1f} ms > maximum {max_p50_ms:.1f} ms"
|
||||
)
|
||||
passed = False
|
||||
|
||||
if max_p99_ms is not None and avg_p99 > max_p99_ms:
|
||||
console.print(
|
||||
f"\n[red]THRESHOLD FAILED:[/red] avg P99 {avg_p99:.1f} ms > maximum {max_p99_ms:.1f} ms"
|
||||
)
|
||||
passed = False
|
||||
|
||||
if passed and (min_rps is not None or max_p50_ms is not None or max_p99_ms is not None):
|
||||
console.print("\n[green]All thresholds passed.[/green]")
|
||||
|
||||
return passed
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Async HTTP benchmark client for MLflow Gateway")
|
||||
parser.add_argument("--url", required=True, help="Gateway invocation URL")
|
||||
parser.add_argument("--requests", type=int, default=2000)
|
||||
parser.add_argument("--max-concurrent", type=int, default=50)
|
||||
parser.add_argument("--runs", type=int, default=3)
|
||||
parser.add_argument(
|
||||
"--min-rps",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Fail (exit 1) if average throughput falls below N req/s",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-p50-ms",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Fail (exit 1) if average P50 latency exceeds N ms",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-p99-ms",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Fail (exit 1) if average P99 latency exceeds N ms",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth-username",
|
||||
default=None,
|
||||
help="Basic auth username. If set together with --auth-password, sent on every request.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth-password",
|
||||
default=None,
|
||||
help="Basic auth password. If set together with --auth-username, sent on every request.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
auth = (
|
||||
aiohttp.BasicAuth(args.auth_username, args.auth_password)
|
||||
if args.auth_username and args.auth_password
|
||||
else None
|
||||
)
|
||||
|
||||
console.print(f"\n[bold]Benchmarking[/bold] {args.url}")
|
||||
console.print(
|
||||
f" {args.requests} requests · {args.max_concurrent} concurrent · {args.runs} runs\n"
|
||||
)
|
||||
results = run_benchmark(args.url, args.requests, args.max_concurrent, args.runs, auth)
|
||||
print_results(results)
|
||||
|
||||
if not check_thresholds(
|
||||
results, min_rps=args.min_rps, max_p50_ms=args.max_p50_ms, max_p99_ms=args.max_p99_ms
|
||||
):
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,68 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.10"
|
||||
# dependencies = ["fastapi>=0.115.0,<1", "uvicorn[standard]>=0.30.0,<1"]
|
||||
# ///
|
||||
"""Fake OpenAI-compatible server for benchmarking.
|
||||
|
||||
Returns synthetic responses after a configurable delay so benchmarks measure
|
||||
MLflow overhead rather than provider latency.
|
||||
|
||||
Run standalone:
|
||||
uv run fake_server.py
|
||||
PORT=9200 uv run fake_server.py
|
||||
|
||||
Or with multiple workers (as launched by run.py):
|
||||
uvicorn fake_server:app --workers 8 --port 9137
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
DELAY_MS = int(os.environ.get("FAKE_RESPONSE_DELAY_MS", "50"))
|
||||
|
||||
|
||||
class ChatRequest(BaseModel):
|
||||
model: str = "gpt-4o-mini"
|
||||
messages: list[dict[str, str]] = Field(min_length=1)
|
||||
stream: bool = False
|
||||
temperature: float = 1.0
|
||||
max_tokens: int = 50
|
||||
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def chat_completions(req: ChatRequest) -> dict[str, Any]:
|
||||
await asyncio.sleep(DELAY_MS / 1000)
|
||||
return {
|
||||
"id": "chatcmpl-fake",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": req.model,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "Hello!"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health() -> dict[str, str]:
|
||||
# Polled by run.py's _wait_for_port to detect when the server is ready.
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
port = int(os.environ.get("PORT", "9137"))
|
||||
host = os.environ.get("FAKE_SERVER_HOST", "127.0.0.1")
|
||||
uvicorn.run("fake_server:app", host=host, port=port, log_level="warning")
|
||||
@@ -0,0 +1,793 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.10"
|
||||
# dependencies = ["aiohttp>=3.13.3,<4", "psycopg2-binary>=2.9,<3", "rich>=14.3.3,<15"]
|
||||
# ///
|
||||
"""MLflow AI Gateway benchmark runner.
|
||||
|
||||
Orchestrates fake OpenAI server, MLflow server(s), optional PostgreSQL and
|
||||
nginx (via Docker), then runs the async benchmark client.
|
||||
|
||||
Usage:
|
||||
uv run run.py # 4 instances, PostgreSQL, nginx (Docker)
|
||||
uv run run.py --instances 1 # single instance, SQLite, no Docker
|
||||
uv run run.py --instances 1 --database postgres
|
||||
uv run run.py --instances 8 --workers 8
|
||||
uv run run.py --url http://... # benchmark an existing endpoint directly
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from collections.abc import Generator
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
import aiohttp # type: ignore[import-not-found]
|
||||
import benchmark as bm # local module; path inserted above
|
||||
from rich.console import Console # type: ignore[import-not-found]
|
||||
from rich.panel import Panel # type: ignore[import-not-found]
|
||||
from rich.progress import ( # type: ignore[import-not-found]
|
||||
Progress,
|
||||
SpinnerColumn,
|
||||
TextColumn,
|
||||
TimeElapsedColumn,
|
||||
)
|
||||
|
||||
SCRIPT_DIR = Path(__file__).parent
|
||||
|
||||
FAKE_SERVER_PORT = 9137
|
||||
FAKE_SERVER_WORKERS = 8
|
||||
MLFLOW_PORT = 5731
|
||||
INSTANCE_BASE_PORT = 5800
|
||||
POSTGRES_PORT = int(os.environ.get("GATEWAY_BENCH_POSTGRES_PORT", "5432"))
|
||||
POSTGRES_PASSWORD = "benchmarkpass"
|
||||
ENDPOINT_NAME = "benchmark-chat"
|
||||
|
||||
_API_SECRET_CREATE = "gateway/secrets/create"
|
||||
_API_MODEL_DEF_CREATE = "gateway/model-definitions/create"
|
||||
_API_ENDPOINT_CREATE = "gateway/endpoints/create"
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def _uv_prefix() -> list[str]:
|
||||
"""Return uv run prefix when inside the mlflow repo, else empty list."""
|
||||
in_repo = (
|
||||
shutil.which("uv")
|
||||
and subprocess.run(
|
||||
["git", "rev-parse", "HEAD"], cwd=SCRIPT_DIR, capture_output=True
|
||||
).returncode
|
||||
== 0
|
||||
)
|
||||
return ["uv", "run", "--no-build-isolation", "--extra", "gateway"] if in_repo else []
|
||||
|
||||
|
||||
def _subprocess_env() -> dict[str, str]:
|
||||
return os.environ | {"OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES"}
|
||||
|
||||
|
||||
def _wait_for_port(port: int, label: str, log_file: Path | None = None, timeout: int = 30) -> None:
|
||||
url = f"http://127.0.0.1:{port}/health"
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
TimeElapsedColumn(),
|
||||
console=console,
|
||||
transient=True,
|
||||
) as progress:
|
||||
progress.add_task(f" Waiting for {label}...", total=None)
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
with urllib.request.urlopen(url, timeout=1):
|
||||
break
|
||||
except Exception:
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
console.print(f" [red]✗ {label} failed to start within {timeout}s[/red]")
|
||||
if log_file and log_file.exists():
|
||||
console.print(" [yellow]Last 20 lines of log:[/yellow]")
|
||||
for line in log_file.read_text().splitlines()[-20:]:
|
||||
console.print(f" [dim]{line}[/dim]")
|
||||
sys.exit(1)
|
||||
console.print(f" [green]✓[/green] {label} ready")
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _start_fake_server(
|
||||
work_dir: str, port: int = FAKE_SERVER_PORT, workers: int = FAKE_SERVER_WORKERS
|
||||
) -> Generator[None, None, None]:
|
||||
prefix = _uv_prefix()
|
||||
log_file = Path(work_dir) / "fake_server.log"
|
||||
with (
|
||||
log_file.open("w") as f,
|
||||
subprocess.Popen(
|
||||
[
|
||||
*prefix,
|
||||
"uvicorn",
|
||||
"fake_server:app",
|
||||
"--workers",
|
||||
str(workers),
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(port),
|
||||
"--log-level",
|
||||
"warning",
|
||||
],
|
||||
cwd=SCRIPT_DIR,
|
||||
stdout=f,
|
||||
stderr=f,
|
||||
env=_subprocess_env(),
|
||||
) as proc,
|
||||
):
|
||||
_wait_for_port(port, "fake OpenAI server", log_file)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
proc.terminate()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _start_mlflow(
|
||||
work_dir: str,
|
||||
port: int,
|
||||
workers: int,
|
||||
backend_uri: str,
|
||||
label: str = "MLflow server",
|
||||
host: str = "127.0.0.1",
|
||||
auth: bool = False,
|
||||
) -> Generator[None, None, None]:
|
||||
prefix = _uv_prefix()
|
||||
# basic-auth requires the `auth` extra (Flask-WTF) at runtime.
|
||||
if auth and prefix:
|
||||
prefix = [*prefix, "--extra", "auth"]
|
||||
# psycopg2-binary lives in the `db` extra.
|
||||
if backend_uri.startswith("postgresql") and prefix:
|
||||
prefix = [*prefix, "--extra", "db"]
|
||||
log_file = Path(work_dir) / f"mlflow-{port}.log"
|
||||
cmd = [
|
||||
*prefix,
|
||||
"mlflow",
|
||||
"server",
|
||||
"--backend-store-uri",
|
||||
backend_uri,
|
||||
"--host",
|
||||
host,
|
||||
"--port",
|
||||
str(port),
|
||||
"--workers",
|
||||
str(workers),
|
||||
"--disable-security-middleware",
|
||||
]
|
||||
if auth:
|
||||
cmd += ["--app-name", "basic-auth"]
|
||||
with (
|
||||
log_file.open("w") as f,
|
||||
subprocess.Popen(cmd, cwd=SCRIPT_DIR, stdout=f, stderr=f, env=_subprocess_env()) as proc,
|
||||
):
|
||||
_wait_for_port(port, label, log_file)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
proc.terminate()
|
||||
|
||||
|
||||
def _check_docker() -> None:
|
||||
try:
|
||||
result = subprocess.run(["docker", "info"], capture_output=True)
|
||||
except FileNotFoundError:
|
||||
console.print(
|
||||
"[red]Docker is not installed. Install it at https://docs.docker.com/get-docker/[/red]"
|
||||
)
|
||||
sys.exit(1)
|
||||
if result.returncode != 0:
|
||||
console.print("[red]Docker daemon is not running. Please start Docker and try again.[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _start_postgres(container_name: str = "benchmark-postgres") -> Generator[str, None, None]:
|
||||
"""Start a PostgreSQL Docker container. Yields the connection URI."""
|
||||
subprocess.run(["docker", "rm", "-f", container_name], capture_output=True)
|
||||
|
||||
with subprocess.Popen(
|
||||
[
|
||||
"docker",
|
||||
"run",
|
||||
"--rm",
|
||||
"--name",
|
||||
container_name,
|
||||
"-e",
|
||||
f"POSTGRES_PASSWORD={POSTGRES_PASSWORD}",
|
||||
"-e",
|
||||
"POSTGRES_DB=mlflow",
|
||||
"-p",
|
||||
f"127.0.0.1:{POSTGRES_PORT}:5432",
|
||||
"postgres:16-alpine",
|
||||
"-c",
|
||||
"max_connections=500",
|
||||
],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
):
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
TimeElapsedColumn(),
|
||||
console=console,
|
||||
transient=True,
|
||||
) as progress:
|
||||
progress.add_task(" Starting PostgreSQL...", total=None)
|
||||
deadline = time.monotonic() + 30
|
||||
while time.monotonic() < deadline:
|
||||
if (
|
||||
subprocess.run(
|
||||
["docker", "exec", container_name, "pg_isready", "-U", "postgres"],
|
||||
capture_output=True,
|
||||
).returncode
|
||||
== 0
|
||||
):
|
||||
break
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
console.print(" [red]✗ PostgreSQL failed to start within 30s[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
console.print(" [green]✓[/green] PostgreSQL ready")
|
||||
try:
|
||||
yield f"postgresql://postgres:{POSTGRES_PASSWORD}@127.0.0.1:{POSTGRES_PORT}/mlflow"
|
||||
finally:
|
||||
subprocess.run(["docker", "kill", container_name], capture_output=True)
|
||||
|
||||
|
||||
def _basic_auth_header(creds: tuple[str, str] | None) -> dict[str, str]:
|
||||
if creds is None:
|
||||
return {}
|
||||
token = base64.b64encode(f"{creds[0]}:{creds[1]}".encode()).decode()
|
||||
return {"Authorization": f"Basic {token}"}
|
||||
|
||||
|
||||
def _api_post(
|
||||
tracking_uri: str,
|
||||
path: str,
|
||||
body: dict[str, Any],
|
||||
creds: tuple[str, str] | None = None,
|
||||
) -> Any:
|
||||
url = f"{tracking_uri.rstrip('/')}/api/3.0/mlflow/{path}"
|
||||
headers = {"Content-Type": "application/json", **_basic_auth_header(creds)}
|
||||
req = urllib.request.Request(url, data=json.dumps(body).encode(), headers=headers)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
return json.loads(resp.read())
|
||||
except urllib.error.HTTPError as e:
|
||||
console.print(f" [red]API error {e.code} at {url}: {e.read().decode()}[/red]")
|
||||
sys.exit(1)
|
||||
except urllib.error.URLError as e:
|
||||
console.print(f" [red]API error at {url}: {e.reason}[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def _setup_endpoint(
|
||||
tracking_uri: str,
|
||||
fake_server_url: str,
|
||||
endpoint_name: str,
|
||||
usage_tracking: bool,
|
||||
creds: tuple[str, str] | None = None,
|
||||
) -> str:
|
||||
"""Create secret → model definition → endpoint. Returns the invocation URL."""
|
||||
console.print(" Creating secret...")
|
||||
secret_id = _api_post(
|
||||
tracking_uri,
|
||||
_API_SECRET_CREATE,
|
||||
{
|
||||
"secret_name": "benchmark-secret",
|
||||
"secret_value": {"api_key": "fake-benchmark-key"},
|
||||
"provider": "openai",
|
||||
"auth_config": {"api_base": fake_server_url},
|
||||
},
|
||||
creds,
|
||||
)["secret"]["secret_id"]
|
||||
|
||||
console.print(" Creating model definition...")
|
||||
model_def_id = _api_post(
|
||||
tracking_uri,
|
||||
_API_MODEL_DEF_CREATE,
|
||||
{
|
||||
"name": "benchmark-model",
|
||||
"secret_id": secret_id,
|
||||
"provider": "openai",
|
||||
"model_name": "gpt-4o-mini",
|
||||
},
|
||||
creds,
|
||||
)["model_definition"]["model_definition_id"]
|
||||
|
||||
console.print(f" Creating endpoint '{endpoint_name}' (usage_tracking={usage_tracking})...")
|
||||
_api_post(
|
||||
tracking_uri,
|
||||
_API_ENDPOINT_CREATE,
|
||||
{
|
||||
"name": endpoint_name,
|
||||
"model_configs": [
|
||||
{"model_definition_id": model_def_id, "linkage_type": "PRIMARY", "weight": 1.0}
|
||||
],
|
||||
"usage_tracking": usage_tracking,
|
||||
},
|
||||
creds,
|
||||
)
|
||||
|
||||
invoke_url = f"{tracking_uri.rstrip('/')}/gateway/{endpoint_name}/mlflow/invocations"
|
||||
console.print(f" [green]✓[/green] Endpoint ready: [cyan]{invoke_url}[/cyan]")
|
||||
return invoke_url
|
||||
|
||||
|
||||
def _sanity_check(url: str, creds: tuple[str, str] | None = None) -> None:
|
||||
console.print(" Sending sanity-check request...")
|
||||
body = json.dumps({"messages": [{"role": "user", "content": "test"}]}).encode()
|
||||
headers = {"Content-Type": "application/json", **_basic_auth_header(creds)}
|
||||
req = urllib.request.Request(url, data=body, headers=headers)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
if resp.status != 200:
|
||||
console.print(f" [red]✗ Sanity check failed: HTTP {resp.status}[/red]")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
console.print(f" [red]✗ Sanity check failed: {e}[/red]")
|
||||
sys.exit(1)
|
||||
console.print(" [green]✓[/green] Sanity check passed")
|
||||
|
||||
|
||||
def _run_benchmark(
|
||||
url: str,
|
||||
n_requests: int,
|
||||
max_concurrent: int,
|
||||
runs: int,
|
||||
min_rps: float | None = None,
|
||||
max_p50_ms: float | None = None,
|
||||
max_p99_ms: float | None = None,
|
||||
output: Path | None = None,
|
||||
creds: tuple[str, str] | None = None,
|
||||
) -> None:
|
||||
auth = aiohttp.BasicAuth(*creds) if creds else None
|
||||
results = bm.run_benchmark(url, n_requests, max_concurrent, runs, auth)
|
||||
bm.print_results(results)
|
||||
if output is not None:
|
||||
output.write_text(json.dumps(bm.results_to_dict(results), indent=2))
|
||||
console.print(f" Results saved to [cyan]{output}[/cyan]")
|
||||
if not bm.check_thresholds(
|
||||
results, min_rps=min_rps, max_p50_ms=max_p50_ms, max_p99_ms=max_p99_ms
|
||||
):
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _start_nginx(
|
||||
work_dir: str, instance_ports: list[int], port: int, container_name: str = "benchmark-nginx"
|
||||
) -> Generator[None, None, None]:
|
||||
nginx_dir = Path(work_dir) / "nginx"
|
||||
conf_d = nginx_dir / "conf.d"
|
||||
conf_d.mkdir(parents=True)
|
||||
|
||||
upstream_lines = "\n".join(f" server host.docker.internal:{p};" for p in instance_ports)
|
||||
(conf_d / "mlflow.conf").write_text(
|
||||
f"upstream mlflow_backends {{\n"
|
||||
f"{upstream_lines}\n"
|
||||
f" keepalive 512;\n"
|
||||
f" keepalive_requests 100000;\n"
|
||||
f" keepalive_timeout 60s;\n"
|
||||
f"}}\n"
|
||||
f"server {{\n"
|
||||
f" listen {port} reuseport backlog=65535;\n"
|
||||
f" location / {{\n"
|
||||
f" proxy_pass http://mlflow_backends;\n"
|
||||
f" proxy_http_version 1.1;\n"
|
||||
f' proxy_set_header Connection "";\n'
|
||||
f" proxy_set_header Host $host;\n"
|
||||
f" proxy_set_header X-Real-IP $remote_addr;\n"
|
||||
f" proxy_connect_timeout 5s;\n"
|
||||
f" proxy_send_timeout 60s;\n"
|
||||
f" proxy_read_timeout 60s;\n"
|
||||
f" }}\n"
|
||||
f"}}\n"
|
||||
)
|
||||
(nginx_dir / "nginx.conf").write_text(
|
||||
"worker_processes auto;\n"
|
||||
"worker_rlimit_nofile 65535;\n"
|
||||
"events {\n"
|
||||
" worker_connections 16384;\n"
|
||||
" use epoll;\n"
|
||||
" multi_accept on;\n"
|
||||
"}\n"
|
||||
"http {\n"
|
||||
" access_log off;\n"
|
||||
" tcp_nodelay on;\n"
|
||||
" keepalive_timeout 65;\n"
|
||||
" keepalive_requests 100000;\n"
|
||||
" reset_timedout_connection on;\n"
|
||||
" include /etc/nginx/conf.d/*.conf;\n"
|
||||
"}\n"
|
||||
)
|
||||
|
||||
subprocess.run(["docker", "rm", "-f", container_name], capture_output=True)
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
TimeElapsedColumn(),
|
||||
console=console,
|
||||
transient=True,
|
||||
) as progress:
|
||||
progress.add_task(" Starting nginx...", total=None)
|
||||
subprocess.run(
|
||||
[
|
||||
"docker",
|
||||
"run",
|
||||
"--rm",
|
||||
"-d",
|
||||
"--name",
|
||||
container_name,
|
||||
"--add-host=host.docker.internal:host-gateway",
|
||||
"--ulimit",
|
||||
"nofile=65535:65535",
|
||||
"-v",
|
||||
f"{nginx_dir / 'nginx.conf'}:/etc/nginx/nginx.conf:ro",
|
||||
"-v",
|
||||
f"{conf_d}:/etc/nginx/conf.d:ro",
|
||||
"-p",
|
||||
f"127.0.0.1:{port}:{port}",
|
||||
"nginx:alpine",
|
||||
],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
)
|
||||
|
||||
deadline = time.monotonic() + 15
|
||||
while time.monotonic() < deadline:
|
||||
if (
|
||||
subprocess.run(
|
||||
["docker", "exec", container_name, "nginx", "-t"], capture_output=True
|
||||
).returncode
|
||||
== 0
|
||||
):
|
||||
break
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
console.print(" [red]✗ nginx failed to start[/red]")
|
||||
sys.exit(1)
|
||||
|
||||
console.print(" [green]✓[/green] nginx ready")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
subprocess.run(["docker", "kill", container_name], capture_output=True)
|
||||
|
||||
|
||||
def cmd_bench(args: argparse.Namespace) -> None:
|
||||
instances = args.instances
|
||||
mode = "1 instance" if instances == 1 else f"{instances} instances, nginx LB"
|
||||
creds = (args.auth_username, args.auth_password) if args.auth else None
|
||||
|
||||
if args.url:
|
||||
console.print(
|
||||
Panel.fit(
|
||||
f"[bold]Gateway Benchmark[/bold] ({mode})\n"
|
||||
f"URL: [cyan]{args.url}[/cyan]\n"
|
||||
f"Auth: {'basic-auth as ' + args.auth_username if creds else 'disabled'}\n"
|
||||
f"Requests: {args.requests} · Concurrency: {args.max_concurrent}"
|
||||
f" · Runs: {args.runs}",
|
||||
border_style="cyan",
|
||||
)
|
||||
)
|
||||
console.print("\n[bold]Running benchmark[/bold]")
|
||||
_run_benchmark(
|
||||
args.url,
|
||||
args.requests,
|
||||
args.max_concurrent,
|
||||
args.runs,
|
||||
args.min_rps,
|
||||
args.max_p50_ms,
|
||||
args.max_p99_ms,
|
||||
args.output,
|
||||
creds,
|
||||
)
|
||||
return
|
||||
|
||||
needs_docker = instances > 1 or args.database == "postgres"
|
||||
if needs_docker:
|
||||
_check_docker()
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="mlflow-bench-") as work_dir:
|
||||
port = args.port
|
||||
fake_port = args.fake_server_port
|
||||
instance_ports = [args.base_port + i for i in range(instances)]
|
||||
|
||||
auth_line = f"basic-auth as {args.auth_username}" if creds else "disabled"
|
||||
if instances == 1:
|
||||
panel = (
|
||||
f"[bold]Gateway Benchmark[/bold] ({mode})\n"
|
||||
f"Workers: {args.workers} · DB: {args.database.upper()} · "
|
||||
f"Usage tracking: {args.usage_tracking} · Auth: {auth_line}\n"
|
||||
f"Requests: {args.requests} · Concurrency: {args.max_concurrent} · "
|
||||
f"Runs: {args.runs} · Fake delay: {args.fake_delay_ms}ms\n"
|
||||
f"Ports: MLflow :{port} · Fake server :{fake_port}"
|
||||
)
|
||||
else:
|
||||
panel = (
|
||||
f"[bold]Gateway Benchmark[/bold] ({mode})\n"
|
||||
f"Workers/instance: {args.workers} · "
|
||||
f"Total workers: {instances * args.workers} · "
|
||||
f"Usage tracking: {args.usage_tracking} · Auth: {auth_line}\n"
|
||||
f"Requests: {args.requests} · Concurrency: {args.max_concurrent} · "
|
||||
f"Runs: {args.runs} · Fake delay: {args.fake_delay_ms}ms\n"
|
||||
f"Ports: instances {instance_ports[0]}–{instance_ports[-1]}"
|
||||
f" · LB :{port} · Fake server :{fake_port}"
|
||||
)
|
||||
console.print(Panel.fit(panel, border_style="cyan"))
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
stack.callback(lambda: console.print("\n[dim]Cleaning up...[/dim]"))
|
||||
|
||||
# Backend
|
||||
if instances > 1 or args.database == "postgres":
|
||||
console.print("\n[bold]PostgreSQL[/bold]")
|
||||
backend_uri = stack.enter_context(_start_postgres())
|
||||
else:
|
||||
db_path = Path(work_dir) / "mlflow.db"
|
||||
backend_uri = f"sqlite:///{db_path}"
|
||||
console.print(f"\n[dim]Using SQLite: {db_path}[/dim]")
|
||||
|
||||
# Servers
|
||||
console.print("\n[bold]Starting servers[/bold]")
|
||||
stack.enter_context(
|
||||
_start_fake_server(work_dir, port=fake_port, workers=FAKE_SERVER_WORKERS)
|
||||
)
|
||||
|
||||
if instances == 1:
|
||||
stack.enter_context(
|
||||
_start_mlflow(work_dir, port, args.workers, backend_uri, auth=args.auth)
|
||||
)
|
||||
|
||||
console.print("\n[bold]Setting up gateway endpoint[/bold]")
|
||||
invoke_url = _setup_endpoint(
|
||||
f"http://127.0.0.1:{port}",
|
||||
f"http://127.0.0.1:{fake_port}/v1",
|
||||
ENDPOINT_NAME,
|
||||
usage_tracking=args.usage_tracking,
|
||||
creds=creds,
|
||||
)
|
||||
_sanity_check(invoke_url, creds)
|
||||
else:
|
||||
# Start instance 0 first — it initializes the DB schema.
|
||||
# All instances share the same PostgreSQL DB, so starting concurrently
|
||||
# can cause CREATE TABLE race conditions.
|
||||
stack.enter_context(
|
||||
_start_mlflow(
|
||||
work_dir,
|
||||
instance_ports[0],
|
||||
args.workers,
|
||||
backend_uri,
|
||||
"MLflow instance 0",
|
||||
host="0.0.0.0",
|
||||
auth=args.auth,
|
||||
)
|
||||
)
|
||||
for i, p in enumerate(instance_ports[1:], start=1):
|
||||
stack.enter_context(
|
||||
_start_mlflow(
|
||||
work_dir,
|
||||
p,
|
||||
args.workers,
|
||||
backend_uri,
|
||||
f"MLflow instance {i}",
|
||||
host="0.0.0.0",
|
||||
auth=args.auth,
|
||||
)
|
||||
)
|
||||
|
||||
console.print("\n[bold]Setting up gateway endpoint[/bold]")
|
||||
_setup_endpoint(
|
||||
f"http://127.0.0.1:{instance_ports[0]}",
|
||||
f"http://127.0.0.1:{fake_port}/v1",
|
||||
ENDPOINT_NAME,
|
||||
usage_tracking=args.usage_tracking,
|
||||
creds=creds,
|
||||
)
|
||||
|
||||
console.print("\n[bold]Starting nginx load balancer[/bold]")
|
||||
nginx_container = "benchmark-nginx"
|
||||
stack.enter_context(
|
||||
_start_nginx(
|
||||
work_dir, instance_ports, port=port, container_name=nginx_container
|
||||
)
|
||||
)
|
||||
subprocess.run(
|
||||
["docker", "exec", nginx_container, "nginx", "-s", "reload"],
|
||||
capture_output=True,
|
||||
)
|
||||
time.sleep(1)
|
||||
|
||||
invoke_url = f"http://127.0.0.1:{port}/gateway/{ENDPOINT_NAME}/mlflow/invocations"
|
||||
_sanity_check(invoke_url, creds)
|
||||
|
||||
console.print("\n[bold]Running benchmark[/bold]")
|
||||
_run_benchmark(
|
||||
invoke_url,
|
||||
args.requests,
|
||||
args.max_concurrent,
|
||||
args.runs,
|
||||
args.min_rps,
|
||||
args.max_p50_ms,
|
||||
args.max_p99_ms,
|
||||
args.output,
|
||||
creds,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="MLflow AI Gateway benchmark",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--url",
|
||||
metavar="URL",
|
||||
help="Benchmark this endpoint URL directly, skipping server setup entirely",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--instances",
|
||||
type=int,
|
||||
default=int(os.environ.get("INSTANCES", "4")),
|
||||
metavar="N",
|
||||
help=(
|
||||
"Number of MLflow instances to run (default: 4). "
|
||||
"Values >1 require Docker (postgres + nginx). "
|
||||
"Use --instances 1 for a single instance with optional SQLite."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--workers",
|
||||
type=int,
|
||||
default=int(os.environ.get("WORKERS_PER_INSTANCE", "4")),
|
||||
metavar="N",
|
||||
help="Gunicorn/uvicorn worker processes per MLflow instance (default: 4)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--database",
|
||||
choices=["sqlite", "postgres"],
|
||||
default="sqlite",
|
||||
help=(
|
||||
"Database to use — only applies when --instances 1. "
|
||||
"'postgres' auto-starts a Docker container. (default: sqlite)"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-usage-tracking",
|
||||
dest="usage_tracking",
|
||||
action="store_false",
|
||||
default=True,
|
||||
help="Disable usage tracking (tracing) on the benchmark endpoint",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=int(os.environ.get("MLFLOW_PORT", str(MLFLOW_PORT))),
|
||||
metavar="N",
|
||||
help=(
|
||||
"Port the benchmark client sends requests to. "
|
||||
"For --instances 1 this is the MLflow port; "
|
||||
"for --instances >1 this is the nginx load balancer port. (default: 5731)"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--base-port",
|
||||
type=int,
|
||||
default=int(os.environ.get("BASE_PORT", str(INSTANCE_BASE_PORT))),
|
||||
metavar="N",
|
||||
help=(
|
||||
"Starting port for MLflow instances in multi mode. "
|
||||
"Instances listen on base-port, base-port+1, … (default: 5800)"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fake-server-port",
|
||||
type=int,
|
||||
metavar="N",
|
||||
default=int(os.environ.get("FAKE_SERVER_PORT", str(FAKE_SERVER_PORT))),
|
||||
help="Port for the fake OpenAI server that simulates provider latency (default: 9137)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--requests",
|
||||
type=int,
|
||||
default=int(os.environ.get("REQUESTS", "2000")),
|
||||
metavar="N",
|
||||
help="Total requests to send per benchmark run (default: 2000)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-concurrent",
|
||||
type=int,
|
||||
default=int(os.environ.get("MAX_CONCURRENT", "50")),
|
||||
metavar="N",
|
||||
help="Maximum number of in-flight requests at any time (default: 50)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--runs",
|
||||
type=int,
|
||||
default=int(os.environ.get("RUNS", "3")),
|
||||
metavar="N",
|
||||
help="Number of timed runs; results are reported per-run and averaged (default: 3)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fake-delay-ms",
|
||||
type=int,
|
||||
default=int(os.environ.get("FAKE_RESPONSE_DELAY_MS", "50")),
|
||||
metavar="N",
|
||||
help=(
|
||||
"Simulated provider latency in ms. Set to 0 to measure pure MLflow overhead "
|
||||
"with no provider delay. (default: 50)"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
default=None,
|
||||
metavar="FILE",
|
||||
help="Write benchmark results as JSON to FILE (useful for CI artifact upload)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-rps",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Exit 1 if average throughput across runs falls below N req/s (CI threshold)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-p50-ms",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Exit 1 if average P50 latency across runs exceeds N ms (CI threshold)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-p99-ms",
|
||||
type=float,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help="Exit 1 if average P99 latency across runs exceeds N ms (CI threshold)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth",
|
||||
action="store_true",
|
||||
default=os.environ.get("AUTH", "").lower() in ("1", "true"),
|
||||
help=(
|
||||
"Start MLflow with --app-name=basic-auth and authenticate every setup + "
|
||||
"benchmark request using --auth-username/--auth-password."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth-username",
|
||||
default=os.environ.get("AUTH_USERNAME", "admin"),
|
||||
help="Basic auth username (default: admin, from basic_auth.ini)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth-password",
|
||||
default=os.environ.get("AUTH_PASSWORD", "password1234"),
|
||||
help="Basic auth password (default: password1234, from basic_auth.ini)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
os.environ["FAKE_RESPONSE_DELAY_MS"] = str(args.fake_delay_ms)
|
||||
cmd_bench(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,11 @@
|
||||
# Tracing Benchmark
|
||||
|
||||
Per-commit tracing perf check. Runs on push to `master`; trend at https://mlflow.github.io/mlflow/dev/benchmarks/tracing/.
|
||||
|
||||
```bash
|
||||
uv run pytest dev/benchmarks/tracing/ --benchmark-only
|
||||
```
|
||||
|
||||
Add a scenario by writing a `test_*` function in `test_trace_perf.py` — it appears in the chart on the next master push. Renaming a test starts a new trend line.
|
||||
|
||||
Setup is modeled after [`opentelemetry-python`'s benchmark workflow](https://github.com/open-telemetry/opentelemetry-python/blob/main/.github/workflows/benchmarks.yml).
|
||||
@@ -0,0 +1,116 @@
|
||||
import json
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
|
||||
from opentelemetry import trace as trace_api
|
||||
from opentelemetry.sdk.resources import Resource as _OTelResource
|
||||
from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
|
||||
from opentelemetry.trace import SpanContext
|
||||
|
||||
from mlflow.entities.span import Span, SpanType, create_mlflow_span
|
||||
from mlflow.entities.trace_info import TraceInfo
|
||||
from mlflow.entities.trace_location import TraceLocation
|
||||
from mlflow.entities.trace_state import TraceState
|
||||
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
||||
from mlflow.tracing.constant import SpanAttributeKey, TraceTagKey
|
||||
from mlflow.tracing.utils import TraceJSONEncoder
|
||||
|
||||
ENV_CHOICES = ["prod", "staging", "dev"]
|
||||
NAME_PREFIXES = ["agent_run", "qa_chain", "rag_pipeline", "summarizer"]
|
||||
WEEK_MS = 7 * 24 * 60 * 60 * 1000
|
||||
|
||||
SEED_TRACES = 1000
|
||||
SEED_SPANS_PER_TRACE = 10
|
||||
|
||||
|
||||
def generate_trace_data(
|
||||
experiment_id: str,
|
||||
num_spans: int,
|
||||
rng: random.Random,
|
||||
) -> tuple[TraceInfo, list[Span]]:
|
||||
trace_id = f"tr-{uuid.uuid4().hex}"
|
||||
request_time = int(time.time() * 1000) - rng.randint(0, WEEK_MS)
|
||||
name_prefix = rng.choice(NAME_PREFIXES)
|
||||
trace_info = TraceInfo(
|
||||
trace_id=trace_id,
|
||||
trace_location=TraceLocation.from_experiment_id(experiment_id),
|
||||
request_time=request_time,
|
||||
state=rng.choice([TraceState.OK, TraceState.OK, TraceState.OK, TraceState.ERROR]),
|
||||
execution_duration=rng.randint(100, 5000),
|
||||
tags={
|
||||
TraceTagKey.TRACE_NAME: f"{name_prefix}_{trace_id[-4:]}",
|
||||
"env": rng.choice(ENV_CHOICES),
|
||||
},
|
||||
)
|
||||
|
||||
span_types = [SpanType.LLM, SpanType.RETRIEVER, SpanType.TOOL, SpanType.CHAIN]
|
||||
base_ns = 1_000_000_000_000
|
||||
spans: list[Span] = []
|
||||
|
||||
for i in range(num_spans):
|
||||
is_root = i == 0
|
||||
span_type = SpanType.AGENT if is_root else rng.choice(span_types)
|
||||
parent_id = None if is_root else rng.choice(range(max(0, i - 3), i))
|
||||
|
||||
trace_num = rng.randint(1, 2**63)
|
||||
ctx = SpanContext(
|
||||
trace_id=trace_num,
|
||||
span_id=i + 1,
|
||||
is_remote=False,
|
||||
trace_flags=trace_api.TraceFlags(1),
|
||||
trace_state=trace_api.TraceState(),
|
||||
)
|
||||
|
||||
parent_ctx = None
|
||||
if parent_id is not None:
|
||||
parent_ctx = SpanContext(
|
||||
trace_id=trace_num,
|
||||
span_id=parent_id + 1,
|
||||
is_remote=False,
|
||||
trace_flags=trace_api.TraceFlags(1),
|
||||
trace_state=trace_api.TraceState(),
|
||||
)
|
||||
|
||||
attrs: dict[str, object] = {}
|
||||
if is_root:
|
||||
attrs[SpanAttributeKey.INPUTS] = json.dumps(
|
||||
{"query": "What is ML?"}, cls=TraceJSONEncoder
|
||||
)
|
||||
attrs[SpanAttributeKey.OUTPUTS] = json.dumps(
|
||||
{"response": "ML is..."}, cls=TraceJSONEncoder
|
||||
)
|
||||
|
||||
otel_span = OTelReadableSpan(
|
||||
name=f"{span_type.lower()}_{i}" if not is_root else "agent_run",
|
||||
context=ctx,
|
||||
parent=parent_ctx,
|
||||
attributes={
|
||||
"mlflow.traceRequestId": json.dumps(trace_id),
|
||||
"mlflow.spanType": json.dumps(span_type, cls=TraceJSONEncoder),
|
||||
**attrs,
|
||||
},
|
||||
start_time=base_ns + i * 10_000_000,
|
||||
end_time=base_ns + i * 10_000_000 + rng.randint(5_000_000, 50_000_000),
|
||||
status=trace_api.Status(trace_api.StatusCode.OK),
|
||||
resource=_OTelResource.get_empty(),
|
||||
)
|
||||
spans.append(create_mlflow_span(otel_span, trace_id, span_type))
|
||||
|
||||
return trace_info, spans
|
||||
|
||||
|
||||
def seed_traces(
|
||||
store: SqlAlchemyStore,
|
||||
experiment_id: str,
|
||||
count: int,
|
||||
spans_per_trace: int,
|
||||
) -> list[str]:
|
||||
rng = random.Random(123)
|
||||
trace_ids: list[str] = []
|
||||
for _ in range(count):
|
||||
ti, sp = generate_trace_data(experiment_id, spans_per_trace, rng)
|
||||
store.start_trace(ti)
|
||||
store.log_spans(experiment_id, sp)
|
||||
trace_ids.append(ti.trace_id)
|
||||
return trace_ids
|
||||
@@ -0,0 +1,39 @@
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from _data import SEED_SPANS_PER_TRACE, SEED_TRACES, seed_traces
|
||||
|
||||
import mlflow
|
||||
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def bench_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
|
||||
return tmp_path_factory.mktemp("bench")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def store(bench_dir: Path) -> SqlAlchemyStore:
|
||||
db_uri = f"sqlite:///{bench_dir / 'mlflow.db'}"
|
||||
(bench_dir / "artifacts").mkdir()
|
||||
artifact_root = (bench_dir / "artifacts").as_uri()
|
||||
return SqlAlchemyStore(db_uri, artifact_root)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def experiment_id(store: SqlAlchemyStore) -> str:
|
||||
return str(store.create_experiment("bench"))
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def seeded(store: SqlAlchemyStore, experiment_id: str) -> list[str]:
|
||||
return seed_traces(store, experiment_id, SEED_TRACES, SEED_SPANS_PER_TRACE)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def e2e_setup(bench_dir: Path) -> Iterator[None]:
|
||||
mlflow.set_tracking_uri(f"sqlite:///{bench_dir / 'e2e.db'}")
|
||||
mlflow.set_experiment("bench_e2e")
|
||||
yield
|
||||
mlflow.flush_trace_async_logging(terminate=True)
|
||||
@@ -0,0 +1,131 @@
|
||||
import random
|
||||
|
||||
from _data import generate_trace_data
|
||||
from pytest_benchmark.fixture import BenchmarkFixture
|
||||
|
||||
import mlflow
|
||||
from mlflow.entities.span import SpanType
|
||||
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
||||
|
||||
DEFAULT_SPANS = 100
|
||||
INGEST_ROUNDS = 20
|
||||
INGEST_WARMUP = 3
|
||||
|
||||
|
||||
def test_ingest(benchmark: BenchmarkFixture, store: SqlAlchemyStore, experiment_id: str) -> None:
|
||||
rng = random.Random(42)
|
||||
|
||||
def setup():
|
||||
ti, sp = generate_trace_data(experiment_id, DEFAULT_SPANS, rng)
|
||||
return (ti, sp), {}
|
||||
|
||||
def do(ti, sp):
|
||||
store.start_trace(ti)
|
||||
store.log_spans(experiment_id, sp)
|
||||
|
||||
benchmark.pedantic(
|
||||
do, setup=setup, iterations=1, rounds=INGEST_ROUNDS, warmup_rounds=INGEST_WARMUP
|
||||
)
|
||||
|
||||
|
||||
def test_search_by_tag(
|
||||
benchmark: BenchmarkFixture,
|
||||
store: SqlAlchemyStore,
|
||||
experiment_id: str,
|
||||
seeded: list[str],
|
||||
) -> None:
|
||||
benchmark(
|
||||
store.search_traces,
|
||||
locations=[experiment_id],
|
||||
max_results=100,
|
||||
filter_string="tag.env = 'prod'",
|
||||
)
|
||||
|
||||
|
||||
def test_search_by_state(
|
||||
benchmark: BenchmarkFixture,
|
||||
store: SqlAlchemyStore,
|
||||
experiment_id: str,
|
||||
seeded: list[str],
|
||||
) -> None:
|
||||
benchmark(
|
||||
store.search_traces,
|
||||
locations=[experiment_id],
|
||||
max_results=100,
|
||||
filter_string="status = 'ERROR'",
|
||||
)
|
||||
|
||||
|
||||
def test_search_by_name_like(
|
||||
benchmark: BenchmarkFixture,
|
||||
store: SqlAlchemyStore,
|
||||
experiment_id: str,
|
||||
seeded: list[str],
|
||||
) -> None:
|
||||
benchmark(
|
||||
store.search_traces,
|
||||
locations=[experiment_id],
|
||||
max_results=100,
|
||||
filter_string="name LIKE 'rag_pipeline%'",
|
||||
)
|
||||
|
||||
|
||||
def test_search_by_timestamp(
|
||||
benchmark: BenchmarkFixture,
|
||||
store: SqlAlchemyStore,
|
||||
experiment_id: str,
|
||||
seeded: list[str],
|
||||
) -> None:
|
||||
benchmark(
|
||||
store.search_traces,
|
||||
locations=[experiment_id],
|
||||
max_results=100,
|
||||
filter_string="timestamp > 0",
|
||||
order_by=["timestamp DESC"],
|
||||
)
|
||||
|
||||
|
||||
def _run_agent_workflow(num_tools: int, num_docs: int, query: str) -> None:
|
||||
with mlflow.start_span(name="agent_run", span_type=SpanType.AGENT) as root:
|
||||
root.set_inputs({"query": query})
|
||||
|
||||
with mlflow.start_span(name="retrieve", span_type=SpanType.RETRIEVER) as retr:
|
||||
retr.set_inputs({"query": query})
|
||||
docs = [
|
||||
{"id": f"doc_{i}", "score": 0.9 - i * 0.01, "text": f"doc text {i} " * 10}
|
||||
for i in range(num_docs)
|
||||
]
|
||||
retr.set_outputs({"documents": docs})
|
||||
|
||||
with mlflow.start_span(name="plan", span_type=SpanType.CHAIN) as planner:
|
||||
planner.set_inputs({"query": query, "num_docs": len(docs)})
|
||||
steps = [f"step_{i}" for i in range(num_tools)]
|
||||
planner.set_outputs({"steps": steps})
|
||||
|
||||
tool_results = []
|
||||
for step in steps:
|
||||
with mlflow.start_span(name=f"tool:{step}", span_type=SpanType.TOOL) as tool:
|
||||
tool.set_inputs({"step": step})
|
||||
result = {"step": step, "status": "ok", "value": len(step)}
|
||||
tool.set_outputs(result)
|
||||
tool_results.append(result)
|
||||
|
||||
with mlflow.start_span(name="summarize", span_type=SpanType.LLM) as summ:
|
||||
summ.set_inputs({"query": query, "tool_results": tool_results})
|
||||
response = f"Answer to {query!r} using {num_docs} docs and {num_tools} tool calls."
|
||||
summ.set_outputs({"response": response})
|
||||
summ.set_attribute("model", "gpt-test")
|
||||
summ.set_attribute("usage.input_tokens", 1234)
|
||||
summ.set_attribute("usage.output_tokens", 567)
|
||||
|
||||
root.set_outputs({"response": response})
|
||||
|
||||
|
||||
def test_e2e_agent(benchmark: BenchmarkFixture, e2e_setup: None) -> None:
|
||||
counter = [0]
|
||||
|
||||
def do():
|
||||
_run_agent_workflow(num_tools=20, num_docs=20, query=f"q-{counter[0]}")
|
||||
counter[0] += 1
|
||||
|
||||
benchmark(do)
|
||||
Reference in New Issue
Block a user