wehub-resource-sync
146 lines
7.9 KiB
Markdown
146 lines
7.9 KiB
Markdown
# 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 |
|