Unable to register Scikit-learn or XGBoost model to unity catalog

[AlkaSaliss]

Hello, I'm following the tutorial provided here https://docs.databricks.com/aws/en/notebooks/source/mlflow/mlflow-classic-ml-e2e-mlflow-3.html for ML model management process using ML FLow, in a unity-catalog enabled workspace, however I'm facing an error I'm unable to understand and fix.

The error is specifically occurring in the section 4 of the tutorial:  Log the model using MLflow:

# Incorporate MLflow evaluation
evaluation_data = X_test.copy()
evaluation_data["label"] = y_test

# Log the model and training metadata results
with mlflow.start_run(run_name="my_xgboost_regression") as run:
    # Extract metrics
    final_train_rmse = np.array(reg.evals_result()["validation_0"]["rmse"])[-1]
    final_test_rmse = np.array(reg.evals_result()["validation_1"]["rmse"])[-1]

    # Extract parameters for logging
    feature_map = {key: value for key, value in reg.get_xgb_params().items() if value is not None}

    # Generate a model signature using the infer_signature utility in MLflow
    # A signature is required to register the model to Unity Catalog 
    # so that the model can be used in SQL queries
    signature = infer_signature(X, reg.predict(X))

    # Log parameters
    mlflow.log_params(feature_map)

    # Log the model to MLflow and register the model to Unity Catalog
    # All model metrics and parameters will be available in Unity Catalog
    model_info = mlflow.xgboost.log_model(
        xgb_model=reg,
        name="xgboost_regression_model",
        input_example=X.iloc[[0]],
        signature=signature,
        registered_model_name="my_catalog.y_schema.xgboost_regression_model",
    )

Here's the stack trace of model logging error:

🔗 View Logged Model at: https://dbc-my-ws-account.cloud.databricks.com/ml/experiments/2706846431648976/models/m-782e22e859f24a63a29e0d9ee9227a6a?o=1366243967608557
/local_disk0/.ephemeral_nfs/envs/pythonEnv-ea9e38f3-62ca-47d0-b399-bd069dfe7414/lib/python3.12/site-packages/xgboost/sklearn.py:1028: UserWarning: [08:35:04] WARNING: /workspace/src/c_api/c_api.cc:1427: Saving model in the UBJSON format as default.  You can use file extension: `json`, `ubj` or `deprecated` to choose between formats.
  self.get_booster().save_model(fname)
Successfully registered model 'my_catalag.my_schema.xgboost_regression_model'.

MlflowException: The following failures occurred while uploading one or more artifacts to s3://my-s3-catalog-bucket/models/95d12042-52cb-4348-a05f-4bbdfe0c71f1/versions/814a5856-83df-4ed2-b4df-dc3dfc5c34ad: {'/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/serving_input_example.json': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')", '/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/conda.yaml': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')", '/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/requirements.txt': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')", '/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/MLmodel': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')", '/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/input_example.json': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')", '/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/model.xgb': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')", '/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/python_env.yaml': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')", '/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/metadata/conda.yaml': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')", '/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/metadata/requirements.txt': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')", '/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/metadata/MLmodel': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')", '/local_disk0/repl_tmp_data/ReplId-198e0-2fc1c-a/tmp4w84_er0/model/metadata/python_env.yaml': "JSONDecodeError('Expecting value: line 1 column 1 (char 0)')"}
File <command-2706846431648989>, line 24
     20 mlflow.log_params(feature_map)
     22 # Log the model to MLflow and register the model to Unity Catalog
     23 # All model metrics and parameters will be available in Unity Catalog
---> 24 model_info = mlflow.xgboost.log_model(
     25     xgb_model=reg,
     26     name="xgboost_regression_model",
     27     input_example=X.iloc[[0]],
     28     signature=signature,
     29     registered_model_name="my_catalog.my_schema.xgboost_regression_model",
     30 )
     32 # Log metrics to the run and model
     33 mlflow.log_metric("train_rmse", final_train_rmse)
File <command-2706846431648989>, line 24
     20 mlflow.log_params(feature_map)
     22 # Log the model to MLflow and register the model to Unity Catalog
     23 # All model metrics and parameters will be available in Unity Catalog
