Problem

You have a registered MLflow model in your workspace. You want to generate predictions on a PySpark DataFrame and save the results to a table. You use MLflow in a notebook on an all-purpose cluster to load the model as an Apache Spark user-defined function (UDF) and encounter the following error. 

Code that produces the error

loaded_model = mlflow.pyfunc.spark_udf(spark, model_uri= model_uri, result_type="float")

df_result = df.withColumn('predictions', loaded_model(struct(*map(col, df.columns))))

df_result.write.mode('overwrite').option("mergeSchema", "true").saveAsTable('<table-name>')

Stack trace for the error

The cluster used to log the model has a different Databricks Runtime version (15.4 LTS ML) than the one used to load the model (14.3 LTS ML). 

PythonException: 
  An exception was thrown from the Python worker. Please see the stack trace below.
Traceback (most recent call last):
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-XXXXX/lib/python3.10/site-packages/mlflow/pyfunc/__init__.py", line 1914, in udf
    loaded_model = mlflow.pyfunc.load_model(local_model_path)
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-XXXXX/lib/python3.10/site-packages/mlflow/pyfunc/__init__.py", line 857, in load_model
    except ModuleNotFoundError as e:
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-XXXXX/lib/python3.10/site-packages/mlflow/pyfunc/model.py", line 468, in _load_pyfunc
    context, python_model, signature = _load_context_model_and_signature(model_path, model_config)
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-XXXXX/lib/python3.10/site-packages/mlflow/pyfunc/model.py", line 450, in _load_context_model_and_signature
    python_model = cloudpickle.load(f)
TypeError: code expected at most 16 arguments, got 18

Cause

The cloudpickle package and Python version you used to log and register the model differ from the versions present in the Databricks Runtime used to load the model. 

Cloudpickle can only be used to send objects between the exact same version of Python. For details, refer to the cloudpickle documentation.

Solution

When you load the model using MLflow to generate single predictions or as a Spark UDF for batch inference, ensure that you’re using a cluster with a Databricks Runtime that matches the Python and cloudpickle versions used to log and register the model.

To check for correspondence, either run the !python --version command in a notebook cell or run the following code snippet in a notebook cell. 

import sys
print(sys.version)