Tutorial 17: Building an OpenAI-Compatible API Server¶
To make our engine usable, we expose it through an HTTP API. By implementing the OpenAI API format, any tool that works with OpenAI's API (LangChain, LlamaIndex, custom clients) works with our engine without modification.
1. API Design Goals¶
- Compatibility: Match OpenAI's request/response JSON schema exactly.
- Streaming: Support both streaming (SSE) and non-streaming responses.
- Concurrency: Handle multiple requests while the scheduler runs.
- Separation: The server layer should not know about MLX internals — it only talks to the scheduler.
2. The Request Schemas¶
Pydantic models define the request shape:
from pydantic import BaseModel, Field
class ChatCompletionRequest(BaseModel):
"""OpenAI-compatible chat completions request."""
model: str
messages: list[dict] # [{"role": "user", "content": "..."}]
temperature: float = 0.0
top_p: float | None = None
max_tokens: int = 64
stream: bool = False
n: int = 1
stop: list[str] | None = None
user: str | None = None
class CompletionRequest(BaseModel):
"""OpenAI-compatible text completion request (simpler than chat)."""
model: str
prompt: str
temperature: float = 0.0
top_p: float | None = None
max_tokens: int = 64
stream: bool = False
The n field controls how many completions to generate per prompt. The stop
field lists stopping sequences.
3. The Threading Model¶
MLX is not thread-safe — all MLX operations must run on a single thread. But HTTP requests arrive concurrently. The solution is a dedicated worker thread for all GPU work:
HTTP Threads (FastAPI)
│
├── submit request → scheduler.submit()
├── poll → scheduler.poll() ← non-blocking
└── read chunks → SSE response
GPU Worker Thread
│
└── while True:
scheduler.step() ← runs on the GPU thread only
sleep(small_delay)
This avoids locks and race conditions by serializing all MLX operations through one thread.
4. The GPU Entry Point¶
All MLX operations run inside a dedicated executor:
from concurrent.futures import ThreadPoolExecutor
# Single-threaded executor — guarantees only one MLX operation at a time
gpu_executor = ThreadPoolExecutor(max_workers=1)
def _gpu_entry(fn, *args):
"""Run a function on the single GPU thread."""
future = gpu_executor.submit(fn, *args)
return future.result()
Why not multiple workers? Two MLX operations running simultaneously on different threads could: 1. Corrupt shared state (model weights, caches). 2. Cause non-deterministic results from concurrent GPU access. 3. Deadlock if two threads try to allocate from the same pool.
The single-threaded guarantee makes the whole system trivially safe.
5. Handling Streaming Responses¶
Streaming uses Server-Sent Events (SSE). The server sends chunks of text
as they become available, delimited by \n\n:
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
async def handle_chat(request: ChatCompletionRequest, body: dict):
"""Streaming chat completion handler."""
prompt_tokens = [tokenize(p["content"]) for p in request.messages]
# Submit to the scheduler
req_id = engine.submit(
prompt_tokens,
max_new_tokens=request.max_tokens,
temp=request.temperature,
top_p=request.top_p,
)
# Stream responses
async def event_generator():
while True:
tokens, status, reason = engine.poll(req_id)
if tokens:
# Yield each token as an SSE event
text = detokenize(tokens)
payload = {
"id": f"chatcmpl-{req_id}",
"object": "chat.completion.chunk",
"model": request.model,
"choices": [{
"index": 0,
"delta": {"content": text},
"finish_reason": None,
}],
}
yield f"data: {json.dumps(payload)}\n\n"
if status == "done":
# Final event with finish_reason
payload = {
"id": f"chatcmpl-{req_id}",
"object": "chat.completion.chunk",
"model": request.model,
"choices": [{
"index": 0,
"delta": {},
"finish_reason": reason,
}],
}
yield f"data: {json.dumps(payload)}\n\n"
yield "data: [DONE]\n\n"
return
await asyncio.sleep(0.01)
return StreamingResponse(event_generator(), media_type="text/event-stream")
SSE Format¶
data: {"choices": [{"delta": {"content": "Hello"}}]}
data: {"choices": [{"delta": {"content": " world"}}]}
data: {"choices": [{"delta": {"content": "!"}}]}
data: {"choices": [{"finish_reason": "stop"}]}
data: [DONE]
The client assembles the deltas to reconstruct the full response.
6. Backpressure¶
A common bug: the server generates tokens faster than the client can consume them. Without backpressure, the event queue grows unboundedly and consumes unlimited memory.
Backpressure strategies: 1. Small batch polling: Read a few tokens at a time, sleep between reads. 2. Queue limit: Refuse new requests if the pending queue is full. 3. Sliding window: Keep only the last N generated tokens in memory.
Our scheduler naturally provides backpressure: it only generates one token per step. If the client is slow, the queue just holds pending (not yet sampled) requests.
7. Non-Streaming Responses¶
For non-streaming requests, we simply generate everything then respond:
async def handle_completion(request: CompletionRequest):
# Run all generation on the GPU thread (blocking)
def run_blocking():
return engine.run_until_complete(prompt_tokens, ...)
tokens, reason = await asyncio.get_event_loop().run_in_executor(
gpu_executor, run_blocking
)
return JSONResponse({
"id": "cmpl-...",
"object": "text_completion",
"model": request.model,
"choices": [{
"text": detokenize(tokens),
"index": 0,
"finish_reason": reason,
}],
"usage": {
"prompt_tokens": len(prompt_tokens),
"completion_tokens": len(tokens),
"total_tokens": len(prompt_tokens) + len(tokens),
},
})
8. The Complete App¶
from fastapi import FastAPI
app = FastAPI(title="anllm")
@app.get("/v1/models")
async def list_models():
"""List available models (required for OpenAI compatibility)."""
return {
"object": "list",
"data": [{
"id": model_name,
"object": "model",
"owned_by": "anllm",
}],
}
@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
if request.stream:
return await handle_chat(request)
return await handle_chat_non_streaming(request)
@app.post("/v1/completions")
async def completions(request: CompletionRequest):
if request.stream:
return await handle_completion_stream(request)
return await handle_completion(request)
9. Concurrency Testing¶
Testing a concurrent server requires verifying: 1. No deadlocks: Multiple requests should not block each other. 2. Correct isolation: Each request gets its own KV cache and output. 3. Non-flaky: The test should not rely on timing.
A common pattern is to use a deterministic fake model for tests:
class FakeModel:
"""Deterministic model for tests — produces predictable token sequences."""
def __call__(self, inputs, offset=0, cache=None, logits_to_keep=None):
# Return fixed logits that always predict token 42
...
This makes tests fast, deterministic, and unaffected by GPU timing.
10. Summary¶
| Component | Purpose |
|---|---|
| Schemas | Validate request format |
| Scheduler | Run generation across requests |
| GPU executor | Serialize all MLX operations on one thread |
| SSE | Stream tokens to clients |
| Models endpoint | Advertise available models |
The server layer is deliberately thin — it does no MLX work, just translates HTTP requests into scheduler calls and scheduler outputs back into API responses.