MLflow’s PyFunc flavor lets you package arbitrary Python code into a reusable inference artifact, but it’s not just for pandas DataFrames.
Here’s a sklearn model saved with MLflow PyFunc, ready for deployment:
import mlflow
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification
# 1. Train a simple model
X, y = make_classification(n_samples=100, n_features=4, random_state=42)
model = RandomForestClassifier(random_state=42)
model.fit(X, y)
# 2. Define a custom PyFunc wrapper
class SklearnModelWrapper(mlflow.pyfunc.PythonModel):
def __init__(self, model):
self.model = model
def predict(self, context, model_input):
# model_input is expected to be a pandas DataFrame by default
# Convert DataFrame to numpy array for sklearn
return self.model.predict(model_input.values)
# 3. Log the model with the custom wrapper
mlflow.pyfunc.log_model(
artifact_path="sklearn_pyfunc",
python_model=SklearnModelWrapper(model),
artifacts={"model_file": "path/to/your/saved/sklearn/model.pkl"} # if you saved it separately
)
# Example of loading and predicting
loaded_model = mlflow.pyfunc.load_model("runs:/<run_id>/sklearn_pyfunc")
sample_data = pd.DataFrame([[0.1, 0.2, 0.3, 0.4]])
predictions = loaded_model.predict(sample_data)
print(predictions)
The core problem PyFunc solves is standardizing the serving of ML models, regardless of the underlying framework. Think of it as a universal adapter. When you log a PyFunc model, MLflow packages your Python code, its dependencies, and any other required artifacts into a self-contained directory. This directory can then be loaded and used for inference in a consistent way, whether it’s a TensorFlow model, a PyTorch model, or even just a simple Python function that performs some data transformation. The predict method is the universal interface.
Internally, when you call mlflow.pyfunc.log_model, MLflow creates a MLmodel file in the specified artifact_path. This file acts as a manifest, describing the model’s flavor (in this case, python_function), the location of the Python code (often a code/ subdirectory), and any other specified artifacts. When you later load this model using mlflow.pyfunc.load_model, MLflow reconstructs the Python environment and makes the predict method available. The context object passed to predict contains information about the loaded model and its artifacts, which can be useful for accessing pre-loaded data or configurations.
The primary lever you control is the predict method within your PythonModel subclass. This method receives a context object and model_input. By default, model_input is a pandas.DataFrame. However, you can configure MLflow to accept other input types. You also manage dependencies by specifying them in an MLmodel file or through conda.yaml when logging. The artifacts parameter in log_model is crucial for including external files like trained model weights, configuration files, or lookup tables that your predict method might need.
The real magic of PyFunc is its ability to abstract away the underlying model framework. You can have a complex tensorflow model and a simple scikit-learn model, both logged as PyFuncs, and they’ll be loaded and invoked with the exact same loaded_model.predict(data) call. This unification is invaluable for creating consistent deployment pipelines and batch inference jobs. MLflow handles the environment setup and dependency resolution behind the scenes.
Many users are unaware that the artifacts dictionary passed to mlflow.pyfunc.log_model can reference any file, not just model weights. You can include configuration files, feature lookup tables, or even small datasets that your predict method needs to access. These artifacts are copied into the MLflow model’s artifact directory and are accessible via context.artifacts within your predict method, enabling more complex inference logic that goes beyond just model scoring.
The next concept you’ll likely encounter is integrating PyFunc models into MLflow’s real-time serving capabilities.