How to package a custom Python ML model for deployment with MLflow

Summary

• Use MLflow's pyfunc flavor to wrap your custom model code, specifying dependencies and input signatures for reproducible deployment
• Log your packaged model with mlflow.pyfunc.log_model() and register it in Unity Catalog for version control and access management
• At serving time, MLflow automatically recreates the environment and runs your model code, ensuring consistency between training and production

How to package a custom Python ML model for deployment with MLflow

MLflow's pyfunc (Python function) flavor lets you package custom model code alongside dependencies. This ensures your model runs the same way in production as during training, avoiding environment drift. Log your model with MLflow, register it in Unity Catalog, and deploy it with Model Serving or batch inference jobs.

What is MLflow pyfunc?

pyfunc is a generic Python model format that wraps your custom inference logic. Instead of relying on framework-specific formats (which only work for that framework), pyfunc lets you define a predict() method that takes any input and returns predictions. This flexibility supports custom preprocessing, ensemble logic, or proprietary algorithms.

Key capabilities to evaluate

Capability	Why it matters
Dependency management	MLflow captures exact package versions so your model runs identically in dev, staging, and production
Input signatures	Defining schemas enables automatic validation and type conversion, preventing silent errors
Framework-agnostic	pyfunc works with any Python library: scikit-learn, custom code, FastAPI apps, or model combinations
Reproducibility	Logging and versioning your model code ensures you can always revert to a known good version
Easy deployment	Registered models work with Model Serving, batch jobs, and feature serving without modification

How Databricks approaches custom model packaging

1. Create a custom model class
Define a class with a predict() method:

import mlflow.pyfunc
import pandas as pd
import numpy as np

class CustomFraudDetector(mlflow.pyfunc.PythonModel):
    def __init__(self, threshold=0.5):
        self.threshold = threshold
        self.feature_scaler = None
    
    def predict(self, context, model_input):
        """Custom prediction logic."""
        # model_input can be pandas DataFrame or numpy array
        if isinstance(model_input, pd.DataFrame):
            features = model_input.values
        else:
            features = model_input
        
        # Your custom logic here
        scores = features @ self.weights + self.bias
        predictions = (scores > self.threshold).astype(int)
        return pd.DataFrame({"fraud_flag": predictions})

2. Log with MLflow
Specify dependencies and an input example:

import mlflow

    #Create model instance
model = CustomFraudDetector(threshold=0.6)
model.weights = trained_weights
model.bias = trained_bias

    #Log with dependencies
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        artifact_path="fraud_model",
        python_model=model,
        pip_requirements=["scikit-learn==1.3.2", "pandas==2.1.1"],
        input_example=pd.DataFrame({"amount": [100], "merchant_risk": [0.2]}),
        signature=mlflow.models.signature.infer_signature(
            pd.DataFrame({"amount": [100], "merchant_risk": [0.2]}),
            pd.DataFrame({"fraud_flag": [0]})
        )
    )
    mlflow.register_model("runs:/<run_id>/fraud_model", "fraud_detector")

3. Use in Model Serving
Deploy the registered model:

from databricks.sdk import WorkspaceClient

client = WorkspaceClient()
endpoint = client.serving_endpoints.create(
    name="fraud_endpoint",
    config={
        "served_models": [{
            "model_name": "models:/my_catalog.ml.fraud_detector",
            "model_version": "1"
        }]
    }
)

Model Serving automatically installs dependencies and calls your predict() method.
4. Use in batch inference
Load and apply your model:

model_uri = "models:/my_catalog.ml.fraud_detector@production"
model = mlflow.pyfunc.load_model(model_uri)

    #Apply to a Spark DataFrame
df_input = spark.read.table("transactions")
predictions = model.predict(df_input.toPandas())

Use cases

• Ensemble models: Combine predictions from multiple base models with custom logic
• Domain-specific preprocessing: Apply custom feature engineering that's not available in standard libraries
• Legacy models: Wrap proprietary or legacy code so it works with modern deployment platforms
• Multi-stage pipelines: Chain custom preprocessing, model inference, and postprocessing in one artifact

FAQs

What if my model has complex dependencies?

Specify all dependencies in pip_requirements or use a conda.yaml file. MLflow creates an isolated environment when loading the model, ensuring dependencies don't conflict with others.

Can I include data files with my model?

Yes. Use the artifacts directory. Log large objects (scalers, lookup tables) as MLflow artifacts and load them in your model's predict() method.

How do I test my model before deployment?

Load locally and call predict():

model = mlflow.pyfunc.load_model("runs:/<run_id>/fraud_model")
test_input = pd.DataFrame({"amount": [100], "merchant_risk": [0.2]})
print(model.predict(test_input))

What if my model needs GPU?

Mlflow works with GPU-dependent packages. Specify them in dependencies. When deploying to Model Serving, select a GPU endpoint type.