Register Models in Unity Catalog with the MLflow Client API

1.Pointing the MLflow client at Unity Catalog

Databricks offers two model registries. The legacy Workspace Model Registry belongs to a single workspace. Models in Unity Catalog is a hosted MLflow Model Registry that brings Unity Catalog governance to models: centralized access control, auditing, lineage, and model discovery across workspaces. It also works with the open-source MLflow Python client, so the calls you use to register a model are the standard MLflow ones. The only question is which registry those calls reach.

The target registry is controlled by the registry URI. If the workspace's default catalog is in Unity Catalog, and you are on a cluster running Databricks Runtime 13.3 LTS or above or you are using MLflow 3, models are created in and loaded from that default catalog automatically. MLflow 3 also changes the default registry URI to databricks-uc. In every other case you set the URI yourself at the top of the notebook. The reverse also applies: a team that wants to stay on the legacy registry sets the URI to databricks.

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

import mlflow
mlflow.set_registry_uri(" ? ")

Sources12

2.What must be in place before you register

Pointing the client at Unity Catalog is necessary but not enough. A registration in Unity Catalog has to pass three gates: the compute, the caller's privileges, and the model itself.

Compute. The cluster must have access to Unity Catalog. For ML workloads, that means Dedicated access mode (formerly single user). On Databricks Runtime 15.4 LTS ML and above, dedicated group access mode also works. Databricks Runtime 13.2 ML and above includes support for models in Unity Catalog. On Databricks Runtime 11.3 LTS and above, you can get it by installing the latest MLflow client with %pip install --upgrade "mlflow-skinny[databricks]".

Privileges. A registered model is a Unity Catalog securable, so the normal catalog and schema privileges apply. Creating a new registered model and adding a version to an existing one need different privileges:

Signature. Every new model version in Unity Catalog must have a model signature. You can get one in either of two ways. Databricks autologging logs models with signatures automatically for many popular frameworks. With MLflow 2.5.0 and above, you can pass an input_example to mlflow.<flavor>.log_model, and MLflow infers the signature from it.

If you hit authorization errors during registration, the docs suggest setting the environment variable MLFLOW_USE_DATABRICKS_SDK_MODEL_ARTIFACTS_REPO_FOR_UC to True. On AWS GovCloud this setting is mandatory.

Sources1

3.Registering a logged model with mlflow.register_model

Once the prerequisites are met, registration is a single call. The docs call it the MLflow Client API register_model() method. It takes two arguments: a URI that points at a model you have already logged, and the three-level name of the registered model in Unity Catalog that the new version belongs to. The URI format depends on your MLflow major version, because MLflow 3 tracks models as first-class LoggedModel objects with their own IDs. MLflow 2.x addresses a model as an artifact path inside a run.

In a notebook you seldom type the ID by hand. In MLflow 3, mlflow.last_logged_model() returns the model you just logged, and its model_uri is ready to pass in. In MLflow 2.x, you build the runs:/ URI from the last active run, as the autologging example below shows.

autolog_run = mlflow.last_active_run()
model_uri = "runs:/{}/model".format(autolog_run.info.run_id)
mlflow.register_model(model_uri, "prod.ml_team.iris_model")

You do not have to create the registered model first. If the name does not exist, register_model creates the registered model and then adds the version to it. The caller still needs CREATE MODEL on the schema for that first registration. With MLflow 3 there is an extra benefit: the new version carries all the parameters and metrics logged with the source LoggedModel, and that performance data appears on the Unity Catalog model version page across all workspaces and experiments.

Sources132

4.Registering at log time, and from the UI

register_model is a separate step after logging. You can also log and register in one go by passing registered_model_name to the flavor's log_model call. The Unity Catalog naming rule still applies here: registered_model_name must be the full three-level name. The example below also covers the signature requirement from earlier. It passes the first training row as input_example, and MLflow infers the signature from that row and the model's prediction on it.

with mlflow.start_run():
    # Train a sklearn model on the iris dataset
    X, y = datasets.load_iris(return_X_y=True, as_frame=True)
    clf = RandomForestClassifier(max_depth=7)
    clf.fit(X, y)
    # Take the first row of the training dataset as the model input example.
    input_example = X.iloc[[0]]
    # Log the model and register it as a new version in UC.
    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",
    )

You can also register without code. On the experiment run page, the Register model button opens a dialog where you choose Unity Catalog as the destination and pick the target model. Registration is not instant, so you may need to refresh the destination model in Unity Catalog until the new version appears.

Sources1

5.What the three-level name does and does not tell you

All the examples register into prod.ml_team, which raises an obvious question: is version 2 of prod.ml_team.iris_model now serving production traffic? Not necessarily. The catalog, schema and registered model a version sits in describe its environment and the governance rules attached to it. For example, privileges can be set so that only admins can delete from the prod catalog. They do not describe its deployment status. Deployment status is tracked separately with model aliases, which belong to a later lesson on promotion. For registration, the takeaway is that choosing the catalog is a governance decision.

One more registration path uses an MlflowClient instance directly. Some versions already live in the legacy Workspace Model Registry. To bring them into Unity Catalog, Databricks recommends copy_model_version() with MLflow client 3.4.0 or later. The client points at the workspace registry as the source, and the destination is a three-level Unity Catalog name. Like register_model, the call creates the destination model if it does not exist.

import mlflow
from mlflow import MlflowClient

# Registry must be set to workspace registry
mlflow.set_registry_uri("databricks")
client = MlflowClient(registry_uri="databricks")

src_model_uri = f"models:/my_wmr_model/1"
uc_migrated_copy = client.copy_model_version(
   src_model_uri, "mycatalog.myschema.my_uc_model"
)

Sources14

Exam traps

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

1. 1.Calling mlflow.register_model on Databricks always registers to Unity Catalog.Why is that wrong?

   That is only automatic when the workspace's default catalog is in Unity Catalog (on DBR 13.3 LTS+ or MLflow 3). Otherwise the client writes to the workspace model registry unless you call mlflow.set_registry_uri("databricks-uc").

   Covered in Pointing the MLflow client at Unity Catalog

2. 2.A model signature is optional when registering to Unity Catalog, just as it was in the workspace registry.Why is that wrong?

   Every new model version in Unity Catalog must have a signature. Get one from autologging, or by passing input_example to log_model on MLflow 2.5.0+.

   Covered in What must be in place before you register

3. 3.registered_model_name in log_model can be a short name such as "iris_model", because only register_model needs the full path.Why is that wrong?

   Passing registered_model_name to log_model also registers to Unity Catalog, so it needs the full catalog.schema.model name too.

   Covered in Registering at log time, and from the UI

4. 4.Registering a version under the prod catalog puts it into production.Why is that wrong?

   The catalog reflects environment and governance rules, not deployment status. Deployment status is managed with model aliases.

   Covered in What the three-level name does and does not tell you