将模型版本从工作区模型注册表迁移到 Unity 目录

Databricks 建议在 Unity Catalog 中使用模型，以增强治理、在工作区和环境之间轻松共享，以及实现更灵活的 MLOps 工作流。 若要将模型版本从工作区模型注册表迁移到 Unity 目录，Databricks 建议使用copy_model_version()与MLflow客户端>=3.4.0：

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

如果 Unity 目录中不存在目标模型，则此 API 调用会创建该模型。

模型签名

Unity 目录中的模型需要签名。 如果工作区模型版本没有签名，Databricks 建议按照 MLflow 文档中的说明创建一个。

为了简化迁移，可以使用环境变量 MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION。 此环境变量仅在使用 copy_model_version() 且需要 MLflow 版本 3.4.0 或更高版本时可用。 如果此环境变量设置为 "true"，则不需要签名。

在没有签名的情况下注册的模型版本有限制。 请参阅 添加或更新现有模型版本的签名。

import os

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

若要将签名添加到现有模型版本，请参阅 MLflow 文档。

将模型版本迁移到 Unity 目录模型的示例脚本

以下脚本演示如何将工作区注册模型中的所有模型版本迁移到目标 Unity 目录模型。 此脚本假定已将环境变量 MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION 设置为“true”，如 模型签名中所述。

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