Migrowanie wersji modelu z rejestru modeli obszarów roboczych do wykazu aparatu Unity

Usługa Databricks zaleca używanie modeli w Unity Catalog dla lepszego zarządzania, łatwiejszego udostępniania w różnych obszarach roboczych i środowiskach oraz bardziej elastycznych przepływów MLOps. Aby przeprowadzić migrację wersji modelu z rejestru modeli obszaru roboczego do Unity Catalog, usługa Databricks zaleca używanie copy_model_version() z 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"
)

Jeśli model docelowy nie istnieje w katalogu aparatu Unity, zostanie utworzony przez to wywołanie interfejsu API.

Podpisy modelu

Modele w wykazie aparatu Unity wymagają podpisu. Jeśli wersja modelu obszaru roboczego nie ma podpisu, usługa Databricks zaleca utworzenie go, postępując zgodnie z instrukcjami w dokumentacji platformy MLflow.

Aby uprościć migrację, możesz użyć zmiennej środowiskowej MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION. Ta zmienna środowiskowa jest dostępna tylko w przypadku użycia copy_model_version() i wymaga wersji MLflow lub nowszej 3.4.0 . Gdy ta zmienna środowiskowa jest ustawiona na "true"wartość , podpis nie jest wymagany.

Wersje modelu zarejestrowane bez podpisów mają ograniczenia. Zobacz Dodawanie lub aktualizowanie podpisu dla istniejącej wersji modelu.

import os

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

Aby dodać podpis do istniejącej wersji modelu, zobacz dokumentację platformy MLflow.

Przykładowy skrypt do migrowania wersji modelu do modelu wykazu aparatu Unity

Poniższy skrypt przedstawia sposób migrowania wszystkich wersji modelu w zarejestrowanym modelu obszaru roboczego do docelowego modelu wykazu aparatu Unity. Ten skrypt zakłada, że dla zmiennej środowiskowej MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION ustawiono wartość "true", zgodnie z opisem w artykule Podpisy 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")