Migrace verzí modelu z registru modelů pracovního prostoru do katalogu Unity

Databricks doporučuje používat modely v katalogu Unity pro vylepšené zásady správného řízení, snadné sdílení pracovních prostorů a prostředí a flexibilnější pracovní postupy MLOps. Pokud chcete migrovat verze modelu z registru modelů pracovního prostoru do katalogu Unity, Databricks doporučuje použít copy_model_version() s klientem 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"
)

Pokud cílový model v katalogu Unity neexistuje, vytvoří se tímto voláním rozhraní API.

Podpisy modelů

Modely v katalogu Unity vyžadují podpis. Pokud verze modelu pracovního prostoru nemá podpis, databricks doporučuje, abyste ho vytvořili podle pokynů v dokumentaci k MLflow.

Pro zjednodušení migrace můžete použít proměnnou MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATIONprostředí . Tato proměnná prostředí je dostupná jenom v případě, že používáte copy_model_version() a vyžaduje MLflow verzi 3.4.0 nebo vyšší. Pokud je tato proměnná prostředí nastavená na "true", podpis se nevyžaduje.

Verze modelů zaregistrované bez podpisů mají omezení. Viz Přidání nebo aktualizace podpisu pro existující verzi modelu.

import os

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

Pokud chcete přidat podpis do existující verze modelu, přečtěte si dokumentaci K MLflow.

Ukázkový skript pro migraci verzí modelu do modelu katalogu Unity

Následující skript ukazuje, jak migrovat všechny verze modelu v registrovaném modelu pracovního prostoru do cílového modelu katalogu Unity. Tento skript předpokládá, že jste nastavili proměnnou MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION prostředí na true, jak je popsáno v podpisech modelu.

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