---> 24 model_info = mlflow.xgboost.log_model(
     25     xgb_model=reg,
     26     name="xgboost_regression_model",
     27     input_example=X.iloc[[0]],
     28     signature=signature,
     29     registered_model_name="my_catalog.my_schema.xgboost_regression_model",
     30 )
     32 # Log metrics to the run and model
     33 mlflow.log_metric("train_rmse", final_train_rmse)
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ea9e38f3-62ca-47d0-b399-bd069dfe7414/lib/python3.12/site-packages/mlflow/xgboost/__init__.py:279, in log_model(xgb_model, artifact_path, conda_env, code_paths, registered_model_name, signature, input_example, await_registration_for, pip_requirements, extra_pip_requirements, model_format, metadata, name, params, tags, model_type, step, model_id, **kwargs)
    225 @format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name=FLAVOR_NAME))
    226 def log_model(
    227     xgb_model,
   (...)
    245     **kwargs,
    246 ):
    247     """Log an XGBoost model as an MLflow artifact for the current run.
    248 
    249     Args:
   (...)
    277         metadata of the logged model.
    278     """
--> 279     return Model.log(
    280         artifact_path=artifact_path,
    281         name=name,
    282         flavor=mlflow.xgboost,
    283         registered_model_name=registered_model_name,
    284         xgb_model=xgb_model,
    285         model_format=model_format,
    286         conda_env=conda_env,
    287         code_paths=code_paths,
    288         signature=signature,
    289         input_example=input_example,
    290         await_registration_for=await_registration_for,
    291         pip_requirements=pip_requirements,
    292         extra_pip_requirements=extra_pip_requirements,
    293         metadata=metadata,
    294         params=params,
    295         tags=tags,
    296         model_type=model_type,
    297         step=step,
    298         model_id=model_id,
    299         **kwargs,
    300     )
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ea9e38f3-62ca-47d0-b399-bd069dfe7414/lib/python3.12/site-packages/mlflow/models/model.py:1357, in Model.log(cls, artifact_path, flavor, registered_model_name, await_registration_for, metadata, run_id, resources, auth_policy, prompts, name, model_type, params, tags, step, model_id, **kwargs)
   1354             _logger.warning("Failed to log model config as params: %s", str(e))
   1356 if registered_model_name is not None:
-> 1357     registered_model = mlflow.tracking._model_registry.fluent._register_model(
   1358         f"models:/{model.model_id}",
   1359         registered_model_name,
   1360         await_registration_for=await_registration_for,
   1361         local_model_path=local_path,
   1362     )
   1363 model_info = mlflow_model.get_model_info(model)
   1364 if registered_model is not None:
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ea9e38f3-62ca-47d0-b399-bd069dfe7414/lib/python3.12/site-packages/mlflow/tracking/_model_registry/fluent.py:213, in _register_model(model_uri, name, await_registration_for, tags, local_model_path, env_pack)
    207     with pack_env_for_databricks_model_serving(
    208         model_uri,
    209         enforce_pip_requirements=True,
    210     ) as artifacts_path_with_env:
    211         client.log_model_artifacts(model_id, artifacts_path_with_env)
--> 213 create_version_response = client._create_model_version(
    214     name=name,
    215     source=source,
    216     run_id=run_id,
    217     tags=tags,
    218     await_creation_for=await_registration_for,
    219     local_model_path=local_model_path,
    220     model_id=model_id,
    221 )
    222 created_message = (
    223     f"Created version '{create_version_response.version}' of model "
    224     f"'{create_version_response.name}'"
    225 )
    226 # Print a link to the UC model version page if the model is in UC.
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ea9e38f3-62ca-47d0-b399-bd069dfe7414/lib/python3.12/site-packages/mlflow/tracking/client.py:4198, in MlflowClient._create_model_version(self, name, source, run_id, tags, run_link, description, await_creation_for, local_model_path, model_id)
   4195         # models:/<model_id> source is not supported by WSMR
   4196         new_source = logged_model.artifact_location
-> 4198 return self._get_registry_client().create_model_version(
   4199     name=name,
   4200     source=new_source,
   4201     run_id=run_id,
   4202     tags=tags,
   4203     run_link=run_link,
   4204     description=description,
   4205     await_creation_for=await_creation_for,
   4206     local_model_path=local_model_path,
   4207     model_id=model_id,
   4208 )
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ea9e38f3-62ca-47d0-b399-bd069dfe7414/lib/python3.12/site-packages/mlflow/telemetry/track.py:22, in record_usage_event.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
     19 @functools.wraps(func)
     20 def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
     21     if is_telemetry_disabled() or _is_telemetry_disabled_for_event(event):
