@skosaraju 

When MLflow logs a PyFunc model, it needs to serialize (pickle) the model and its dependencies. The error occurs because dbutils is not serializable and cannot be pickled. Even though you're only using it in the DatabricksParams class to extract parameters, the entire instance of DatabricksParams (including any references to dbutils) might be getting captured in the closure and MLflow is trying to serialize it.
The key fix is to extract the parameters from dbutils before passing them to your model training function, rather than passing the dbutils object itself. Here's how to modify your code:

# COMMAND ----------
if __name__ == '__main__':
# Extract parameters as a dictionary BEFORE passing to your training function
params_dict = {
"param1": dbutils.widgets.get("param1") if dbutils.widgets.get("param1") else "default_value1",
"param2": dbutils.widgets.get("param2") if dbutils.widgets.get("param2") else "default_value2",
# Add other parameters as needed
}

# Pass the extracted parameters instead of dbutils
run_training(params_dict)

Then modify your DatabricksParams class to accept a dictionary instead of dbutils:

class DatabricksParams:
def __init__(self, params_dict):
self.params = params_dict
# Any other initialization

# Your existing methods that use self.params