Migrar versões de modelos do Workspace Model Registry para o Unity Catalog

Databricks MLOps recomenda o uso de modelos em Unity Catalog para melhorar a governança, facilitar o compartilhamento entre espaços e ambientes de trabalho e tornar mais flexível o fluxo de trabalho. Para migrar versões de modelos do Workspace Model Registry para o Unity Catalog, a Databricks recomenda usar copy_model_version() com o cliente MLflow >= 3.4.0:

Python

from mlflow import MLflowClient

# Registry must be set to workspace registry
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 Unity Catalog, ele será criado por essa chamada de API.

Assinaturas de modelo

Os modelos no Unity Catalog exigem uma assinatura. Se a versão do modelo workspace não tiver uma assinatura, o site Databricks recomenda que o senhor crie uma seguindo as instruções da documentaçãoMLflow.

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

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.

Python

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.

Exemplo de script para migrar versões de modelo para um modelo do Unity Catalog

O script a seguir mostra como migrar todas as versões do modelo no seu modelo registrado workspace para um modelo de destino Unity Catalog. Esse script pressupõe que o senhor tenha definido a variável de ambiente MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION como "true", conforme descrito em Assinaturas de modelo.

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

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