Migrera modellversioner från Workspace Model Registry till Unity Catalog

Databricks rekommenderar att du använder Modeller i Unity Catalog för bättre styrning, enkel delning mellan arbetsytor och miljöer och mer flexibla MLOps-arbetsflöden. För att migrera modellversioner från Arbetsytemodellregistret till Unity Catalog rekommenderar Databricks att du använder copy_model_version() med MLflow-klienten >= 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"
)

Om målmodellen inte finns i Unity Catalog skapas den av det här API-anropet.

Modellsignaturer

Modeller i Unity Catalog kräver en signatur. Om arbetsytemodellversionen inte har någon signatur rekommenderar Databricks att du skapar en genom att följa anvisningarna i MLflow-dokumentationen.

För att förenkla migreringen kan du använda miljövariabeln MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION. Den här miljövariabeln är endast tillgänglig när du använder copy_model_version() och kräver MLflow-version 3.4.0 eller senare. När den här miljövariabeln är inställd på "true"krävs ingen signatur.

Modellversioner som är registrerade utan signaturer har begränsningar. Se Lägga till eller uppdatera en signatur för en befintlig modellversion.

import os

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

Information om hur du lägger till en signatur i en befintlig modellversion finns i MLflow-dokumentationen.

Exempelskript för att migrera modellversioner till en Unity Catalog-modell

Följande skript visar hur du migrerar alla modellversioner i din registrerade arbetsyta till en Unity Catalog-målmodell. Det här skriptet förutsätter att du har angett miljövariabeln MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION till "true" enligt beskrivningen i Modellsignaturer.

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