{"slug": "upgrading-mlflow-pipelines-across-a-python-version-boundary", "title": "Upgrading MLflow Pipelines Across a Python Version Boundary", "summary": "A team migrated production MLflow models from Databricks Runtime 14.3 (Python 3.10) to 17.3 (Python 3.12) without retraining, using a \"Load and Repack\" strategy to recover cloudpickle-serialized artifacts that failed with the error \"TypeError: code() argument 13 must be str, not int.\" The failure occurred because MLflow's cloudpickle serialization embeds Python version metadata and bytecode, which shifted between 3.10 and 3.12, breaking inference on models pulled from the registry. Retraining was not viable at the team's scale, since the product classification system alone uses multiple TB of training data and runs five model types per category — CatBoost, TF-IDF regression, a PyTorch Lightning neural network on all-mpnet-base-v2 embeddings, an image model, and a voting ensemble — governed in Databricks Unity Catalog.", "body_md": "# Upgrading MLflow Pipelines Across a Python Version Boundary\n\nHow 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.\n\nPicture 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:\n\n```\nTypeError: code() argument 13 must be str, not int.\n```\n\nThis 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).\n\nThis 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.\n\n## When cloudpickle meets a Python upgrade[¶](#when-cloudpickle-meets-a-python-upgrade)\n\nOur 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.\n\nHowever, 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.\n\nWhen 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.\n\n## Mapping our ML landscape[¶](#mapping-our-ml-landscape)\n\nTo 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.\n\nIn 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:\n\n- **CatBoost** : a tree-based classifier.\n- **TF-IDF Regression** : a logistic regression model operating on cleaned raw text.\n- **Neural Network (NN)** : a PyTorch Lightning model that processes dense vector representations from[`all-mpnet-base-v2`](https://huggingface.co/sentence-transformers/all-mpnet-base-v2) , an open-source sentence embedding model.\n- **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.\n- **Voting Ensemble** : a final logistic regression model that aggregates the predictions of the base models into a single, optimized binary tag assignment.\n\nOur 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.\n\nRetraining this entire matrix across all environments and categories from scratch would have consumed immense time and compute resources.\n\n## The \"Load and Repack\" strategy[¶](#the-load-and-repack-strategy)\n\nBypassing 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.\n\n### Stage 1: Load[¶](#stage-1-load)\n\nWe 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.\n\n### Stage 2: Repack[¶](#stage-2-repack)\n\nWith 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.\n\nBy 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.\n\nEach 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.\n\n## Future-proofing: long-term architectural wins[¶](#future-proofing-long-term-architectural-wins)\n\nWhile 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.\n\n### Standardized wrapper interfaces[¶](#standardized-wrapper-interfaces)\n\nPreviously, 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.\n\n### Self-contained model artifacts[¶](#self-contained-model-artifacts)\n\nIn 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](https://github.com/mlflow/mlflow/releases/tag/v3.11.1), MLflow auto-detects `uv.lock` and `pyproject.toml` and exports a precise, reproducible set of requirements at log time via `uv export`.\n\nBy 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.\n\n## Next iteration: eliminating pickle[¶](#next-iteration-eliminating-pickle)\n\nWhile 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:\n\n### Native serialization formats[¶](#native-serialization-formats)\n\nInstead of wrapping inner models, we can migrate to MLflow's native serialization flavors, which avoid `cloudpickle`:\n\n- `mlflow.catboost.log_model()` saves directly to CatBoost's native .cbm format\n- `mlflow.pytorch.log_model(..., serialization_format=\"pt2\")` utilizes`torch.export.save` (available in PyTorch ≥ 2.4)\n- `mlflow.sklearn.log_model(..., serialization_format=\"skops\")` secures TF-IDF and Voting ensemble models natively\n\n### Adopting \"Models from Code\"[¶](#adopting-models-from-code)\n\nTo 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.\n\n## Lessons learned[¶](#lessons-learned)\n\n**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.\n\n**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.\n\n**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.\n\nBy 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.\n\n*We're hiring! Do you like working in an ever evolving organization such as Zalando? Consider joining our teams as a [Machine Learning Engineer](https://jobs.zalando.com/en/jobs?category=Software+Engineering+-+Machine+Learning&utm_source=eng_blog&utm_content=upgrading-mlflow-pipelines)!*", "url": "https://wpnews.pro/news/upgrading-mlflow-pipelines-across-a-python-version-boundary", "canonical_source": "https://engineering.zalando.com/posts/2026/10/upgrading-mlflow-pipelines.html", "published_at": "2026-10-29 23:00:00+00:00", "updated_at": "2026-10-30 23:48:03.438479+00:00", "lang": "en", "topics": ["mlops", "machine-learning", "ai-infrastructure"], "entities": ["MLflow", "Databricks", "Python 3.10", "Python 3.12", "cloudpickle", "CatBoost", "PyTorch Lightning", "all-mpnet-base-v2"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/upgrading-mlflow-pipelines-across-a-python-version-boundary", "markdown": "https://wpnews.pro/news/upgrading-mlflow-pipelines-across-a-python-version-boundary.md", "text": "https://wpnews.pro/news/upgrading-mlflow-pipelines-across-a-python-version-boundary.txt", "jsonld": "https://wpnews.pro/news/upgrading-mlflow-pipelines-across-a-python-version-boundary.jsonld"}}