• October 13, 2026
• 11 min read

Pickles go stale: saving ML models that survive dependency upgrades

Introduction

I have often seen data scientists treat model deployment as an afterthought, including how the model is saved (serialized) and loaded (deserialized) in the production environment. As a consequence, the model is often a pickled scikit-learn pipeline, with data preprocessing and the whole shebang baked into a single large object. The caveat is that the pickle protocol depends on the identical versions of packages between the environment where it was saved and where it's unpickled, and herein lies a culprit, as ML models seldom live in a vacuum¹ — security updates especially might require the deployment environment to have its dependencies updated². While this will certainly work for now, it might not load cleanly once the dependencies of the production environment are updated.

This post addresses techniques and design decisions for the cases where you cannot guarantee that the environment where the model is deployed stays consistent between the training and production environments.

Example of a model breaking

Let's make an example of a pickle file breaking because we saved a model using one version of scikit-learn and try loading it with a newer version:

imputer_save.py

# /// script
# dependencies = ["scikit-learn==1.7.2", "pandas"]  # (1)
# ///
import pickle

import pandas as pd
from sklearn.impute import SimpleImputer

imputer = SimpleImputer().fit(pd.DataFrame({"a": [1.0, None, 3.0]}))
print(imputer.transform(pd.DataFrame({"a": [None]})))
# Output:
# [[2.]]

with open("imputer.pkl", "wb") as f:
    pickle.dump(imputer, f)

1. Run this snippet using uv run imputer_save.py to use this exact version of scikit-learn.

imputer_load.py

# /// script
# dependencies = ["scikit-learn==1.8.0", "pandas"]
# ///
import pickle

import pandas as pd

with open("imputer.pkl", "rb") as f:
    imputer = pickle.load(f)

print(imputer.transform(pd.DataFrame({"a": [None]})))

We save the model using scikit-learn 1.7.2 and load it with 1.8.0; unfortunately for us the API of the SimpleImputer changed between the versions, and the call fails when we try to predict with the pickled model under a new version of scikit-learn.

