Upgrade ML workflows to target models in Unity Catalog

This article explains how to migrate and upgrade existing Databricks workflows to use models in Unity Catalog.

Requirements

Required privileges

To execute a model training, deployment, or inference workflow in Unity Catalog, the principal running the workflow must have USE CATALOG and USE SCHEMA privileges on the catalog and schema that hold the model.

The following privileges are also required:

• To create a model, the principal must have the CREATE MODEL privilege.
• To load or deploy a model, the principal must have the EXECUTE privilege on the registered model.

Only the owner of a registered model can do the following:

• Create a new model version.
• Set an alias on a registered model.

Compute requirements

The compute resource specified for the workflow must have access to Unity Catalog. See Access modes.

Create parallel training, deployment, and inference workflows

To upgrade model training and inference workflows to Unity Catalog, Databricks recommends an incremental approach in which you create a parallel training, deployment, and inference pipeline that leverage models in Unity Catalog. When you’re comfortable with the results using Unity Catalog, you can switch downstream consumers to read the batch inference output, or increase the traffic routed to models in Unity Catalog in serving endpoints.

Model training workflow

Clone your model training workflow. Confirm that the principal running the workflow and the compute specified for the workflow meet the Requirements.

Next, modify the model training code in the cloned workflow. You might need to clone the notebook run by the workflow, or create and target a new git branch in the cloned workflow. Follow these steps to install the necessary version of MLflow and configure the client to target Unity Catalog in your training code. Then, update the model training code to register models to Unity Catalog. See Train and register Unity Catalog-compatible models.

Model deployment workflow

Clone your model deployment workflow. Confirm that the principal running the workflow and the compute specified for the workflow meet the Requirements.

If you have model validation logic in your deployment workflow, update it to load model versions from UC. Use aliases to manage production model rollouts.

Model inference workflow

Batch inference workflow

Clone the batch inference workflow. Confirm that the principal running the workflow and the compute specified for the workflow meet the Requirements.

Model serving workflow

If you are using Mosaic AI Model Serving, you don’t need to clone your existing endpoint. Instead, use the traffic split feature to start routing a small fraction of traffic to models in Unity Catalog. As you review the results using Unity Catalog, increase the amount of traffic until all of the traffic is rerouted.

Promote a model across environments

Promoting a model across environments works differently with models in Unity Catalog. For details, see Promote a model across environments.

Use job webhooks for manual approval for model deployment

Databricks recommends that you automate model deployment if possible, using appropriate checks and tests during the model deployment process. However, if you do need to perform manual approvals to deploy production models, you can use job notifications to call out to external CI/CD systems to request manual approval for deploying a model, after your model training job completes successfully. After manual approval is provided, your CI/CD system can then deploy the model version to serve traffic, for example by setting the “Champion” alias on it.