Serve an ML Model on CPU: ONNX Export, FastAPI, Dynamic Batching and a p95 Latency Budget

"""Serving Harrow County's tile classifier on one CPU core: ONNX export, a fast runtime, an HTTP API, dynamic
batching and a latency budget."""
import asyncio
import base64
import binascii
import random
import time

import numpy as np
import onnxruntime as ort
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from model import MEAN, STD

ONNX_PATH = "tile_model.onnx"
TILE_BYTES = 64 * 64 * 3
MAX_TILES = 32


# ---------- Step 1: export ----------
def export_onnx(model, path=ONNX_PATH):
    """Export the model to ONNX with torch.onnx.export(..., dynamo=False): one input named "tiles" (N, 3, 64, 64)
    float32, one output named "logits"; the batch dimension must be dynamic (dynamic_axes) so one file serves any
    batch size."""
    # TODO (Step 1): torch.onnx.export(model, (example,), path, input_names=..., output_names=...,
    # dynamic_axes=..., dynamo=False) with a (1, 3, 64, 64) example; return path.
    raise NotImplementedError("Step 1: write export_onnx()")


def max_difference(model, session, x):
    """The largest absolute difference between the PyTorch model's logits and the ONNX session's on x
    (a float32 tensor), as a float."""
    # TODO (Step 1): run both on x and return the largest absolute difference as a float.
    raise NotImplementedError("Step 1: write max_difference()")


# ---------- Step 2: a fast runtime ----------
def preprocess(tiles):
    """uint8 (n, 64, 64, 3) -> float32 (n, 3, 64, 64): / 255, minus MEAN, divided by STD, channels first,
    contiguous. numpy only: the server does not need torch."""
    raise NotImplementedError("preprocess() arrives in Step 2")


class Classifier:
    """ONNX Runtime session on one thread, plus the class names."""

    def __init__(self, path, classes):
        self.classes = classes
        raise NotImplementedError("Classifier.__init__() arrives in Step 2")

    def predict(self, tiles):
        """[{"label": class name, "confidence": softmax probability of that class}] for every tile (uint8)."""
        raise NotImplementedError("Classifier.predict() arrives in Step 2")


def per_tile_ms(fn, tiles, repeats=20):
    """Median milliseconds per tile of fn(tiles) over `repeats` timed calls, after 2 untimed warm-up calls."""
    raise NotImplementedError("per_tile_ms() arrives in Step 2")


# ---------- Step 3: an HTTP API ----------
class PredictRequest(BaseModel):
    tiles: list[str]      # each a base64 string of the tile's 12,288 raw RGB bytes (64 x 64 x 3, row by row)


def decode_tile(b64):
    """The (64, 64, 3) uint8 array for one base64 tile. Raises ValueError if it is not valid base64 or not
    exactly TILE_BYTES bytes."""
    raise NotImplementedError("decode_tile() arrives in Step 3")


def create_app(classifier, batcher=None):
    """GET /health -> {"status": "ok", "classes": [...]}.
    POST /predict with PredictRequest -> {"predictions": [...]} (one per tile, from classifier.predict, or from
    batcher.submit for each tile when a batcher is given). Answer 422 with a helpful "detail" when there are no
    tiles, more than MAX_TILES, or a tile that decode_tile() rejects (say which one)."""
    raise NotImplementedError("create_app() arrives in Step 3")


# ---------- Step 4: dynamic batching ----------
class MicroBatcher:
    """Collects single tiles from many concurrent requests into batches: the worker takes the first waiting tile,
    then keeps taking more until it has max_batch or max_wait_ms has passed since the first, runs ONE
    classifier.predict() on the whole batch (in a thread, so new requests keep arriving), and hands every caller
    its own result. If predict() raises, every caller in that batch gets the exception."""

    def __init__(self, classifier, max_batch=16, max_wait_ms=2.0):
        self.classifier, self.max_batch, self.max_wait = classifier, max_batch, max_wait_ms / 1000
        self.queue, self.worker, self.batch_sizes = None, None, []

    async def submit(self, tile):
        """The prediction for one tile, once its batch has run. Starts the worker on first use."""
        raise NotImplementedError("MicroBatcher.submit() arrives in Step 4")

    async def _run(self):
        raise NotImplementedError("MicroBatcher._run() arrives in Step 4")

    def stop(self):
        if self.worker:
            self.worker.cancel()


# ---------- Step 5: a latency budget ----------
def percentile(values, q):
    """The q-th percentile (0-100) of values, linear interpolation (numpy's default), as a float."""
    raise NotImplementedError("percentile() arrives in Step 5")


async def load_test(batcher, tile, rate, n, seed=0):
    """Send n requests of one tile as a Poisson stream at `rate` per second: arrival times are running sums of
    random.Random(seed).expovariate(rate). Every 5 ms, start every request that is due, each as its own task that
    measures from its start to its answer. Returns {"p50", "p95" (ms), "throughput" (answers per second)}."""
    raise NotImplementedError("load_test() arrives in Step 5")


def pick_config(results, budget_ms):
    """results: {(max_batch, max_wait_ms): {"quiet": stats, "peak": stats}}. Among the configurations whose p95
    is within budget_ms under BOTH loads, the one with the lowest peak p95. None if no configuration fits."""
    raise NotImplementedError("pick_config() arrives in Step 5")


CONFIG = {"max_batch": 1, "max_wait_ms": 0}