Step-by-Step Study Material

Step 1: Understanding the Core Problem

What is cloudpickle and why does MLflow use it?

When you save a machine learning model with MLflow using mlflow.pyfunc.log_model(), MLflow needs to freeze your entire Python object into a file so it can be loaded later.

It uses a library called cloudpickle to do this.

Think of cloudpickle like taking a photograph of your Python object:

Your Python Model Object
        │
        ▼
   cloudpickle
        │
        ▼
  Serialized File (.pkl)
  ┌─────────────────────┐
  │ Model weights       │
  │ Python bytecode     │  ← This is the problem
  │ Python version info │  ← This too
  │ Wrapper logic       │
  └─────────────────────┘

Why does a Python version upgrade break this?

Python compiles your code into bytecode — a lower-level set of instructions the interpreter executes.

When Python 3.12 tries to read a file serialized with Python 3.10 bytecode, it encounters instructions it cannot parse correctly.

Result:

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

Key insight: The model weights themselves are fine. The problem is the wrapper around them was frozen with version-specific bytecode.

---

Step 2: Understanding Why "Just Retrain" Was Not Viable

Before learning the solution, you need to appreciate the scale constraint that made retraining impractical.

The ML System Architecture at Zalando

For each product category, five model types run in parallel:

Product Category
       │
       ├── CatBoost (tree-based classifier)
       ├── TF-IDF Regression (text-based logistic regression)
       ├── Neural Network (PyTorch Lightning + sentence embeddings)
       ├── Image Model (PyTorch Lightning + image embeddings)
       └── Voting Ensemble (aggregates all above predictions)

Infrastructure Requirements Per Model

Training data: Multiple terabytes per run.

Retraining everything across Dev, Staging, and Production environments from scratch would consume enormous time and compute cost.

Key insight: The solution must recover existing trained weights, not discard them.

---

Step 3: The "Load and Repack" Strategy

This is the core technical solution. It has two distinct stages.

Stage 1: Load (Using the Old Environment)

The goal here is to read the broken pickle files using the same Python version that created them.

┌─────────────────────────────────────┐
│  Temporary Migration Job            │
│  Environment: DBR 14.3 (Python 3.10)│
│                                     │
│  Old pickle file ──► Load model     │
│                       │             │
│                       ▼             │
│              Extract raw weights    │
│              (CatBoost, PyTorch,    │
│               sklearn, etc.)        │
└─────────────────────────────────────┘

Why this works: You are using the same Python version that created the file, so the bytecode is compatible.

What you extract: The raw model weights, completely separated from the old pyfunc wrapper.

Stage 2: Repack (Using the New Environment)

Now you take those extracted weights and wrap them in a new, clean wrapper written in pure Python.

Raw Model Weights
       │
       ▼
New Clean Wrapper (model_wrappers.py)
  ┌────────────────────────────────┐
  │ class CatBoostModelWrapper     │
  │   (mlflow.pyfunc.PythonModel)  │
  │                                │
  │  - Pure Python logic           │
  │  - No old bytecode             │
  │  - Explicit dependencies       │
  └────────────────────────────────┘
       │
       ▼
mlflow.pyfunc.log_model()  ← Logged fresh with Python 3.12
       │
       ▼
Promoted to @Production alias

Key actions in Stage 2:

• Unwrap inner artifacts from the old pyfunc wrapper
• Create standardized wrapper classes in model_wrappers.py
• Re-log to MLflow with explicit, pinned dependencies
• Infer fresh MLflow model signatures using dummy DataFrames
• Promote the new version to the @Production alias

Key insight: You preserved the trained weights (expensive to recreate) while replacing only the serialization wrapper (cheap to recreate).

---

Step 4: Architectural Improvements Made Post-Migration

The migration was not just a fix — it prompted two lasting improvements.

Improvement 1: Standardized Wrapper Interfaces

Before migration (problematic pattern):

# Ad-hoc, inconsistent per model type
mlflow.catboost.log_model(model, ...)   # framework-native flavor
mlflow.pytorch.log_model(model, ...)    # different interface
# Each had different predict() behavior

