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