Migración de versiones de modelo del Registro de modelos del área de trabajo al catálogo de Unity

Databricks recomienda usar Modelos en Unity Catalog para mejorar la gobernanza, facilitar el uso compartido entre áreas de trabajo y entornos, y flujos de trabajo de MLOps más flexibles. Para migrar versiones de modelo del Registro de modelos de Workspace al catálogo de Unity, Databricks recomienda usar copy_model_version() con el 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"
)

Si el modelo de destino no existe en el catálogo de Unity, se crea mediante esta llamada API.

Firmas de modelo

Los modelos del catálogo de Unity requieren una firma. Si la versión del modelo de área de trabajo no tiene una firma, Databricks recomienda crear una siguiendo las instrucciones de la documentación de MLflow.

Para simplificar la migración, puede usar la variable MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATIONde entorno . Esta variable de entorno solo está disponible cuando se usa copy_model_version() y se requiere la versión 3.4.0 de MLflow o superior. Cuando esta variable de entorno se establece en "true", no se requiere una firma.

Las versiones de modelo registradas sin firmas tienen limitaciones. Consulte Agregar o actualizar una firma para una versión de modelo existente.

import os

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

Para agregar una firma a una versión de modelo existente, consulte la documentación de MLflow.

Script de ejemplo para migrar versiones de modelo a un modelo de Catálogo de Unity

El siguiente script muestra cómo migrar todas las versiones del modelo del modelo registrado del área de trabajo a un modelo de catálogo de Unity de destino. Este script supone que ha establecido la variable MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION de entorno en "true", como se describe en Firmas 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")