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

aswanson
New Contributor

‎02-24-2025 02:45 PM

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.

0 Kudos

mark_ott
Databricks Employee
Databricks Employee

‎11-06-2025 03:52 AM

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

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

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

  1. 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.

0 Kudos