Upgrading MLflow Pipelines Across a Python Version Boundary

How we migrated production MLflow models from Python 3.10 to 3.12 without retraining, using a "Load and Repack" strategy to recover cloudpickle-serialized artifacts and future-proof our ML architecture.

Picture this: you are running automated daily Machine Learning pipelines at scale, and it is time for a critical environment upgrade, such as moving to a new Python version. Consequently, you need to upgrade your Databricks Runtime from LTS 14.3 to LTS 17.3. However, the upgrade cannot proceed because the logs show a cryptic error:

TypeError: code() argument 13 must be str, not int.

This was the exact challenge our team faced recently during a Databricks upgrade that uses MLflow. At our core, we manage daily automated pipelines spanning data preparation, training, and inference across separate Dev, Staging, and Prod Databricks workspaces managed via Declarative Automation Bundles (DAB).

This post covers our journey identifying the root cause of Python version serialization failures, our "Load and Repack" technique to recover production models, and the architectural refinements we made to ensure our MLflow artifacts are forward-compatible.

When cloudpickle meets a Python upgrade¶ #

Our upgrade path was clear: transition our machine learning workflows from DBR 14.3 (Python 3.10) to DBR 17.3 (Python 3.12). Python 3.12 introduces excellent language features, such as PEP 695's simplified generic types. This syntax is fully embraced by Pydantic V2 to enforce data validity at runtime using type hints, which prompted us to implement more comprehensive type hints across our codebase.

However, modernizing the runtime introduced a breaking change for our legacy machine learning models in the registry. The root cause is how MLflow traditionally serializes Python models.

When MLflow logs a Python model, it uses cloudpickle to freeze the entire wrapper object, including inner models like CatBoost classifiers or PyTorch modules, into a single serialized artifact. Crucially, this process embeds Python version metadata and bytecode directly into that file. Because bytecode structures shifted significantly between Python 3.10 and 3.12, the new interpreter could no longer safely parse the older pickle files. This resulted in an immediate TypeError whenever the upgraded inference logic attempted to pull an existing model from the registry.

Mapping our ML landscape¶ #

To understand why a naive approach like "just retrain all the models" was not viable, consider the scale we operate at: our product classification system alone uses multiple TB of training data.

In our case, that system processes and classifies product data across a large and continuously growing catalogue. For every active product category, our system trains and runs five distinct model types in parallel:

• CatBoost : a tree-based classifier.
• TF-IDF Regression : a logistic regression model operating on cleaned raw text.
• Neural Network (NN) : a PyTorch Lightning model that processes dense vector representations fromall-mpnet-base-v2 , an open-source sentence embedding model.
• Image Model : another PyTorch Lightning classifier that processes dense image embeddings, trained only for image-eligible categories and integrated into the voting ensemble via late fusion with a configurable weight parameter.
• Voting Ensemble : a final logistic regression model that aggregates the predictions of the base models into a single, optimized binary tag assignment.

Our infrastructure uses Databricks to train these category models concurrently. Heavy NN, Image Model and CatBoost tasks run on dedicated GPU clusters, while TF-IDF and Voting tasks run on CPU clusters. Models are governed inside the Databricks Unity Catalog using secure Service Principals and strictly named aliases (like @Production) rather than numeric version tags.

Retraining this entire matrix across all environments and categories from scratch would have consumed immense time and compute resources.

The "Load and Repack" strategy¶ #

Bypassing the costly retraining phase, we developed a methodology for migrating and relogging our models to ensure operational continuity. We orchestrated this recovery across two foundational stages.

Stage 1: Load¶

We created a temporary, one-time migration job using an environment that matched DBR 14.3. By installing the old libraries directly into this environment, we avoided version conflicts. This allowed the system to easily read the problematic pickle files and separate the raw model weights from their restrictive pyfunc wrappers.

Stage 2: Repack¶