$ uv run imputer_save.py
$ uv run imputer_load.py
~/.cache/uv/environments-v2/imputer-load-2e047ca73b91fb2f/lib/python3.14/site-packages/sklearn/base.py:463:
InconsistentVersionWarning: Trying to unpickle estimator SimpleImputer from version 1.7.2 when using version 1.8.0.
This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:
https://scikit-learn.org/stable/model_persistence.html#security-maintainability-limitations
  warnings.warn(
Traceback (most recent call last):
  File "~/src/blog_allanlrh/blog_allanlrh/docs/posts/pickles-go-stale/imputer_load.py", line 9, in <module>
    print(imputer.transform([[None]]))
          ~~~~~~~~~~~~~~~~~^^^^^^^^^^
  File "~/.cache/uv/environments-v2/imputer-load-2e047ca73b91fb2f/lib/python3.14/site-packages/sklearn/utils/_set_output.py", line 316, in wrapped
    data_to_wrap = f(self, X, *args, **kwargs)
  File "~/.cache/uv/environments-v2/imputer-load-2e047ca73b91fb2f/lib/python3.14/site-packages/sklearn/impute/_base.py", line 632, in transform
    self._fill_dtype, copy=False
    ^^^^^^^^^^^^^^^^
AttributeError: 'SimpleImputer' object has no attribute '_fill_dtype'. Did you mean: '_fit_dtype'?

Write savers and loaders

Fortunately, the error above can be circumvented by renaming a property in the loaded model prior to calling .transform():

imputer_load_fixed.py

# /// script
# dependencies = ["scikit-learn==1.8.0", "pandas"]
# ///
import pickle

import pandas as pd

with open("imputer.pkl", "rb") as f:
    imputer = pickle.load(f)

# 1.8.0's fit() sets the new _fill_dtype to the same value as the existing _fit_dtype.
imputer._fill_dtype = imputer._fit_dtype

print(imputer.transform(pd.DataFrame({"a": [None]})))
# Output:
# [[2.]]

This hack works, but we will have to apply fixes like this wherever the model is loaded, and that is frustrating and error-prone. And I'd also argue that the model is the "information expert" here, so it's the right place to keep the code responsible for saving and loading it.

We can mitigate problems like this by treating the serialization and deserialization as a problem addressed from the very start of the model lifecycle, explicitly introducing a storage format with version identifiers.

But first consider that we only need the transform() method in production; .fit() and a handful of methods are only used during the model development. To that end, let's make a small wrapper for the model, exposing just what we need while also handling serialization and deserialization.

We define a wrapper-class that holds the actual model (a SimpleImputer in this example), and a list of features that it expects. The class exposes the .transform() method and methods for saving and loading the model; see the examples below. A hardcoded version identifier for the wrapper class and the version of scikit-learn is saved alongside the model. The model wrapper identifier will become useful as the wrapper class evolves, and the scikit-learn version is crucial for determining if we should apply the patch or not. Besides the versions, we also save the list of features, allowing us to make a .transform() which accepts a DataFrame which have at least the columns the model needs, and if any is missing, we will see the very familiar KeyError from Pandas.

imputer_simple_wrapped.py

import base64
import json
import pickle
from collections.abc import Sequence
from pathlib import Path
from typing import Protocol, Self

import pandas as pd
import sklearn
from packaging.version import Version


class SupportsTransform(Protocol):
    def transform(self, X: pd.DataFrame) -> pd.DataFrame: ...


class SerializeableImputer:
    """
    Wrapper for a model object supporting the transform method.
    Implements saving, loading and calling of the transform method.
    """

    # Bump when the saved format changes, and handle older versions in load()
    MODEL_WRAPPER_VERSION = 1

    def __init__(self, imputer: SupportsTransform, features: Sequence[str]) -> None:
        self.imputer = imputer
        # Our own copy, so the caller can't change it behind our back
        self.features = list(features)

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        # Impute only our features, in the order the model expects, and keep the
        # caller's other columns, column order and index as they were
        X_out = X.copy()
        X_out[self.features] = self.imputer.transform(X[self.features])
        return X_out

    def save(self, path: Path) -> None:
        """Save the features, versions and pickled model to a JSON file."""
        serialized = dict(
            features=self.features,
            model_pickle=base64.b64encode(pickle.dumps(self.imputer)).decode(),
            MODEL_WRAPPER_VERSION=self.MODEL_WRAPPER_VERSION,
            sklearn_version=sklearn.__version__,
        )
        path.write_text(json.dumps(serialized))

    @classmethod
    def load(cls, path: Path) -> Self:
        """Alternative constructor: load a saved model, patching it if needed."""
        raw = json.loads(path.read_bytes())
        MODEL_WRAPPER_VERSION = raw["MODEL_WRAPPER_VERSION"]
        if MODEL_WRAPPER_VERSION > cls.MODEL_WRAPPER_VERSION:
            raise ValueError(
                f"Model saved with model version {MODEL_WRAPPER_VERSION}, "
                f"but this code only supports up to version {cls.MODEL_WRAPPER_VERSION}"
            )
        imputer = pickle.loads(base64.b64decode(raw["model_pickle"]))
        # scikit-learn 1.8.0 added _fill_dtype to SimpleImputer. Models pickled with an
        # older version lack it, so we set it the same way 1.8.0's fit() does.
        saved_sklearn = Version(raw["sklearn_version"])
        if saved_sklearn < Version("1.8.0"):
            imputer._fill_dtype = imputer._fit_dtype
        return cls(
            imputer=imputer,
            features=raw["features"],
        )

Now, let's see it in action, training and saving the model with scikit-learn 1.7.2 and loading and patching it under 1.8.0.

imputer_wrapped_save.py

# /// script
# dependencies = ["scikit-learn==1.7.2", "pandas", "packaging"]
# ///
from pathlib import Path

import pandas as pd
from imputer_simple_wrapped import SerializeableImputer
from sklearn.impute import SimpleImputer

X = pd.DataFrame({"a": [1.0, None, 3.0]})
imputer = SimpleImputer().fit(X)
print(imputer.transform(pd.DataFrame({"a": [None]})))
# Output:
# [[2.]]

SerializeableImputer(imputer, features=["a"]).save(Path("imputer_wrapped.json"))

imputer_wrapped_load.py

# /// script
# dependencies = ["scikit-learn==1.8.0", "pandas", "packaging"]
# ///
from pathlib import Path

import pandas as pd
from imputer_simple_wrapped import SerializeableImputer

imputer = SerializeableImputer.load(Path("imputer_wrapped.json"))
print(imputer.transform(pd.DataFrame({"a": [None]})))
# Output:
#      a
# 0  2.0

Use AI for this

Implementing and testing a wrapper class like this is the kind of task coding agents excel at. You should of course make it test that the model produces identical results after a round-trip conversion compared to the un-wrapped one — which gives the agent a very concrete test case to validate its implementation.

Version agnostic formats

We have so far relied on pickle files, which is the reason for the headache in the first place.

A better choice is a truly version agnostic format like ONNX, which saves a model to a version agnostic format compatible with its own runtime. The example with the SimpleImputer is basic enough that ONNX can handle it without our wrapper class, but I'm a fan of using it anyway as it gives me more control of the process, and lets me define an interface for saving, loading and inference, models independent of how it's implemented.

First, a minimal working example:

imputer_onnx.py

# /// script
# dependencies = ["scikit-learn==1.7.2", "pandas", "skl2onnx", "onnxruntime"]
# ///
import onnxruntime
import pandas as pd
import skl2onnx
from sklearn.impute import SimpleImputer

X_train = pd.DataFrame({"a": [1.0, None, 3.0]}, dtype="float32")

# Train the base SimpleImputer...
imputer = SimpleImputer().fit(X_train)
# ...convert to onnx...
imputer_onnx = skl2onnx.to_onnx(model=imputer, X=X_train[:1])
# ...and load it into a session to get an executeable object
imputer_session = onnxruntime.InferenceSession(imputer_onnx.SerializeToString())

X_test = pd.DataFrame({"a": [None]}, dtype="float32")
inputs = {col: X_test[[col]].to_numpy() for col in X_test.columns}

print(imputer_session.run(output_names=None, input_feed=inputs)[0])
# Output:
# [[2.]]
print(imputer.transform(X_test))
# Output:
# [[2.]]

Now, the example above requires the caller of the model to adapt to the ONNX format when saving, loading and calling the model, and I think that work should be done upfront when the model is developed.

imputer_onnx_wrapped.py

import base64
import json
from collections.abc import Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Self

import numpy as np
import onnxruntime as ort
import pandas as pd

if TYPE_CHECKING:
    from sklearn.base import BaseEstimator


class SerializeableImputer:
    # Default for new conversions; ONNX Runtime's Imputer operator only supports float32
    DTYPE = "float32"
    # Bump when the saved format changes, and handle older versions in load()
    MODEL_WRAPPER_VERSION = 1

    def __init__(self, onnx_model: bytes, features: Sequence[str], dtype: str = DTYPE) -> None:
        self.onnx_model = onnx_model
        # Our own copy, so the caller can't change it behind our back
        self.features = list(features)
        self.dtype = dtype
        self.session = ort.InferenceSession(onnx_model)

    @classmethod
    def from_scikit_learn_model(
        cls, imputer: BaseEstimator, sample_data: pd.DataFrame
    ) -> Self:
        """
        Alternative constructor, allows creating a SerializeableImputer from a
        scikit-learn model and sample data.
        """
        # Only needed when converting, so loading works without scikit-learn and skl2onnx
        from skl2onnx import to_onnx

        # We only save the ONNX-converted model, so that this class always use
        # that rather than the Scikit Learn model. That ensure the model behaves
        # the same both before and after a serilisation-deserialization
        # round-trip
        onnx_model = to_onnx(imputer, sample_data.head(1).astype(cls.DTYPE))
        return cls(onnx_model.SerializeToString(), list(sample_data.columns))

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        # We handle the data conversion for the user of this class, making the
        # underlying model implementation transparrent to them.
        inputs = {
            feature: X[[feature]].to_numpy(dtype=self.dtype) for feature in self.features
        }
        (output,) = self.session.run(None, inputs)
        # ONNX Runtime may also return sparse tensors, lists and dicts
        if not isinstance(output, np.ndarray):
            raise TypeError(f"Unexpected output type: {type(output)}")
        # Assumes one output column per feature, i.e. no all-missing features were dropped in fit
        # Keep the caller's other columns, column order and index as they were
        X_out = X.copy()
        X_out[self.features] = output
        return X_out

    def save(self, path: str | Path) -> None:
        """Save the features, version, dtype and ONNX model to a JSON file."""
        serialized = {
            "features": self.features,
            "MODEL_WRAPPER_VERSION": self.MODEL_WRAPPER_VERSION,
            "dtype": self.dtype,
            "onnx_model": base64.b64encode(self.onnx_model).decode(),
        }
        Path(path).write_text(json.dumps(serialized, indent=2))

    @classmethod
    def load(cls, path: str | Path) -> Self:
        """Alternative constructor: load a model saved with save()."""
        raw = json.loads(Path(path).read_text())
        MODEL_WRAPPER_VERSION = raw["MODEL_WRAPPER_VERSION"]
        if MODEL_WRAPPER_VERSION > cls.MODEL_WRAPPER_VERSION:
            raise ValueError(
                f"Model saved with model version {MODEL_WRAPPER_VERSION}, "
                f"but this code only supports up to version {cls.MODEL_WRAPPER_VERSION}"
            )
        return cls(
            onnx_model=base64.b64decode(raw["onnx_model"]),
            features=raw["features"],
            dtype=raw["dtype"],
        )

Using the model above becomes

imputer_onnx_wrapped_save.py

# /// script
# dependencies = ["scikit-learn==1.7.2", "pandas", "skl2onnx", "onnxruntime"]
# ///
import pandas as pd
from imputer_onnx_wrapped import SerializeableImputer
from sklearn.impute import SimpleImputer

X = pd.DataFrame({"a": [1.0, None, 3.0]}, dtype="float32")
imputer = SimpleImputer().fit(X)
print(imputer.transform(pd.DataFrame({"a": [None]})))
# Output:
# [[2.]]

wrapped = SerializeableImputer.from_scikit_learn_model(imputer, X)
wrapped.save("imputer_onnx_wrapped.json")

imputer_onnx_wrapped_load.py

# /// script
# dependencies = ["onnxruntime", "pandas"]
# ///
import pandas as pd
from imputer_onnx_wrapped import SerializeableImputer

imputer = SerializeableImputer.load("imputer_onnx_wrapped.json")
print(imputer.transform(pd.DataFrame({"a": [None]})))
# Output:
#      a
# 0  2.0

Pickle, ONNX, or something else?

It really comes down to how confident you are about the production environment staying consistent with your training environment, and what risks you are willing to accept vs. the complexity of your code, because ONNX adds a bit of complexity.

When you use pickles, I'd encourage you to minimize when they are used if another and simple solution exists; in the examples above we store the list of required feature names in a JSON list, and write some simple Python code that works with that. That is way less likely to break than a custom scikit-learn Pipeline element that does the same operation, including encoding the feature names.

ONNX can't encode all models, so that is another reason why we might have to fall back to pickle files, and there can be small numerical differences between ONNX and scikit-learn, though these are usually negligible.

It's also worth checking if the model you use has a native, version-agnostic serialization format, such as the JSON format for XGBoost and LightGBM (you might have to grab the underlying LightGBM Booster object if you are using the scikit-learn compatible interface). PyTorch has a compiler that allows you to load your models in future versions of PyTorch as well. If you are using a linear model, e.g. linear regression, consider just storing the weights and implementing the algorithm by hand.

Also, consider if you want to use scikit-learn pipelines if a dependency-version agnostic format is important for you, because those objects contain many smaller ones, and each one is prone to breaking when the API is updated. If you need a lot of custom behaviour, I have good experience with implementing a model class from scratch, having that own the underlying ML model and enforcing business rules etc. onto the model output — and it can implement saves and loaders as well.

Conclusion

We have seen how pickled objects can break when the dependencies are updated, and how we can use version-agnostic formats to circumvent the problem. In cases where we can't use a version-agnostic serialization format, we sometimes can get away with patching the loaded objects — provided that we tested that the behaviour doesn't change.

The main idea is that we should anticipate a change of the environment when developing the model, and define savers and loaders for the model by wrapping it in a class, which gives us a piece of code, that we can update and evolve as the environment evolves. And crucially, that code stays close to the model, and isn't replicated across the various places where the model is deployed. This also makes rollbacks safe: if a model retrained under newer dependencies has to be rolled back, the older artifact is loaded into the updated production environment, and the loader knows how to handle it, instead of the rollback itself breaking production.

Of course, if you already have a setup that allows you to deploy the model with the exact dependencies that were used during training, then by all means, that's what you should do — this approach is for when that's not an option.

---

1. Some do, because they are packaged in a docker container which handles only the model and some way of calling it from within the container, and while this is a good way of handling issues around model serialization and a stable environment, the company you work for might not allow you to deploy containers with code containing known CVEs, no matter if they are reachable from the internet or not. And some teams simply don't have the infrastructure to support a setup like that. ↩

2. While the ML centric ecosystem of packages is not as exposed security-wise as something like httpx, transitive dependencies have a tendency to force updates of seemingly unrelated packages. Not to mention internal packages that are shared between several projects which might need to be updated for various reasons. ↩

3. An example could be that the SimpleImputer class is replaced with a scikit-learn Pipeline wrapping the SimpleImputer, or a custom model object with features that the downstream consumers of the code rely on, and the loader method can ensure that we patch the model with sane defaults, allowing older saved model artifacts to keep working. ↩