Skip to content

ZALANDO 2026-10-29 Tier 2

Read original ↗

Upgrading MLflow Pipelines Across a Python Version Boundary

Summary

Zalando's product-classification team runs daily automated ML pipelines (data preparation, training, inference) on Databricks, across separate Dev, Staging and Prod workspaces deployed with Declarative Automation Bundles (DAB). A planned upgrade of the Databricks Runtime from LTS 14.3 (Python 3.10) to LTS 17.3 (Python 3.12) failed with TypeError: code() argument 13 must be str, not int. The root cause: MLflow logs Python models by cloudpickling the whole wrapper object, which embeds Python bytecode and version metadata in the artifact. Bytecode layout changed between 3.10 and 3.12, so the new interpreter could not load any model already in the registry. Retraining was not viable (multiple TB of training data; five model types per product category, across all environments). The team instead ran a one-time "Load and Repack" migration: load each model in an old-runtime (DBR 14.3) environment, extract the inner artifacts from the pyfunc wrapper, re-wrap them in a standard file-based wrapper class, relog with pinned dependencies, and move the @Production alias to the new version. Long-term fixes: a uniform wrapper interface, self-contained artifacts (code via code_paths, pinned deps via uv), and a plan to remove pickle entirely through native framework formats plus MLflow "Models from Code".

Key takeaways

  1. Pickle couples a stored artifact to the interpreter that wrote it. cloudpickle serializes code objects, so the artifact carries Python bytecode. A runtime upgrade across a bytecode change makes every stored model unreadable at once. This is a backward-compatibility failure of the storage format, not of the model code. (Source: sources/2026-10-29-zalando-upgrading-mlflow-pipelines-across-a-python-version-boundary)
  2. The failure is fleet-wide and happens at the training→serving handoff. Upgraded inference jobs failed as soon as they pulled any existing model from the registry. The registry is the boundary between training and serving, and its artifact format became the incompatible contract (see concepts/training-serving-boundary).
  3. Recompute cost decides the migration strategy. The ensemble has five model types per active category — CatBoost, TF-IDF logistic regression, a PyTorch Lightning NN on all-mpnet-base-v2 text embeddings, an image model (late fusion, configurable weight), and a voting-ensemble logistic regression. Heavy models train on GPU clusters, light ones on CPU clusters. Retraining this matrix in every environment was too expensive, so the team chose to migrate the artifacts.
  4. Load and Repack = read with the old runtime, write with a portable format. Stage 1 runs a temporary job on a DBR 14.3-matched environment with the old libraries installed, so the old pickles deserialize. Stage 2 unwraps the inner models, puts them in a file-based wrapper class (model_wrappers.py), relogs them to MLflow, and moves the @Production alias. Each model type needed its own extraction strategy. Dummy DataFrame inputs let the tool infer new MLflow model signatures.
  5. Alias-based promotion makes the cutover a pointer move. Models in Unity Catalog are addressed by named aliases (@Production), not numeric versions, and accessed through Service Principals. Consumers did not change: the repacked version got the alias.
  6. Standard wrapper interface. Before, models used framework-native MLflow flavors with ad-hoc wrapper logic. Now every model is a mlflow.pyfunc.PythonModel subclass (CatBoostModelWrapper, NNModelWrapper, …) with one uniform predict() interface.
  7. Self-contained artifacts. The wrapper module ships inside the artifact via code_paths; package versions are pinned. Since MLflow 3.11.1, MLflow detects uv.lock / pyproject.toml and exports exact requirements at log time via uv export. Each logged model then declares its own environment (related to, but weaker than, concepts/reproducible-builds).
  8. The wrapper itself is still pickled — next step removes pickle. Planned changes: native flavors (mlflow.catboost.log_model → .cbm; mlflow.pytorch.log_model(..., serialization_format="pt2") → torch.export.save, PyTorch ≥ 2.4; mlflow.sklearn.log_model(..., serialization_format="skops")), and Models from Code: MLflow stores the wrapper as Python source, and load_context() loads the native weights at runtime. With source code plus native weights, the artifact no longer depends on bytecode.
  9. Lessons as stated by the team. Coordinate library changes with Applied Scientists before you schedule runtime upgrades. Keep manifests current (pyproject.toml, uv sync --upgrade) so Python, PySpark and Delta Lake stay aligned. Prefer native formats over pickled Python objects.

Systems / concepts / patterns extracted

Operational numbers

  • Runtime: DBR LTS 14.3 (Python 3.10) → DBR LTS 17.3 (Python 3.12).
  • Training data for product classification: "multiple TB".
  • Model types per active product category: 5 (CatBoost, TF-IDF LR, NN, image model, voting ensemble).
  • Environments: 3 Databricks workspaces (Dev, Staging, Prod).
  • MLflow ≥ 3.11.1 for automatic uv.lock export; PyTorch ≥ 2.4 for torch.export-based pt2 format.
  • Migration job: run once per environment.

Caveats

  • The post gives no counts of models migrated, no migration runtime, and no compute-cost comparison against retraining.
  • The pickle-free design (native flavors + Models from Code) is planned, not yet in production.
  • The repacked wrapper still uses cloudpickle, so the current state reduces the risk but does not remove it.
  • Per-model-type extraction strategies are named but not described.

Source

Last updated · 854 distilled / 2,448 read