icsy7867/complexity-router repository overview

⁠ModernBERT Complexity Router REST API

This service wraps predict_complexity_chunked_v2.py with FastAPI and server-side dynamic batching.

⁠Request flow

1. HTTP request is validated by FastAPI.
2. Request enters a bounded asyncio.Queue.
3. The batching worker waits up to BATCH_WAIT_MS for nearby requests.
4. Up to MAX_BATCH_SIZE conversations are preprocessed together.
5. All history chunks from those conversations are flattened into encoder batches.
6. All final-user sequences are encoded in batches.
7. The conversation router head is evaluated as one vectorized batch.
8. Each HTTP request receives its own score and latency information.

Only one model inference worker executes against a model instance, preventing a burst of HTTP requests from causing 100 simultaneous calls into the same model.

⁠Model layout

Mount the existing trained directory at /models/router:

/models/router/
├── encoder/
│   └── config.json
├── router_model.pt
├── router_config.json
└── tokenizer/

The encoder weights do not need to be duplicated in encoder/model.safetensors; the complete trained state is loaded from router_model.pt by the existing predictor.

⁠CPU image

docker build -f Dockerfile.cpu -t modernbert-router:cpu .

docker run --rm \
  -p 8000:8000 \
  -v "$PWD/complexity-router-v2/best:/models/router:ro" \
  -e CPU_THREADS=8 \
  -e MAX_BATCH_SIZE=16 \
  -e MAX_ENCODER_BATCH_SIZE=32 \
  modernbert-router:cpu

CPU defaults to FP32. BF16 can be forced with -e DTYPE=bf16 if the CPU supports it and it benchmarks better.

⁠CUDA image

docker build -f Dockerfile.cuda -t modernbert-router:cuda .

docker run --rm \
  --gpus all \
  -p 8000:8000 \
  -v "$PWD/complexity-router-v2/best:/models/router:ro" \
  modernbert-router:cuda

CUDA defaults to BF16.

The CUDA image is based on pytorch/pytorch:2.11.0-cuda12.8-cudnn9-runtime.

⁠Request

curl -s http://127.0.0.1:8000/v1/complexity \
  -H 'Content-Type: application/json' \
  -d '{
    "messages": [
      {
        "role": "user",
        "content": "Write a Python function that parses a CSV file."
      }
    ]
  }' | jq

Example response shape:

{
  "complexity": 2.4173,
  "request_id": "...",
  "device": "cuda",
  "dtype": "torch.bfloat16",
  "batch_size": 12,
  "queue_ms": 2.31,
  "preprocessing_ms": 0.74,
  "encoder_ms": 31.82,
  "router_head_ms": 1.09,
  "service_ms": 34.12,
  "total_ms": 36.43,
  "history_chunks": 0,
  "final_tokens": 17
}

⁠Endpoints

• POST /v1/complexity - score a conversation.
• GET /health - process health.
• GET /ready - confirms model is loaded and reports batching configuration.
• GET /metrics - lightweight queue/batch counters in JSON.

Use /ready for Kubernetes readiness checks. Use /health for liveness checks.

⁠Important environment variables

⁠100-request burst test

Install the test dependency on the client:

pip install httpx

Then:

python burst_test.py \
  --url http://127.0.0.1:8000/v1/complexity \
  --requests 100 \
  --concurrency 100

The script reports throughput, HTTP statuses, p50/p95/p99/max latency, and observed batch size.

A bounded queue protects the process from bursts. It does not create unlimited compute capacity: if requests arrive faster than the model can drain them, latency grows. When the queue is completely full the API returns HTTP 503 with Retry-After: 1 instead of consuming unbounded memory.

⁠Kubernetes

Run one Uvicorn worker per pod. Do not set --workers above 1 unless you explicitly want one complete model copy per worker.

For more throughput, scale pods horizontally. Each pod gets its own resident model and batching queue. Configure the readiness probe against /ready so a pod receives traffic only after its model is loaded.

For a GPU deployment, request one GPU per pod unless intentionally using another topology.