Migrate model versions from Workspace Model Registry to Unity Catalog

Databricks recommends using Models in Unity Catalog for improved governance, easy sharing across workspaces and environments, and more flexible MLOps workflows. To migrate model versions from the Workspace Model Registry to Unity Catalog, Databricks recommends using copy_model_version() with MLflow client >= 3.4.0:

Python

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"
)

If the destination model does not exist in Unity Catalog, it is created by this API call.

Model signatures

Models in Unity Catalog require a signature. If the workspace model version doesn't have a signature, Databricks recommends that you create one by following the instructions in the MLflow documentation.

To simplify migration, you can use the environment variable MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION. This environment variable is only available when you use copy_model_version() and requires MLflow version 3.4.0 or above. When this environment variable is set to "true", a signature is not required.

Model versions that are registered without signatures have limitations. See Add or update a signature for an existing model version.

Python

import os

os.environ["MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION"] = "true"

To add a signature to an existing model version, see the MLflow documentation.

Example script to migrate model versions to a Unity Catalog model

The following script shows how to migrate all of the model versions in your workspace registered model to a destination Unity Catalog model. This script assumes that you have set the environment variable MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION to "true" as described in Model signatures.

Python

import mlflow
from mlflow import MlflowClient
from mlflow.exceptions import MlflowException
from mlflow.models import ModelSignature
from mlflow.types.schema import Schema, ColSpec, AnyType

mlflow.set_registry_uri("databricks")
workspace_client = MlflowClient(registry_uri="databricks")
uc_client = MlflowClient(registry_uri="databricks-uc")


# Make a placeholder model that can be used to increment the version number
def make_placeholder_model() -> str:
    class _Placeholder(mlflow.pyfunc.PythonModel):
        def predict(self, ctx, x):
            return None

    with mlflow.start_run() as run:
        schema = Schema([ColSpec(AnyType())])
        model = mlflow.pyfunc.log_model(
            name="m",
            python_model=_Placeholder(),
            signature=ModelSignature(inputs=schema, outputs=schema),
        )
        return f"models:/{model.model_id}"


# Check if the source model has a particular version number
def workspace_model_exists(name: str, version: int) -> bool:
    try:
        workspace_client.get_model_version(name, str(version))
        return True
    except MlflowException as e:
        if e.error_code == "RESOURCE_DOES_NOT_EXIST":
            # Convert the RESOURCE_DOES_NOT_EXIST error into False
            return False
        # Raise all other exceptions
        raise e


# Copy model versions from a source Databricks workspace-registered model to
# a destination Databricks Unity Catalog registered model
def copy_model_versions_to_uc(src: str, dst: str) -> None:
    latest_versions = workspace_client.get_latest_versions(src)
    max_version_number = max(int(v.version) for v in latest_versions)
    placeholder_model = make_placeholder_model()

    for v in range(1, max_version_number + 1):
        if workspace_model_exists(src, v):
            workspace_client.copy_model_version(f"models:/{src}/{str(v)}", dst)
        else:
            # Create and immediately delete a placeholder model version to increment
            # the version counter on the UC model, so the version numbers on the UC
            # model match those on the workspace registered model.
            mv = uc_client.create_model_version(dst, placeholder_model)
            uc_client.delete_model_version(dst, mv.version)


copy_model_versions_to_uc("my_workspace_model", "mycatalog.myschema.my_uc_model")