Overview
This notebook demonstrates how to use model registry APIs to upgrade existing models in the workspace Model Registry to Unity Catalog.
Limitations
- The migration helpers in this notebook do not configure model permissions in Unity Catalog. Upgraded models are owned by the user running the notebook.
- Model versions in Unity Catalog do not support fixed stages. Instead, stages are migrated to aliases. The "Champion" alias is assigned to the latest model version in the "Production" stage (if any), and the "Challenger" alias is assigned to the latest model version in the "Staging" stage.
- Model versions in Unity Catalog must have a model signature, and model names in Unity Catalog must be valid Unity Catalog identifiers (AWS|Azure|GCP). The APIs in this notebook print the names of all model and versions that cannot be upgraded.
- Model version upgrade is not idempotent. If you rerun the migration helpers with the same destination schema, the notebook recreates model versions that already exist. Consider selecting a new destination schema name rather than attempting to resume a failed migration.
Requirements
- A cluster running Databricks Runtime 13.0 or above with access to Unity Catalog (AWS|Azure|GCP).
- Permissions to create models in at least one Unity Catalog schema. In particular, you need
USE CATALOG
permissions on the parent catalog, and bothUSE SCHEMA
andCREATE MODEL
permissions on the parent schema. If you get permissions errors while running the notebook, ask a catalog or schema owner or admin for access.
Define internal APIs for migration.
Define migration APIs
You can call these APIs individually to migrate individual registered models or their constituent versions.
Call helpers to migrate models to UC
Call upgrade_registered_model
on individual models to migrate them individually, or list models using mlflow.search_registered_models
and iterate through to migrate. For example, you can tag registered models you'd like to migrate, use search_registered_models
to find those models, and call upgrade_registered_model
to migrate each of them.
Migrate model versions without signature
When you migrate registered models to Unity Catalog using the helpers defined above, some model versions might not migrate because they lack a model signature.
You can manually migrate model versions without a signature by passing signature
to the
upgrade_model_version
helper function. For more information about how to construct a model signature, see the MLflow docs.
First, construct a model signature.
Then, call upgrade_model_version
, passing the signature
explicitly