Pickles go stale: saving ML models that survive dependency upgrades
Introduction & Example of a model breaking
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 trained and where it's unpickled, and herein lies a culprit, as ML models seldom live in a vacuum1 – security updates especially might require the deployment environment to have its dependencies updated 2. 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.
I'd argue that this problem becomes more manageable if we break down the large pickle files into smaller ones, and replace the pickle-components with version-agnostic formats where we can... but first, 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:
# /// 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)
- Run this snippet using
uv run imputer_save.pyto use this exact version of scikit-learn.
# /// 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():
# /// 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 such are only used during 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 is also saved alongside the model and list of features, and is used by the loading logic to identify which version of the wrapper saved a particular artifact, and thus determine whether any fixes are needed to make the code run with evolutions of the model3. Saving the version of the package that the ML model belongs to (scikit-learn in this case) could also be beneficial, but it's omitted from this example for the sake of brevity.
import base64
import json
import pickle
from pathlib import Path
from typing import Protocol, Self
import pandas as pd
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()
VERSION = 1
def __init__(self, imputer: SupportsTransform, features: list[str]) -> None:
self.imputer = imputer
self.features = features
def transform(self, X: pd.DataFrame) -> pd.DataFrame:
XX = X[self.features]
return self.imputer.transform(XX)
def save(self, path: Path) -> None:
serialized = dict(
features=self.features,
model_pickle=base64.b64encode(pickle.dumps(self.imputer)).decode(),
version=self.VERSION,
)
path.write_text(json.dumps(serialized))
@classmethod
def load(cls, path: Path) -> Self:
raw = json.loads(path.read_bytes())
version = raw["version"]
if version > cls.VERSION:
raise ValueError(
f"Model saved with format version {version}, "
f"but this code only supports up to version {cls.VERSION}"
)
imputer = pickle.loads(base64.b64decode(raw["model_pickle"]))
# This special case handling not be in the first version of this class,
# but as the code evolves, we can handle backwards compatability issues
# like this.
if version == 1:
imputer._fill_dtype = imputer._fit_dtype
return cls(
imputer=imputer,
features=raw["features"],
)
# /// script
# dependencies = ["scikit-learn==1.7.2", "pandas"]
# ///
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"))
# /// script
# dependencies = ["scikit-learn==1.8.0", "pandas"]
# ///
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:
# [[2.]]
Use AI for this
Implementing and testing a wrapper class like this is very well defined, and 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.
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:
# /// 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.
import base64
import json
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()
VERSION = 1
def __init__(
self, onnx_model: bytes, features: list[str], dtype: str = DTYPE
) -> None:
self.onnx_model = onnx_model
self.features = 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
return pd.DataFrame(output, columns=self.features, index=X.index)
def save(self, path: str | Path) -> None:
serialized = {
"features": self.features,
"version": self.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:
raw = json.loads(Path(path).read_text())
version = raw["version"]
if version > cls.VERSION:
raise ValueError(
f"Model saved with format version {version}, "
f"but this code only supports up to version {cls.VERSION}"
)
return cls(
onnx_model=base64.b64decode(raw["onnx_model"]),
features=raw["features"],
dtype=raw["dtype"],
)
Using the model above becomes
# /// 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")
# /// 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 or ONNX?
Now, while ONNX is great, it can't encode all models, which is why we might have to fall back to using pickle files. Additionally, there can be small numerical differences in the predictions between the pickle files and the ONNX re-implementation of the model, though these usually will be 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 implement 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.
-
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. ↩
-
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. ↩
-
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. ↩