With the models loaded, we unwrapped the inner artifacts from the old pyfunc wrapper. We encapsulated the loaded model artifacts within a streamlined, file-based wrapper class in our model_wrappers.py module. This architectural shift prioritized pure Python logic, removing dependencies on outdated pickle structures to ensure native compatibility with Python 3.12.

By relogging this refined wrapper to MLflow with dependencies managed via uv, we preserved the trained artifacts while promoting the upgraded model back to the @Production alias.

Each of our model types originally used different serialization patterns, so our migration tool applied specialized extraction strategies per model type. With new dummy DataFrame inputs, we could also infer fresh MLflow model signatures.

Future-proofing: long-term architectural wins¶ #

While our migration tool was designed to run only once per environment, solving this immediate blocker yielded two architectural improvements for our production ML systems going forward.

Standardized wrapper interfaces¶

Previously, models were logged using framework-native MLflow flavors with ad-hoc wrapper logic. We introduced a dedicated model_wrappers.py file, establishing consistent mlflow.pyfunc.PythonModel subclasses (e.g. CatBoostModelWrapper, NNModelWrapper). This ensures a uniform predict() interface regardless of the underlying model type. Every model developed post-migration natively adopts this pattern.

Self-contained model artifacts¶

In our legacy workflow, MLflow artifacts lacked embedded code dependencies and explicit manifest files. To fix this, we pinned package versions explicitly and included model_wrappers.py directly via code_paths at log time. We have since adopted uv for dependency management. Since MLflow 3.11.1, MLflow auto-detects uv.lock and pyproject.toml and exports a precise, reproducible set of requirements at log time via uv export.

By embedding specialized dependencies and internal deserialization logic directly within the artifact, every logged model is now fully autonomous. This architectural shift ensures that future environment transitions or runtime upgrades proceed smoothly, eliminating the need for costly retraining or manual recovery efforts.

Next iteration: eliminating pickle¶ #

While our "Load and Repack" strategy solved the immediate DBR upgrade issue, encapsulating models inside custom PythonModel subclasses still relies on cloudpickle to serialize the wrapper object itself. To ensure we never face a version boundary issue again, we are evolving our architecture to eliminate pickle entirely through two advanced MLflow patterns:

Native serialization formats¶

Instead of wrapping inner models, we can migrate to MLflow's native serialization flavors, which avoid cloudpickle:

• mlflow.catboost.log_model() saves directly to CatBoost's native .cbm format
• mlflow.pytorch.log_model(..., serialization_format="pt2") utilizestorch.export.save (available in PyTorch ≥ 2.4)
• mlflow.sklearn.log_model(..., serialization_format="skops") secures TF-IDF and Voting ensemble models natively

Adopting "Models from Code"¶

To maintain our unified predict() interface without pickling the wrapper, we can transition to MLflow's Models from Code concept. Rather than logging an instantiated object, MLflow stores the wrapper as a raw Python source file. At runtime, the wrapper script uses the load_context() method to fetch the native model artifacts (like the .cbm file). Because the wrapper is executed as pure source code and the weights are native, the entire artifact becomes independent from future Python bytecode changes.

Lessons learned¶ #

Coordinate library shifts early. Upgrading a major runtime means foundational open-source libraries change under the hood as well. Always align early with your Applied Scientists to review respective package lists before scheduling data engineering upgrades.

Prioritize rigorous dependency management. Establish a disciplined workflow when configuring your environments. By thoroughly reviewing LTS documentation and modernizing your manifests, such as updating pyproject.toml and running uv sync --upgrade, you ensure that your Python, PySpark, and Delta Lake layers remain harmonized and resilient against system-wide changes.

Embrace native formats over pickles. Whenever possible, utilize native framework serializers (like .cbm or pt2) and avoid serializing Python objects. The closer your artifacts remain to pure configuration and native weights, the easier your next runtime upgrade will be.

By pivoting from a disruptive serialization bottleneck toward a unified architectural standard, future runtime transitions are now engineered to be as seamless as a simple version increment.

We're hiring! Do you like working in an ever evolving organization such as Zalando? Consider joining our teams as a Machine Learning Engineer!