ML Model Deployment skill for Codex

A model is deployed properly when the exact artifact that passed evaluation is the one serving (pinned by hash and version), preprocessing is shared code rather than a rewrite, the service rejects malformed input, readiness reflects that the model is loaded, p95 latency and memory are measured under realistic concurrency, and a rollback is one command.

Inspect first

1. Read AGENTS.md, the training code, and how the model is saved today (pickle, joblib, torch.save, save_pretrained, MLflow). Find the preprocessing used at training time.
2. Collect requirements: request rate, p95 latency budget, CPU vs GPU, batch vs online, input size, where it runs (Kubernetes, VM, serverless, managed endpoint).
3. Freeze dependency versions from the training environment (pip freeze or the lockfile). Version skew between training and serving (scikit-learn, tokenizers, numpy) is the most common cause of silent prediction changes.

Choose the serving stack

Situation	Default
One sklearn / XGBoost / small torch model, modest traffic	FastAPI + the native model, or ONNX Runtime for speed
GPU model, needs request batching, multiple models in one service	BentoML
LLM text generation	Not here: use vllm-serving (or ollama-local-llms for local)
Many teams, model registry, canaries built in	Managed endpoint (SageMaker, Vertex AI, Azure ML) or KServe
Offline scoring of a table	Batch job, no service at all
Cross-language or edge deployment	Export to ONNX

Artifact and manifest

Ship a directory, not a lone pickle:

models/churn/2026-10-20-a1b2c3/
├── model.onnx            # or model.joblib / model.safetensors
├── preprocess.py         # or a serialized sklearn Pipeline that includes preprocessing
├── manifest.json
└── requirements.lock

{
  "name": "churn",
  "version": "2026-10-20-a1b2c3",
  "sha256": "<hash of model file>",
  "framework": "onnxruntime==1.20.*",
  "input_schema": {"tenure_months": "int", "plan": "str", "monthly_spend": "float"},
  "metrics": {"auc": 0.871, "eval_set": "holdout_2026_09"},
  "training_commit": "<git sha>"
}

Verify the hash at load time. Never unpickle files from untrusted sources: pickle and joblib execute code on load. Prefer ONNX or safetensors when the artifact crosses a trust boundary.

FastAPI service

from contextlib import asynccontextmanager
import hashlib, json, pathlib
import numpy as np, onnxruntime as ort
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

MODEL_DIR = pathlib.Path("/models/churn/current")
state: dict = {}

class Features(BaseModel):
    tenure_months: int = Field(ge=0, le=600)
    plan: str = Field(pattern="^(basic|pro|enterprise)$")
    monthly_spend: float = Field(ge=0)

class Prediction(BaseModel):
    churn_probability: float
    model_version: str

@asynccontextmanager
async def lifespan(app: FastAPI):
    manifest = json.loads((MODEL_DIR / "manifest.json").read_text())
    blob = (MODEL_DIR / "model.onnx").read_bytes()
    if hashlib.sha256(blob).hexdigest() != manifest["sha256"]:
        raise RuntimeError("model hash mismatch")
    opts = ort.SessionOptions(); opts.intra_op_num_threads = 2
    state["sess"] = ort.InferenceSession(blob, opts, providers=["CPUExecutionProvider"])
    state["manifest"] = manifest
    yield
    state.clear()

app = FastAPI(lifespan=lifespan)

@app.get("/healthz")
def healthz(): return {"ok": True}                       # liveness: process is up

@app.get("/readyz")
def readyz():                                            # readiness: model loaded
    if "sess" not in state: raise HTTPException(503, "loading")
    return {"model_version": state["manifest"]["version"]}

@app.post("/predict", response_model=Prediction)
def predict(f: Features):                                # sync def: runs in threadpool
    x = preprocess(f)                                    # same function used in training
    # skl2onnx classifier exported with ZipMap disabled: outputs = [labels, probabilities]
    prob = state["sess"].run(None, {"input": x})[1][0][1]
    return Prediction(churn_probability=float(prob), model_version=state["manifest"]["version"])

• Inference is CPU-bound: use def endpoints (threadpool) or offload; never run model code inside async def without run_in_executor, it blocks the event loop.
• Size workers to cores: one model copy per worker process. With 4 workers and a 2 GB model you need 8 GB plus headroom. Cap intra-op threads so workers do not oversubscribe.
• Return the model version in every response for debugging and audits.

BentoML (batching, GPU, containers)

import bentoml, numpy as np

@bentoml.service(resources={"gpu": 1}, traffic={"timeout": 30})
class Classifier:
    def __init__(self):
        import torch
        self.model = torch.jit.load("model.pt").eval().cuda()

    @bentoml.api(batchable=True, max_batch_size=64, max_latency_ms=20)
    def predict(self, inputs: np.ndarray) -> np.ndarray:
        import torch
        with torch.inference_mode():
            return self.model(torch.from_numpy(inputs).cuda()).cpu().numpy()

bentoml serve service:Classifier locally, bentoml build then bentoml containerize for an image (needs Docker and network). Adaptive batching is where GPU throughput comes from; tune max_latency_ms against the latency budget.

ONNX export

import torch
model.eval()
example = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, (example,), "model.onnx", dynamo=True,
                  input_names=["input"], output_names=["logits"],
                  dynamic_shapes={"x": {0: torch.export.Dim("batch")}})

• Use the argument name of forward as the key in dynamic_shapes (here x). Older code uses dynamic_axes with the TorchScript exporter; both work, the dynamo exporter is the current default path.
• scikit-learn models export with skl2onnx (to_onnx(pipeline, X[:1].astype(np.float32), ...)); disable the classifier’s ZipMap output through the converter’s zipmap option so probabilities come back as an array, not a list of dicts. Include the preprocessing pipeline in the export.
• Check onnx.checker.check_model and compare outputs with the original framework (parity test below) before deploying. GPU inference needs onnxruntime-gpu and CUDAExecutionProvider.

Parity, latency and rollout

Parity test (in CI): run 1,000 held-out rows through the training pipeline and through the deployed service or exported model; assert np.allclose(a, b, atol=1e-4) for scores and identical labels for >= 99.9% of rows.

Latency: benchmark with realistic payloads and concurrency (see load-testing for k6). Report p50/p95/p99, throughput at the target concurrency, peak RSS and GPU memory, and cold-start time (load to ready).

Rollout:

• Build an image tagged with the model version; keep the previous image.
• Deploy to staging, run the parity and smoke tests against it.
• Production deploys and traffic shifts need explicit user confirmation. Use a canary (5-10% traffic) or shadow mode, compare prediction distributions and error rates, then promote. Rollback = redeploy the previous tag.
• Monitor input drift and prediction distribution after release.

Workflow

• Read AGENTS.md, training code, current artifact format, dependency locks
• Choose serving stack from the table
• Package artifact directory with manifest and hash
• Share preprocessing code between training and serving
• Build service with validation, health and readiness, version in responses
• Parity test and latency benchmark
• Container image (see dockerfile-best-practices), probes, resource limits
• Staging, then confirmed canary rollout with rollback ready

Gotchas

• Re-implementing preprocessing in the service “from memory” is the top cause of training-serving skew. Import it or serialize it with the model.
• Loading the model per request instead of once at startup adds seconds of latency.
• Readiness that returns 200 before the model loads sends traffic to a cold pod.
• Kubernetes CPU limits below the thread count cause throttling and p99 spikes; align thread settings with the CPU request.
• Float32 vs float64 differences between frameworks break exact-equality tests; use tolerances.