upgrade-models-to-unity-catalog(Python)

Loading...

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 both USE SCHEMA and CREATE 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.

      Configure catalog and schema

      To create a new catalog or schema, use CREATE CATALOG (AWS|Azure|GCP) or CREATE SCHEMA (AWS|Azure|GCP), respectively.

      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