Sign inStart now ›
Inference /inference.yaml

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
model:
  module: "my_serve"            # scripts/my_serve.py — dotted path for subfolders
  class_name: "YourPredictor"   # the class to instantiate
  init_args:                    # → YourPredictor(**init_args)
    num_classes: 3
  path: "model/weights.pt"      # optional — see Weights below

cluster:
  num_workers: 1
  device: "auto"                # your __init__ picks cuda > cpu
  gpus_per_worker: 1
  cpus_per_worker: 8
  min_vram_gb_per_worker: 6     # free VRAM each selected worker must have
  max_ongoing_requests: 8       # per-replica concurrency cap

requirements:
  - "transformers>=4.56,<5"     # your extras — see the note below

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.
note
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.

Previousserve.pyNext Examples gallery