mlflow--mlflow
wehub-resource-sync
794 行
27 KiB
Python
794 行
27 KiB
Python
# /// 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()
|