Migrer des versions de modèle du Registre de modèles d’espace de travail vers le catalogue Unity

Databricks recommande d’utiliser des modèles dans Unity Catalog pour une meilleure gouvernance, un partage facile entre les espaces de travail et les environnements, et des workflows MLOps plus flexibles. Pour migrer des versions de modèle à partir du registre de modèles d’espace de travail vers le catalogue Unity, Databricks recommande d’utiliser copy_model_version() avec le client >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 le modèle de destination n’existe pas dans le catalogue Unity, il est créé par cet appel d’API.

Signatures de modèle

Les modèles du catalogue Unity nécessitent une signature. Si la version du modèle d’espace de travail n’a pas de signature, Databricks vous recommande de en créer une en suivant les instructions de la documentation MLflow.

Pour simplifier la migration, vous pouvez utiliser la variable MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATIONd’environnement. Cette variable d’environnement est disponible uniquement lorsque vous utilisez copy_model_version() et que vous avez besoin de la version 3.4.0 MLflow ou ultérieure. Lorsque cette variable d’environnement est définie "true"sur , une signature n’est pas requise.

Les versions de modèle inscrites sans signatures ont des limitations. Consultez Ajouter ou mettre à jour une signature pour une version de modèle existante.

import os

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

Pour ajouter une signature à une version de modèle existante, consultez la documentation MLflow.

Exemple de script pour migrer des versions de modèle vers un modèle de catalogue Unity

Le script suivant montre comment migrer toutes les versions de modèle de votre espace de travail enregistrées vers un modèle de catalogue Unity de destination. Ce script suppose que vous avez défini la variable MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION d’environnement sur « true » comme décrit dans les signatures de modèle.

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