Databricks 建議使用 Unity 目錄中的模型來改善治理、輕鬆地跨工作區和環境共用,以及實現更有彈性的 MLOps 工作流程。 若要將模型版本從工作區環境模型註冊表移轉至 Unity 目錄,Databricks 建議使用 copy_model_version() 搭配 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"
)
如果目的地模型不存在於 Unity 目錄中,則會由此 API 呼叫建立。
模型簽章
Unity 目錄中的模型需要簽章。 如果工作區模型版本沒有簽章,Databricks 建議您遵循 MLflow 檔中的指示來建立簽章。
若要簡化移轉,您可以使用環境變數 MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION。 只有在您使用 copy_model_version() 且需要 MLflow 版本 3.4.0 或更新版本時,才能使用此環境變數。 當此環境變數設為 "true"時,不需要簽章。
在沒有簽章的情況下註冊的模型版本有限制。 請參閱新增 或更新現有模型版本的簽章。
import os
os.environ["MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION"] = "true"
若要將簽章新增至現有的模型版本,請參閱 MLflow 檔。
將模型版本移轉至 Unity 目錄模型的範例腳本
下列指令碼示範如何將工作區註冊模型中的所有模型版本移轉至目的地 Unity 目錄模型。 此指令碼假設您已將環境變數 MLFLOW_SKIP_SIGNATURE_CHECK_FOR_UC_REGISTRY_MIGRATION 設定為 “true”,如 模型簽章中所述。
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")