ML in Rust Best Practices
Reproducible, fast, memory-safe ML inference and serving in Rust.
How to Use This List
- Apply A-B when adding a new model to the repo.
- Walk C-D before exposing any HTTP endpoint.
- Revisit E after GPU OOM, accuracy regressions, or latency incidents.
A - Model Lifecycle
- Pin model revision, tokenizer.json, and crate versions in one lockfile manifest. Drift breaks silently at decode time.
- Store ONNX/GGUF checksums alongside artifacts. Reject downloads that fail hash verify.
- Document preprocessing (mean/std, NCHW, max length) in model card. Serving must mirror training exactly.
- Run golden batch tests against Python reference outputs. Epsilon thresholds on logits or embeddings.
- Version API responses with
model_versionfield. Clients know when behavior changed.
B - Inference Correctness
- Validate input shapes before allocating tensors. Return 400, not panic, on bad client data.
- Normalize embeddings for cosine search. L2 norm on insert and query paths.
- Cap
max_tokensand context length server-side. Client params are hints, not policy. - Apply chat templates before tokenization for LLMs. Template mismatch garbles generations.
- Fail startup if model files missing. Do not serve traffic with lazy load errors mid-request.
C - Performance and Resources
- Use
spawn_blockingor dedicated pool for forwards. Never block Tokio worker threads on matmul. - Probe GPU at boot and log fallback to CPU. Operators need visible degraded mode.
- Limit concurrent inference with semaphores. Protect VRAM from request stampedes.
- Batch embed jobs during offline index builds. GPU likes batch > 1.
- Right-size quant level (Q4/Q8) vs quality SLA. Measure perplexity or task metric, not guess.
D - Serving and Security
- Authenticate
/predictand generation routes. GPUs are expensive to share with the internet. - Set request body limits and timeouts. Multi-megabyte JSON is a DoS vector.
- Separate liveness and readiness probes. Readiness runs cheap dummy inference.
- Stream LLM tokens for perceived latency. Buffering full completions frustrates users.
- Redact PII from logs; hash prompts if needed for audit. Compliance beats debug convenience.
E - Operations Culture
- Track time-to-first-token and tokens/sec for LLMs. Product SLOs live here, not just HTTP 200 rate.
- Record CPU/GPU memory after model load in metrics. Catch leaks across deploys.
- Feature-flag CPU-only CI; GPU images separate. Builds must not require NVIDIA in PR checks.
- Plan index rebuild when embedding model changes. Old vectors are incompatible with new geometry.
- Keep training in Python when faster; own inference in Rust. Play to ecosystem strengths.
FAQs
What is the top production mistake?
Running synchronous model forward inside async Axum handlers without spawn_blocking, stalling unrelated routes.
How many golden tests?
At least one fixed input per model output head - classification logits, embedding vector, or first-token logit.
Rust for training?
Fine for small models and research; large LLM training usually stays Python with Rust owning inference.
When to use ort vs candle?
ort for exported ONNX from any framework; candle for HF Rust ports and tight Rust-only stacks.
How to document GPU needs?
VRAM at load, per-request KV growth, and max concurrent sessions in runbook tied to hardware SKU.
Tokenizer updates?
Treat tokenizer.json changes like model changes - rerun golden tests and bump model_version.
RAG specifics?
Version chunking policy and embedding model together; changing one without the other hurts recall.
CI without GPU?
CPU EP only in PR pipeline; nightly GPU job for perf regression on larger fixtures.
Error handling?
Result through blocking tasks; map to HTTP status without leaking tensor shape errors to clients.
Where to learn serving patterns?
See Serving Models with Axum for handler layout.
Related
- Serving Models with Axum - HTTP inference
- ONNX Runtime (ort) - deployment runtime
- LLM Inference & Serving - generation ops
- ML in Rust Basics - section map
Stack versions: This page was written for Rust 1.97.0 (edition 2024), Tokio 1.x, Axum 0.8, serde 1.0, sqlx 0.8, clap 4, and Polars 0.46+.