Intermediate

MLflow Models

Learn the MLflow Models format — a standard way to package models from any framework for portable deployment.

MLflow Models Format

An MLflow Model is a directory containing a MLmodel file that describes the model's flavors (framework-specific formats) along with the serialized model artifacts. This allows the same model to be loaded and served by multiple tools.

YAML — MLmodel file structure

artifact_path: model
flavors:
  python_function:
    env: conda.yaml
    loader_module: mlflow.sklearn
    model_path: model.pkl
    python_version: 3.11.0
  sklearn:
    code: null
    pickled_model: model.pkl
    serialization_format: cloudpickle
    sklearn_version: 1.4.0
model_uuid: abc123def456
signature:
  inputs: '[{"name": "age", "type": "double"}, {"name": "income", "type": "double"}]'
  outputs: '[{"type": "long"}]'

Model Flavors

Flavors are MLflow's way of supporting multiple ML frameworks. Each flavor defines how to save and load a model:

Flavor	Module	Frameworks
sklearn	mlflow.sklearn	scikit-learn
pytorch	mlflow.pytorch	PyTorch
tensorflow	mlflow.tensorflow	TensorFlow / Keras
xgboost	mlflow.xgboost	XGBoost
lightgbm	mlflow.lightgbm	LightGBM
transformers	mlflow.transformers	Hugging Face Transformers
pyfunc	mlflow.pyfunc	Any Python model (universal)

Model Signature and Input Example

Python — Defining model signature

from mlflow.models import infer_signature

with mlflow.start_run():
    model = RandomForestClassifier(n_estimators=100)
    model.fit(X_train, y_train)

    # Infer signature from data
    predictions = model.predict(X_test)
    signature = infer_signature(X_test, predictions)

    # Log model with signature and input example
    mlflow.sklearn.log_model(
        model,
        artifact_path="model",
        signature=signature,
        input_example=X_test.iloc[:3],  # First 3 rows as example
    )

✅

Always include a signature: Model signatures enable input validation at serving time, catching schema mismatches before they cause errors. The input example helps users understand the expected format.

Logging Models

Python — Logging models from different frameworks

# scikit-learn
mlflow.sklearn.log_model(sklearn_model, "sklearn-model")

# PyTorch
mlflow.pytorch.log_model(pytorch_model, "pytorch-model")

# TensorFlow
mlflow.tensorflow.log_model(tf_model, "tf-model")

# XGBoost
mlflow.xgboost.log_model(xgb_model, "xgb-model")

# Hugging Face Transformers
mlflow.transformers.log_model(
    transformers_model=pipeline,
    artifact_path="hf-model",
    task="text-classification",
)

Loading Models

Python — Loading models

# Load as generic Python function (works with any flavor)
model = mlflow.pyfunc.load_model("runs:/abc123/model")
predictions = model.predict(new_data)

# Load as native framework model
sklearn_model = mlflow.sklearn.load_model("runs:/abc123/model")
pytorch_model = mlflow.pytorch.load_model("runs:/def456/model")

# Load from model registry
model = mlflow.pyfunc.load_model("models:/churn-predictor/Production")

# Load from a specific version
model = mlflow.pyfunc.load_model("models:/churn-predictor/3")

Custom Model Wrappers

For complex models with custom preprocessing or multi-model ensembles, create a custom pyfunc wrapper:

Python — Custom model wrapper

import mlflow.pyfunc
import pandas as pd

class ChurnPredictor(mlflow.pyfunc.PythonModel):
    """Custom model that combines preprocessing with prediction."""

    def load_context(self, context):
        """Load artifacts when the model is loaded."""
        import joblib
        self.model = joblib.load(context.artifacts["model_path"])
        self.scaler = joblib.load(context.artifacts["scaler_path"])
        self.feature_names = context.artifacts["feature_names"]

    def predict(self, context, model_input: pd.DataFrame) -> pd.DataFrame:
        """Preprocess and predict."""
        # Apply scaling
        scaled_input = self.scaler.transform(model_input[self.feature_names])

        # Get predictions and probabilities
        predictions = self.model.predict(scaled_input)
        probabilities = self.model.predict_proba(scaled_input)[:, 1]

        return pd.DataFrame({
            "churn_prediction": predictions,
            "churn_probability": probabilities,
        })

# Log the custom model
with mlflow.start_run():
    artifacts = {
        "model_path": "artifacts/model.pkl",
        "scaler_path": "artifacts/scaler.pkl",
        "feature_names": "artifacts/features.json",
    }

    mlflow.pyfunc.log_model(
        artifact_path="model",
        python_model=ChurnPredictor(),
        artifacts=artifacts,
        conda_env="conda.yaml",
    )

Model Packaging

MLflow models can be packaged for different deployment targets:

Shell — Packaging models

# Build a Docker image from a model
mlflow models build-docker -m "runs:/abc123/model" -n "my-model-image"

# Generate a pip-installable package
mlflow models generate-dockerfile -m "runs:/abc123/model" -d "./dockerfile-dir"

← Previous Projects Next → Model Registry