Pickle/joblib.dump a pre-processing function defined in a notebook

aswanson — Mon, 24 Feb 2025 22:45:09 GMT

I've built a custom MLFlow model class which I know functions. As part of a given run the model class uses `joblib.dump` to store necessary parameters on the databricks DBFS before logging them as artifacts in the MLFlow run. This works fine when using functions defined within the libraries contained in the custom model class, but I run into SPARK-5063 CONTEXT_ONLY_VALID_ON_DRIVER errors if I use functions defined in the notebook in the model parameters.

This extends to trivial python functions defined in the notebook such as:

```

import joblib

def tmpfun(val😞

return val + 'bar'

joblib.dump(tmpfun, 'tmp.pkl')

```

It seems like the spark context is being injected into the function call or something, but I have no idea how to isolate the required functions such that they can be loaded later to rebuild the model.

Re: Pickle/joblib.dump a pre-processing function defined in a notebook

mark_ott — Thu, 06 Nov 2025 11:52:38 GMT

The error you’re seeing—SPARK-5063 CONTEXT_ONLY_VALID_ON_DRIVER—arises when trying to serialize or use objects (such as functions) defined in Databricks notebooks from workers rather than the driver. This issue is especially common with Python functions defined “dynamically” in interactive notebook cells, which can sometimes capture the Spark context or notebook environment. When joblib tries to pickle these, Spark context references may incorrectly be serialized, leading to errors when later deserialized or run in a different context.

Why Notebook-Defined Functions Cause This

Notebook Scope: Functions defined in notebooks are dynamically created in the current interpreter session. They often reside outside proper module namespaces, making serialization fraught.
Spark Context Capture: If the function (or its closure) refers in any way (even implicitly) to variables or objects from the notebook session (such as the Spark context object), it can lead to serialization issues because Spark contexts are only valid on the driver, not worker nodes.
Pickle Protocol Limitations: joblib uses pickle for serialization. Pickle can struggle with closures, lambda functions, or dynamically defined functions outside importable modules.

Best Practices for Serializing Functions

The safest solution to your issue is to ensure that any function (or object) you intend to serialize with joblib/joblib.dump is:

Defined in a separate Python module (i.e., in a .py file imported into the notebook), NOT inside notebook cells
Does NOT close over Spark context or other notebook-specific variables
Uses only pure-Python code with no references to ephemeral notebook objects

Example Solution

Create a Python module (e.g., mymodule.py😞

python

# mymodule.py
def tmpfun(val):
    return val + 'bar'

Import and use in your notebook:

python

import joblib
from mymodule import tmpfun

joblib.dump(tmpfun, 'tmp.pkl')  # This will work!

Why?

When you define tmpfun in mymodule.py, it's statically importable, and joblib/pickle can properly resolve its reference. There’s no unintended closure over notebook scope or Spark context.

Additional Strategies

Avoid Lambda Functions or Inner Functions: Top-level module functions are safest for serialization.
Custom Serialization: For very complex objects, define custom __getstate__ and __setstate__ methods.
Do Not Serialize Spark Context: If your function must use Spark, pass the context as an argument at the time of execution, not closure.
Check Imports in Model Class: Ensure all helpers are imported from python files.

Practical Debug Steps

Confirm your function is top-level in a module.
Use cloudpickle as an alternative to joblib; it handles interactive functions better but is still safer with modules.
Once serialized, test by deserializing in a fresh notebook session.

Reference Table

Function Location	joblib.dump Success?
Notebook cell	Often fails [SPARK-5063]
Python module import	Success
Lambda or closure	Often fails

Summary

To resolve SPARK-5063 and reliably rebuild your MLFlow model, move all serializable functions into static Python modules and import them into your notebook. Avoid closures, notebook cell definitions, and Spark context capture in functions you need to joblib.dump.

topic Re: Pickle/joblib.dump a pre-processing function defined in a notebook in Machine Learning