DocsInferenceServe custom modelsinference.yaml
Inference · Serve custom models

inference.yaml

The custom-mode yaml — module, class_name, init_args, weights, and cluster sizing.

The full file

The yaml binds your class to the platform: which file, which class, what arguments, what hardware. No engine field means custom mode.

inference.yaml

Field reference

FieldRequired?What it does
model.moduleyesthe Python module under scripts/, no .py. Dotted paths reach subfolders: "predictors.unet"scripts/predictors/unet.py. Checked against your files at submit.
model.class_nameyesthe class to instantiate. Its presence in the module file is checked at submit too.
model.init_argsdefaults to {}passed verbatim as keyword arguments to your __init__. YAML scalars, lists, and maps only.
model.pathoptionala weights file; when set, your load_weights() is called with the mounted local path. Leave it out if __init__ pulls weights from HuggingFace Hub itself.
cluster.num_workersyesworker machines = replicas. Each runs one instance of your class.
cluster.devicedefaults to "gpu""auto" lets your code pick cuda over cpu; "gpu" requires a GPU worker.
cluster.gpus_per_worker / cpus_per_workerdefault 1 / 8what each worker must expose.
cluster.min_vram_gb_per_workeroptionalfree-VRAM floor for worker selection. Set it to weights + activations, rounded up.
cluster.max_ongoing_requestsdefaults to 16how many requests can be inside one replica's predict() at once. Keep low (1–8) for heavy models, higher for fast ones.
cluster.modedefaults to "data_parallel"the only valid value in custom mode — each replica is a full independent copy.
requirementsoptionalpip installs on top of the base image.
i
The base image already ships torch, fastapi, httpx, and pillow. Don't re-pin those; list only what your model genuinely adds (the SAM example adds transformers and numpy). Every extra line slows boot.

Weights: two patterns

Pull from the Hub. Skip model.path and download inside __init__ (from_pretrained(...)). Simplest, and right for public checkpoints.

Ship your own file. Upload a weights file with the bundle, set model.path, and implement load_weights(path). The platform mounts the file and hands you the local path after __init__. Right for fine-tunes and private checkpoints.

Submit mechanics for both (including the weights upload): Submit inference via SDK.