ML Inference Server

A FastAPI service that serves a DistilBERT sentiment model (SST-2) with ONNX Runtime on CPU. It adds a hand-rolled asyncio batching queue, an exact-match Redis cache, and Prometheus metrics.

On one CPU node, batching doubles throughput (18 → 37 req/s) and removes the 503 overload errors. Cache hits are served at 653 req/s. Numbers and method are in Benchmark.

Architecture

Client → FastAPI (/predict)
           │
           ├── Redis cache (exact-match, SHA-256 key) → HIT: return immediately
           │
           └── MISS: asyncio.Queue → Batching worker (50ms window, max 32)
                                         │
                                         └── ONNX Runtime inference (thread pool)
                                                  │
                                                  └── Redis SET + return

• Dynamic batching: requests arriving within a 50ms window are run as one batch. The window is a true deadline measured from the first request.
• Overload protection: HTTP 503 when the queue exceeds 500 items.
• Rate limiting: 60 req/min per IP by default (RATE_LIMIT=off disables).
• Observability: Prometheus metrics at /metrics, JSON summary at /metrics/summary, structured JSON logs with request IDs.

Benchmark (measured, single-node CPU)

Locust, 100 users, 60s per mode, on an AMD Ryzen 7 5800H with 1 uvicorn worker. Reproduce with bash benchmarks/run.sh after docker compose up. Full method in benchmarks/results.md.

Mode	Throughput	P50	P99	Failures
Naive (no batching)	18 req/s	4500ms	10000ms	3.5% (503s)
Dynamic batching (max 32)	37 req/s	2300ms	5200ms	0%
Batching + Redis cache (hits)	653 req/s	130ms	210ms	0%

Caveats: one uvicorn worker, CPU inference, and the Locust client runs on the same machine as the server (it competes for CPU). The cache mode sends only the 5 warm-up texts, so every request is a cache hit. The cache is exact-match on input text — not semantic similarity.

Setup

docker compose up --build          # app :8000, redis :6379, prometheus :9090

Local dev without Docker:

python -m venv .venv
source .venv/bin/activate          # Windows: .venv\Scripts\activate
pip install -r requirements-build.txt -r requirements.txt -r requirements-dev.txt
python scripts/export_onnx.py      # one-time: exports DistilBERT to onnx_model/
REDIS_URL= uvicorn app.main:app --reload   # empty REDIS_URL runs without the cache

Tests

pytest -q

Four integration tests cover health, prediction, a cache hit on a repeated input, and concurrent requests being batched together. The cache test needs Redis on localhost:6379 and is skipped otherwise. CI (GitHub Actions) runs the tests against a Redis 7 service and builds the Docker image on every push.

Configuration

Env var	Default	Notes
REDIS_URL	redis://localhost:6379	Empty string disables caching
MAX_BATCH_SIZE	32	1 gives naive per-request serving
MAX_WAIT_MS	50	Batching window
RATE_LIMIT	60/minute	off for load tests
CACHE_TTL	3600	Cache entry TTL (seconds)

API

POST /predict — {"text": "..."} → {"text": "...", "label": "POSITIVE", "score": 0.999, "cached": false, "batch_size": 8}

GET /health — model load state, Redis connection, queue depth. GET /metrics/summary — cache hit rate, average batch size, p50/p99 latency. GET /metrics — raw Prometheus metrics.

Load test

pip install -r requirements-dev.txt
locust -f tests/locustfile.py --headless -u 100 -r 20 -t 60s --host http://localhost:8000

BENCH_MODE controls the workload: miss (unique texts → measures batching), hit (5 fixed texts → measures the cache), or mixed (default).