Signatures and input examples

MLflow · Practical · 10 min

Schema enforcement catches wrong columns and types before they silently produce bad predictions.

A model signature records the schema of a model's inputs, outputs and optional inference params. It is stored in the MLmodel file, shown in the UI, and enforced when the model is called through pyfunc or served.

Create one by inference:

from mlflow.models import infer_signature
signature = infer_signature(X_train, model.predict(X_train))

or let log_model(..., input_example=X.head()) infer it automatically. Signatures can be column-based (named, typed columns from a DataFrame or dict) or tensor-based (dtype and shape, with -1 for variable dimensions).

What enforcement does on predict:

  • Missing required columns raise an error; extra columns are dropped with a warning.
  • Types are converted where safe (int to float) and rejected otherwise.
  • Columns marked optional may be absent.

Params let callers change inference behaviour per request without changing the input data, e.g. {"threshold": 0.7} or a max_tokens for an LLM. Declare them with infer_signature(X, y, params={"threshold": 0.5}); the default comes from the signature.

Integer columns with missing values are a classic trap: pandas turns them into floats at inference time and enforcement fails. Cast such columns to float before training and logging.

Example

import mlflow
import pandas as pd
from mlflow.models import infer_signature
from sklearn.linear_model import LogisticRegression

X = pd.DataFrame({"age": [25.0, 40.0, 33.0, 51.0], "visits": [1.0, 5.0, 2.0, 8.0]})
y = [0, 1, 0, 1]
model = LogisticRegression().fit(X, y)

signature = infer_signature(X, model.predict(X))
print(signature.inputs)   # ['age': double (required), 'visits': double (required)]

with mlflow.start_run():
    info = mlflow.sklearn.log_model(model, name="clf", signature=signature,
                                    input_example=X.head(2))

loaded = mlflow.pyfunc.load_model(info.model_uri)
try:
    loaded.predict(pd.DataFrame({"age": [30.0]}))   # 'visits' missing
except mlflow.exceptions.MlflowException as e:
    print("rejected:", "visits" in str(e))           # rejected: True

Remember

  • Signatures describe inputs, outputs and inference params
  • pyfunc and serving enforce the schema on every call
  • Infer from real data with infer_signature or an input_example

Watch out for

  • Logging a model trained on integer columns that can be null at inference, so schema enforcement rejects float inputs.

Check yourself

  1. A served model's signature requires columns a and b. A request sends a, b and c. What happens by default?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: mlflow.org/docs/latest/ml/model/signatures/