Migrieren von Modellversionen aus der Arbeitsbereichsmodellregistrierung zum Unity-Katalog

Databricks empfiehlt die Verwendung von Modellen in Unity Catalog, um die Governance zu verbessern, die Freigabe über Arbeitsbereiche und Umgebungen hinweg zu erleichtern und flexiblere MLOps-Workflows zu ermöglichen. Um Modellversionen aus der Arbeitsbereichsmodellregistrierung zum Unity-Katalog zu migrieren, empfiehlt Databricks die Verwendung copy_model_version() mit dem MLflow-Client >= 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"
)

Wenn das Zielmodell nicht im Unity-Katalog vorhanden ist, wird es von diesem API-Aufruf erstellt.

Modellsignaturen

Modelle im Unity-Katalog erfordern eine Signatur. Wenn die Arbeitsbereichsmodellversion keine Signatur aufweist, empfiehlt Databricks, eine zu erstellen, indem Sie die Anweisungen in der MLflow-Dokumentation befolgen.

Um die Migration zu vereinfachen, können Sie die Umgebungsvariable MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATIONverwenden. Diese Umgebungsvariable ist nur verfügbar, wenn Sie copy_model_version() verwenden und MLflow-Version 3.4.0 oder höher erforderlich ist. Wenn diese Umgebungsvariable auf "true" festgelegt ist, ist keine Signatur erforderlich.

Modellversionen, die ohne Signaturen registriert sind, weisen Einschränkungen auf. Siehe Hinzufügen oder Aktualisieren einer Signatur für eine vorhandene Modellversion.

import os

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

Informationen zum Hinzufügen einer Signatur zu einer vorhandenen Modellversion finden Sie in der MLflow-Dokumentation.

Beispielskript zum Migrieren von Modellversionen zu einem Unity-Katalogmodell

Das folgende Skript zeigt, wie Sie alle Modellversionen in Ihrem Arbeitsbereich registrierten Modell zu einem Unity-Katalog-Zielmodell migrieren. Bei diesem Skript wird davon ausgegangen, dass Sie die Umgebungsvariable MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION auf "true" festgelegt haben, wie in den Modellsignaturen beschrieben.

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