Register a Model to Unity Catalog with MLflow

1.Models in Unity Catalog and where MLflow sends them

Once a chain or model works in a notebook, the next step toward serving it is to register it. Registration turns a logged artifact into a governed, versioned object. On Databricks, the recommended place for it is Models in Unity Catalog, a hosted version of the MLflow Model Registry. Because it lives in Unity Catalog, it gets centralized access control, auditing, lineage, and model discovery across workspaces. It also works with the open-source MLflow Python client, so you register models with the same MLflow calls you already use to log them.

A model in Unity Catalog sits inside a catalog and a schema, the same way a table does. That's why MLflow refers to it by a three-level name such as prod.ml_team.iris_model. If you call an API that expects a registered model name, such as mlflow.register_model, you pass that full name. There is one shortcut: if your workspace's default catalog is in Unity Catalog, you can pass just <model> and MLflow fills in the default catalog and schema. If Unity Catalog is enabled but the default catalog is not in Unity Catalog, you must spell out all three levels.

Whether you need to configure anything depends on your setup. If the workspace's default catalog is in Unity Catalog and you're on Databricks Runtime 13.3 LTS or above, or on MLflow 3, models are created in and loaded from that default catalog automatically. In other workspaces, the MLflow client writes to the legacy workspace model registry. To send it to Unity Catalog instead, set the registry URI explicitly:

import mlflow
mlflow.set_registry_uri("databricks-uc")

import mlflow
mlflow.set_registry_uri(" ? ")

Sources12

2.Compute and privileges you need before registering

Before you register anything, three conditions have to be met. First, Unity Catalog must be enabled in the workspace. Second, the compute must have access to Unity Catalog, which for ML workloads means Dedicated access mode (formerly single user). On Databricks Runtime 15.4 LTS ML and above, dedicated group access mode also works. Third, you need a runtime that supports it: Databricks Runtime 13.2 ML and above includes support. On Databricks Runtime 11.3 LTS and above, you can install the latest client with %pip install --upgrade "mlflow-skinny[databricks]".

Privileges come next, and the exam likes to test the difference between creating a registered model and adding a version to one. To create a registered model, you need CREATE MODEL and USE SCHEMA on the schema, plus USE CATALOG on the catalog. To add versions to an existing registered model, you must either own it or have CREATE MODEL VERSION on it, along with the same USE SCHEMA and USE CATALOG privileges. An admin can grant the schema-level privilege in Catalog Explorer or with SQL:

GRANT CREATE MODEL ON SCHEMA <schema-name> TO <principal>

If registration fails with an authorization error even though the privileges look right, the docs suggest setting the environment variable MLFLOW_USE_DATABRICKS_SDK_MODEL_ARTIFACTS_REPO_FOR_UC to True. On Databricks on AWS GovCloud, this setting is required rather than optional. It can't be used for models shared with OpenSharing that use default storage.

Sources1

3.Two ways to register: register_model() or registered_model_name

MLflow gives you two code paths. The first is to log the model and then register it in a separate step with mlflow.register_model(). It takes a model URI and the three-level name. In MLflow 3, the URI points to the logged model (models:/<model_id>). In MLflow 2.x, it points to a run artifact (runs:/<run_id>/model). If the registered model doesn't exist yet, MLflow creates it, then adds a new version to it.

logged_model = mlflow.last_logged_model()
mlflow.register_model(logged_model.model_uri, "prod.ml_team.iris_model")

The second path logs and registers in one call by passing registered_model_name to the flavor's log_model(). This is the same argument you'd pass whether the flavor is sklearn, langchain, or transformers. Because this registers straight into Unity Catalog, the value has to be the full three-level name.

mlflow.sklearn.log_model(
        sk_model=clf,
        name="model",
        # The signature is automatically inferred from the input example and its predicted output.
        input_example=input_example,
        # Use three-level name to register model in Unity Catalog.
        registered_model_name="prod.ml_team.iris_model",
    )

Sources1

4.Every new version needs a signature

This is the most common reason registration fails. The legacy workspace registry accepted models without signatures, but Unity Catalog does not accept new versions without one. There are two easy ways to get one. Databricks autologging logs signatures automatically for many popular frameworks. With MLflow 2.5.0 and above, you can pass input_example to mlflow.<flavor>.log_model and MLflow infers the signature from that example and the model's prediction on it. That's what the input_example line in the log_model() example above does: it uses one row of training data.

Some older model versions have no signature, and they come with limitations. You can add or update a signature on an existing version through the MLflow documentation's procedure.

Sources1

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.On MLflow 3 you always have to call mlflow.set_registry_uri("databricks-uc") before registering, or the model goes to the legacy registry.Why is that wrong?

   In MLflow 3 the default registry URI is already databricks-uc. You only need set_registry_uri in workspaces that still default to the workspace registry, or if you deliberately want the legacy registry, in which case you set it to databricks.

   Covered in Models in Unity Catalog and where MLflow sends them

2. 2.CREATE MODEL on a schema lets you add versions to any registered model in that schema.Why is that wrong?

   CREATE MODEL covers creating new registered models. Adding a version to an existing one requires owning it or holding CREATE MODEL VERSION on it, plus USE SCHEMA and USE CATALOG.

   Covered in Compute and privileges you need before registering

3. 3.A signature is optional metadata, so a chain logged without one can still be registered as a new Unity Catalog version.Why is that wrong?

   Unity Catalog requires a signature on new model versions. Log with an input_example (MLflow 2.5.0+) or use autologging so one is inferred.

   Covered in Every new version needs a signature