mlflow How to save a sklearn pipeline with custom transformer?

mlflow

sklearn

machine learning

custom transformer

model deployment

mlflow How to save a sklearn pipeline with custom transformer?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Deep Learning Machine Learning DevOps & Deployment

Introduction

Saving a scikit-learn pipeline with a custom transformer in MLflow requires that the custom transformer class is importable at load time. The standard approach is to define your custom transformer in a separate module (not in a notebook or __main__), then log the pipeline with mlflow.sklearn.log_model(). If the transformer lives in __main__, MLflow cannot deserialize it on load. The solution is to either package the transformer as a module or use MLflow's code_path parameter to include the source file.

The Problem

python

1import mlflow
2from sklearn.pipeline import Pipeline
3from sklearn.base import BaseEstimator, TransformerMixin
4from sklearn.linear_model import LogisticRegression
5
6# Custom transformer defined in the notebook or script
7class TextCleaner(BaseEstimator, TransformerMixin):
8    def fit(self, X, y=None):
9        return self
10
11    def transform(self, X):
12        return [text.lower().strip() for text in X]
13
14pipeline = Pipeline([
15    ('cleaner', TextCleaner()),
16    ('model', LogisticRegression())
17])
18
19# This saves successfully...
20mlflow.sklearn.log_model(pipeline, "model")
21
22# But loading fails:
23model = mlflow.sklearn.load_model("runs:/<run_id>/model")
24# ModuleNotFoundError: No module named '__main__'
25# or
26# AttributeError: Can't get attribute 'TextCleaner' on <module '__main__'>

