Eseguire la migrazione delle versioni del modello dall'area di lavoro al catalogo unity

Databricks consiglia di usare modelli in Unity Catalog per migliorare la governance, semplificare la condivisione tra aree di lavoro e ambienti e flussi di lavoro MLOps più flessibili. Per eseguire la migrazione delle versioni del modello dal Registro modelli nell'area di lavoro al Catalogo Unity, Databricks consiglia di usare copy_model_version() con il 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"
)

Se il modello di destinazione non esiste in Unity Catalog, viene creato da questa chiamata API.

Firme del modello

I modelli nel catalogo Unity richiedono una firma. Se la versione del modello dell'area di lavoro non ha una firma, Databricks consiglia di crearne una seguendo le istruzioni nella documentazione di MLflow.

Per semplificare la migrazione, è possibile usare la variabile MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATIONdi ambiente . Questa variabile di ambiente è disponibile solo quando si usa copy_model_version() e richiede la versione 3.4.0 MLflow o successiva. Quando questa variabile di ambiente è impostata su "true", non è necessaria una firma.

Le versioni del modello registrate senza firme presentano limitazioni. Vedere Aggiungere o aggiornare una firma per una versione del modello esistente.

import os

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

Per aggiungere una firma a una versione del modello esistente, vedere la documentazione di MLflow.

Script di esempio per eseguire la migrazione delle versioni del modello a un modello di catalogo Unity

Lo script seguente illustra come eseguire la migrazione di tutte le versioni del modello nell'area di lavoro registrata in un modello di catalogo Unity di destinazione. Questo script presuppone che la variabile di ambiente sia stata impostata su "true" come descritto in MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION modello.

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