---> 22         return func(*args, **kwargs)
     24     success = True
     25     start_time = time.time()
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ea9e38f3-62ca-47d0-b399-bd069dfe7414/lib/python3.12/site-packages/mlflow/tracking/_model_registry/client.py:266, in ModelRegistryClient.create_model_version(self, name, source, run_id, tags, run_link, description, await_creation_for, local_model_path, model_id)
    264 arg_names = _get_arg_names(self.store.create_model_version)
    265 if "local_model_path" in arg_names:
--> 266     mv = self.store.create_model_version(
    267         name,
    268         source,
    269         run_id,
    270         tags,
    271         run_link,
    272         description,
    273         local_model_path=local_model_path,
    274         model_id=model_id,
    275     )
    276 else:
    277     # Fall back to calling create_model_version without
    278     # local_model_path since old model registry store implementations may not
    279     # support the local_model_path argument.
    280     mv = self.store.create_model_version(
    281         name, source, run_id, tags, run_link, description, model_id=model_id
    282     )
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ea9e38f3-62ca-47d0-b399-bd069dfe7414/lib/python3.12/site-packages/mlflow/store/_unity_catalog/registry/rest_store.py:983, in UcModelRegistryStore.create_model_version(self, name, source, run_id, tags, run_link, description, local_model_path, model_id)
    978 model_version = self._call_endpoint(
    979     CreateModelVersionRequest, req_body, extra_headers=extra_headers
    980 ).model_version
    982 store = self._get_artifact_repo(model_version, full_name)
--> 983 store.log_artifacts(local_dir=local_model_dir, artifact_path="")
    984 finalized_mv = self._finalize_model_version(
    985     name=full_name, version=model_version.version
    986 )
    987 return model_version_from_uc_proto(finalized_mv)
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-ea9e38f3-62ca-47d0-b399-bd069dfe7414/lib/python3.12/site-packages/mlflow/store/artifact/cloud_artifact_repo.py:168, in CloudArtifactRepository.log_artifacts(self, local_dir, artifact_path)
    165             failed_uploads[src_file_path] = repr(e)
    167 if len(failed_uploads) > 0:
--> 168     raise MlflowException(
    169         message=(
    170             "The following failures occurred while uploading one or more artifacts"
    171             f" to {self.artifact_uri}: {failed_uploads}"
    172         )
    173     )

[Gecofer]

Hi @AlkaSaliss!

I ran into the same issue and found the root cause, I think it's related to the use of name instead of artifact_path when calling mlflow.xgboost.log_model(...). Even though the documentation says https://mlflow.org/docs/latest/api_reference/python_api/mlflow.xgboost.html:

artifact_path – Deprecated. Use name instead.

That only applies to the new unified MLflow API (Model.log(...)), not to flavor-specific loggers like mlflow.xgboost.log_model(). When using name="xgboost_regression_model", MLflow doesn't properly resolve the artifact path, which causes it to create empty artifact files (e.g., MLmodel, conda.yaml, input_example.json, etc.).

"JSONDecodeError('Expecting value: line 1 column 1 (char 0)')",

You need to replace the name parameter with artifact_path, like this:

model_info = mlflow.xgboost.log_model(
    xgb_model=reg,
    artifact_path="xgboost_regression_model",  <-- THIS
    input_example=X.iloc[[0]],
    signature=signature,
    registered_model_name="my_catalog.y_schema.xgboost_regression_model",
)

Hope this helps 🙂

Gema

[AlkaSaliss]

Thanks @Gecofer  for the hint, I've tried with artifact_path, but still facing the same error 😭.

Maybe the root cause is different in my case. I'll continue investigating

[gbhatia]

Maybe add missing: 

mlflow.set_tracking_uri("databricks")

mlflow.set_registry_uri("databricks")

[joelramirezai]

You need to ensure that your Unity Catalog catalog and schema already exist, that you have the necessary permissions to use them, and that you update the code to reference your own catalog and schema names. You must also run on a classic cluster with a Databricks Runtime that is compatible with MLflow 3.0 (DBR 17.x or higher). Finally, specify the desired registered model name so it clearly reflects how you want the model to appear in your catalog.

I performed tests in DBR 18.0 ML personal compute cluster