I am using mlflow to register my custom model using a simple code below. The DatabricksParams extracts all the params from dbutils and sets the params dictionary and dbutils is not used anywhere else within the rest of my code base. The code fails when I call the mlflow.pyfunc.log_model(). Can you please help what might be causing this?

Exception:

An unexpected error occurred: You cannot use dbutils within a spark job You cannot use dbutils within a spark job or otherwise pickle it. If you need to use getArguments within a spark job, you have to get the argument before using it in the job. For example, if you have the following code: myRdd.map(lambda i: dbutils.args.getArgument("X") + str(i)) Then you should use it this way: argX = dbutils.args.getArgument("X") myRdd.map(lambda i: argX + str(i))

Code Snippet:

# COMMAND ----------
if __name__ == '__main__':
    params = DatabricksParams(dbutils)
    run_training(params)