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: TrueRemember
- Signatures describe inputs, outputs and inference params
- pyfunc and serving enforce the schema on every call
- Infer from real data with
infer_signatureor aninput_example
Watch out for
- Logging a model trained on integer columns that can be null at inference, so schema enforcement rejects float inputs.