The issue is that pickle (used by MLflow's sklearn flavor) stores the module path of the class. If the class is in __main__, it cannot be found during deserialization in a different process.

Solution 1: Define Transformer in a Separate Module

python

1# custom_transformers.py (separate file)
2from sklearn.base import BaseEstimator, TransformerMixin
3
4class TextCleaner(BaseEstimator, TransformerMixin):
5    def fit(self, X, y=None):
6        return self
7
8    def transform(self, X):
9        return [text.lower().strip() for text in X]

python

1# train.py
2import mlflow
3from sklearn.pipeline import Pipeline
4from sklearn.linear_model import LogisticRegression
5from custom_transformers import TextCleaner  # Import from module
6
7pipeline = Pipeline([
8    ('cleaner', TextCleaner()),
9    ('model', LogisticRegression())
10])
11
12pipeline.fit(X_train, y_train)
13
14with mlflow.start_run():
15    mlflow.sklearn.log_model(pipeline, "model")
16
17# Loading works because TextCleaner is in custom_transformers module
18model = mlflow.sklearn.load_model("runs:/<run_id>/model")

The key requirement is that custom_transformers must be importable when loading the model — it must be installed as a package or on the Python path.

Solution 2: Use code_path Parameter

Include the source file with the model artifact:

python

1# train.py
2import mlflow
3from custom_transformers import TextCleaner
4
5pipeline = Pipeline([
6    ('cleaner', TextCleaner()),
7    ('model', LogisticRegression())
8])
9
10pipeline.fit(X_train, y_train)
11
12with mlflow.start_run():
13    mlflow.sklearn.log_model(
14        pipeline,
15        "model",
16        code_path=["custom_transformers.py"]  # Include the source file
17    )

When loading, MLflow adds the code_path files to sys.path, making the custom classes importable:

python

model = mlflow.sklearn.load_model("runs:/<run_id>/model")
# MLflow automatically adds custom_transformers.py to the path

Solution 3: Package as a Python Module

For production, package your transformers as an installable module:

1my_ml_package/
2    setup.py
3    my_ml_package/
4        __init__.py
5        transformers.py    # TextCleaner defined here

python

1# setup.py
2from setuptools import setup, find_packages
3
4setup(
5    name="my-ml-package",
6    version="1.0.0",
7    packages=find_packages(),
8)

python

1# train.py
2from my_ml_package.transformers import TextCleaner
3
4# Log with pip_requirements to ensure the package is installed on load
5mlflow.sklearn.log_model(
6    pipeline,
7    "model",
8    pip_requirements=["my-ml-package==1.0.0"]
9)

Solution 4: Use MLflow's Pyfunc Flavor

For maximum portability, wrap the pipeline in a custom pyfunc model:

python

1import mlflow.pyfunc
2import cloudpickle
3
4class SklearnPipelineWrapper(mlflow.pyfunc.PythonModel):
5    def __init__(self, pipeline):
6        self.pipeline = pipeline
7
8    def predict(self, context, model_input):
9        return self.pipeline.predict(model_input)
10
11pipeline.fit(X_train, y_train)
12
13with mlflow.start_run():
14    mlflow.pyfunc.log_model(
15        "model",
16        python_model=SklearnPipelineWrapper(pipeline),
17        code_path=["custom_transformers.py"],
18        pip_requirements=["scikit-learn", "pandas"]
19    )

Custom Transformer with Parameters

python

1# custom_transformers.py
2from sklearn.base import BaseEstimator, TransformerMixin
3import numpy as np
4
5class OutlierRemover(BaseEstimator, TransformerMixin):
6    def __init__(self, threshold=3.0):
7        self.threshold = threshold  # Must match parameter name in __init__
8
9    def fit(self, X, y=None):
10        self.mean_ = np.mean(X, axis=0)
11        self.std_ = np.std(X, axis=0)
12        return self
13
14    def transform(self, X):
15        z_scores = np.abs((X - self.mean_) / self.std_)
16        mask = (z_scores < self.threshold).all(axis=1)
17        return X[mask]
18
19    def get_params(self, deep=True):
20        return {"threshold": self.threshold}
21
22    def set_params(self, **params):
23        for key, value in params.items():
24            setattr(self, key, value)
25        return self

Scikit-learn requires get_params() and set_params() to work with Pipeline, GridSearchCV, and serialization. BaseEstimator provides default implementations that use __init__ parameter names.

Complete Example

python

1# custom_transformers.py
2from sklearn.base import BaseEstimator, TransformerMixin
3import pandas as pd
4
5class FeatureSelector(BaseEstimator, TransformerMixin):
6    def __init__(self, columns=None):
7        self.columns = columns
8
9    def fit(self, X, y=None):
10        return self
11
12    def transform(self, X):
13        return X[self.columns] if self.columns else X
14
15# train.py
16import mlflow
17from sklearn.pipeline import Pipeline
18from sklearn.preprocessing import StandardScaler
19from sklearn.linear_model import LogisticRegression
20from custom_transformers import FeatureSelector
21
22pipeline = Pipeline([
23    ('selector', FeatureSelector(columns=['age', 'income'])),
24    ('scaler', StandardScaler()),
25    ('model', LogisticRegression())
26])
27
28pipeline.fit(X_train, y_train)
29accuracy = pipeline.score(X_test, y_test)
30
31with mlflow.start_run():
32    mlflow.log_metric("accuracy", accuracy)
33    mlflow.sklearn.log_model(
34        pipeline,
35        "model",
36        code_path=["custom_transformers.py"],
37        registered_model_name="my-classifier"
38    )

Common Pitfalls

Defining custom transformers in __main__: Pickle stores the module path. If the class is in __main__ (a notebook cell or the main script), it cannot be deserialized in a different process. Always define custom transformers in importable modules.
Forgetting code_path: Without code_path, the model artifact does not include your custom code. Loading the model in a different environment fails with ModuleNotFoundError. Always specify code_path or install the transformer module.
Mismatched __init__ parameters and get_params: Scikit-learn's default get_params() inspects __init__ parameter names. If __init__ takes self.threshold but you set self._threshold, cloning and serialization break. The attribute name must match the parameter name.
Stateful transformers without fit: If your transformer learns from data (e.g., computes mean/std), those fitted attributes must be serializable. Avoid storing non-picklable objects (database connections, file handles) as fitted attributes.
Version mismatches: The scikit-learn version used to save the model must match the version used to load it. Log the version with mlflow.log_param("sklearn_version", sklearn.__version__) and specify it in pip_requirements.

Summary

Define custom transformers in separate importable modules, not in __main__
Use code_path=["my_module.py"] in mlflow.sklearn.log_model() to bundle custom code
For production, package transformers as installable Python packages
Inherit from BaseEstimator and TransformerMixin for scikit-learn compatibility
Ensure __init__ parameter names match attribute names for get_params() to work
Use mlflow.pyfunc.log_model() for maximum portability across environments

System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

mlflow How to save a sklearn pipeline with custom transformer?

ML System Design practice on Codemia

Introduction

The Problem

Solution 1: Define Transformer in a Separate Module

Solution 2: Use code_path Parameter

Solution 3: Package as a Python Module

Solution 4: Use MLflow's Pyfunc Flavor

Custom Transformer with Parameters

Complete Example

Common Pitfalls

Summary

Related reading

System Design Fundamentals

Track what you have practised

ML System Design practice on Codemia