After migration (standardized pattern):

# model_wrappers.py
class CatBoostModelWrapper(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        # consistent interface
        ...

class NNModelWrapper(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        # same interface, different internals
        ...

Benefit: Every model type now has a uniform predict() interface regardless of the underlying framework.

Improvement 2: Self-Contained Model Artifacts

Before migration:

MLflow Artifact
├── model.pkl          ← weights + wrapper (tightly coupled)
└── (no explicit deps) ← dependencies assumed from environment

After migration:

MLflow Artifact
├── model weights      ← weights only
├── model_wrappers.py  ← wrapper code (via code_paths)
├── pyproject.toml     ← explicit dependency manifest
└── uv.lock            ← pinned, reproducible dependencies

How this is achieved:

mlflow.pyfunc.log_model(
    artifact_path="model",
    python_model=CatBoostModelWrapper(),
    code_paths=["model_wrappers.py"],  # embed wrapper code
    # MLflow 3.11.1+ auto-detects uv.lock and pyproject.toml
)

Key insight: The artifact now carries everything it needs to run. It does not depend on the environment being configured correctly.

---

Step 5: Future-Proofing — Eliminating Pickle Entirely

Even after the migration, the new wrappers still use cloudpickle to serialize the wrapper object itself. The next evolution eliminates pickle completely.

Option A: Native Serialization Formats

Instead of wrapping models in a pyfunc, use MLflow's native flavors that save directly to framework-native formats:

These formats store only weights and configuration — no Python bytecode embedded.

Option B: Models from Code

This is the most robust approach for maintaining a unified predict() interface:

Traditional approach:
  Serialize Python object ──► .pkl file ──► Load object at runtime

Models from Code approach:
  Store Python source file ──► .py file ──► Execute source at runtime
                                            + load native weights separately

How it works conceptually:

# wrapper_script.py (stored as raw source, not pickled)
class MyModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # Load native weights at runtime
        self.model = catboost.CatBoostClassifier()
        self.model.load_model(context.artifacts["model_weights"])
    
    def predict(self, context, model_input):
        return self.model.predict(model_input)

At runtime, MLflow executes this file as pure source code. Since there is no pickle involved, Python version changes to bytecode format are completely irrelevant.

---

Step 6: Lessons Learned — Practical Takeaways

These are actionable principles derived from the experience:

Lesson 1: Coordinate Library Shifts Early

Before scheduling a runtime upgrade, review all package dependencies with your team. Major runtime upgrades change foundational libraries underneath you.

Lesson 2: Prioritize Rigorous Dependency Management

# Recommended workflow
uv sync --upgrade          # update dependencies
# Review pyproject.toml changes
# Verify Python, PySpark, Delta Lake compatibility

Lesson 3: Embrace Native Formats Over Pickles

Avoid:

Python Object → cloudpickle → .pkl file
(bytecode-dependent, version-fragile)

Prefer:

Model Weights → native format → .cbm / .pt2 / skops file
(bytecode-independent, version-resilient)

---

Summary: The Complete Mental Model

PROBLEM
Python 3.10 → 3.12 upgrade
cloudpickle embeds bytecode → TypeError on load

         ↓

IMMEDIATE SOLUTION: Load and Repack
Stage 1: Load with old Python 3.10 environment
         Extract raw weights from broken pickle
Stage 2: Repack with new Python 3.12 environment
         Create clean wrapper → re-log → promote to @Production

         ↓

ARCHITECTURAL IMPROVEMENTS
1. Standardized wrappers (model_wrappers.py)
2. Self-contained artifacts (code_paths + uv.lock)

         ↓

LONG-TERM GOAL: Eliminate Pickle
Option A: Native formats (.cbm, pt2, skops)
Option B: Models from Code (store source, not objects)

Core principle: The closer your ML artifacts stay to pure weights and configuration (rather than serialized Python objects), the more resilient they are to environment changes.