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
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.
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]
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:
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:
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
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)
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:
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 )
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
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
