Production model serving has two recurring sources of friction: packaging the runtime environment and exposing a stable prediction API. Research code often assumes a particular Python version, GPU runtime, CUDA/cuDNN combination, model weight layout, and preprocessing pipeline. Service code then has to turn that local convention into a container, an HTTP server, input validation, file handling, concurrency control, cancellation, readiness checks, and an OpenAPI contract.
A schema-driven model container treats the prediction function as the product boundary. The model author declares a container environment, implements a typed predictor class, and lets the runtime derive the HTTP API, request validation, output encoding, and operational endpoints from that declaration. Replicate Cog is a concrete implementation of this pattern: cog.yaml describes the Docker image, a Predictor class describes setup and prediction, and the built container runs a Rust-based prediction server that exposes synchronous, asynchronous, idempotent, health-check, cancellation, and OpenAPI endpoints.
The transferable lesson is not the particular command-line tool. The useful architecture is the separation between model code, environment specification, schema extraction, and serving runtime. That separation lets teams package GPU models consistently without rewriting every model behind a custom Flask app or hand-maintained OpenAPI file.
The container specification captures the runtime ingredients that are usually hidden in notebooks: Python version, system packages, requirements files, GPU enablement, CUDA overrides, setup commands, SDK version pins, and optional image naming. When GPU support is enabled, the build system can choose compatible base images and pass GPU flags at execution time.
The predictor class declares two different lifecycle phases. setup() performs expensive one-time work such as loading weights, downloading artifacts, or constructing preprocessing objects. predict() runs per request and uses type annotations plus input metadata to describe the external API.
The serving process reads the model contract and exposes a standardized HTTP surface. Inputs become JSON fields, file paths become upload or URL values, outputs become JSON or files, custom metrics become response fields, and the same contract can generate an OpenAPI document.
Instead of maintaining a separate REST schema that can drift from the Python prediction function, derive validation, documentation, idempotency, and client behavior from the typed predictor signature. This makes model packaging closer to compiling a service boundary than hand-writing glue code.
The environment file names the build inputs and the predictor entry point. The prediction runtime can then build an image, install dependencies, and know which class to load when the server starts.
build:
gpu: true
python_version: "3.13"
python_requirements: requirements.txt
system_packages:
- "ffmpeg"
predict: "predict.py:Predictor"
concurrency:
max: 4The predictor separates cold-start work from request work. This avoids reloading weights for every prediction and gives readiness probes a meaningful lifecycle: starting, ready, busy, setup failed, or unhealthy.
from cog import BasePredictor, Input, Path
class Predictor(BasePredictor):
def setup(self):
self.model = load_model("./weights")
def predict(
self,
image: Path = Input(description="Input image"),
scale: float = Input(default=1.0, ge=0.25, le=4.0),
) -> Path:
tensor = preprocess(image)
result = self.model(tensor, scale=scale)
return postprocess(result)A robust prediction server distinguishes “wait for the result” from “start work and notify me later.” It also needs retry-safe request creation, because clients and load balancers can repeat requests after timeouts. An idempotent endpoint with a caller-provided prediction ID prevents duplicate expensive model runs.
Model teams can ship image, audio, video, embedding, or language models without writing a bespoke web service for each one. The predictor owns preprocessing and postprocessing, while the container builder handles dependency installation and GPU runtime selection.
A platform team can require every model package to expose the same health check, OpenAPI schema, prediction endpoint, metrics field, and cancellation behavior. That makes model containers easier to route, test, monitor, and roll back.
Asynchronous predictions with webhooks fit workflows where images, videos, or long generations complete after the client has disconnected. Output and log events can be emitted progressively, while the final completion event gives the workflow engine a terminal state.
Walk the predictor signature, extract argument names, Python types, default values, validation constraints, descriptions, and return annotations, then build a schema from the result. The generated schema should be the same source of truth used by request validation.
def schema_from_predictor(fn):
sig = inspect.signature(fn)
fields = {}
for name, param in sig.parameters.items():
annotation = param.annotation
input_meta = param.default if isinstance(param.default, InputSpec) else None
fields[name] = {
"type": map_python_type(annotation),
"required": param.default is inspect._empty,
"description": getattr(input_meta, "description", None),
"minimum": getattr(input_meta, "ge", None),
"maximum": getattr(input_meta, "le", None),
}
return {"type": "object", "properties": fields}Represent prediction lifecycle as a small state machine: starting, processing, succeeded, failed, canceled. Webhooks, logs, output chunks, metrics, and cancellation should all attach to that lifecycle rather than being ad hoc side effects.
For file outputs, return a path-like object from the predictor and let the serving runtime encode or upload the file. In production, prefer a configured upload destination over base64 data URLs when outputs are large.
BUSY or conflict behavior when all slots are occupied.