Migrar versões de modelo do Registro de Modelo do Workspace para o Catálogo do Unity

O Databricks recomenda usar os modelos no Catálogo do Unity para governança aprimorada, fácil compartilhamento entre os espaços de trabalho e ambientes e fluxos de trabalho MLOps mais flexíveis. Para migrar versões de modelo do Registro de Modelo do Workspace para o Catálogo do Unity, o Databricks recomenda usar copy_model_version() com o cliente >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"
)

Se o modelo de destino não existir no Catálogo do Unity, ele será criado por essa chamada à API.

Assinaturas de modelo

Os modelos no Catálogo do Unity exigem uma assinatura. Se a versão do modelo de workspace não tiver uma assinatura, o Databricks recomendará que você crie uma seguindo as instruções na documentação do MLflow.

Para simplificar a migração, você pode usar a variável MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATIONde ambiente. Essa variável de ambiente só está disponível quando você usa copy_model_version() e requer a versão 3.4.0 do MLflow ou superior. Quando essa variável de ambiente é definida como "true", uma assinatura não é necessária.

As versões de modelo registradas sem assinaturas têm limitações. Consulte Adicionar ou atualizar uma assinatura para uma versão de modelo existente.

import os

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

Para adicionar uma assinatura a uma versão de modelo existente, consulte a documentação do MLflow.

Script de exemplo para migrar versões de modelo para um modelo do Catálogo do Unity

O script a seguir mostra como migrar todas as versões do modelo em seu modelo registrado do workspace para um modelo de catálogo do Unity de destino. Esse script pressupõe que você tenha definido a variável MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION de ambiente como "true", conforme descrito nas assinaturas de modelo.

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