跨工作区共享模型

项目
04/01/2024

Azure Databricks 支持在多个工作区之间共享模型。例如，可以在开发工作区中开发和记录模型，然后访问该模型并将其与单独的生产工作区中的模型进行比较。多个团队共享对模型的访问权限或组织有多个工作区来处理不同的开发阶段时，这非常有用。对于跨工作区模型开发和部署，Databricks 建议采用部署代码方法，将模型训练代码部署到多个环境。

在多工作区情况下，可以使用远程模型注册表跨 Azure Databricks 工作区访问模型。例如，数据科学家可以使用只读访问权限访问生产模型注册表，将其开发中的模型与当前生产模型进行比较。多工作区设置示例如下所示。

多个工作区

对远程注册表的访问由令牌控制。每个需要访问权限的用户或脚本都会在远程注册表中创建个人访问令牌，并将该令牌复制到本地工作区的机密管理器中。发送到远程注册表工作区的每个 API 请求都必须包含访问令牌；MLflow 提供了一个简单的机制来指定要在执行模型注册表操作时使用的机密。

注意

作为安全最佳做法，在使用自动化工具、系统、脚本和应用进行身份验证时，Databricks 建议使用属于服务主体（而不是工作区用户）的个人访问令牌。若要为服务主体创建令牌，请参阅管理服务主体的令牌。

模型注册表的所有客户端和 Fluent API 方法都支持远程工作区。

要求

跨工作区使用模型注册表需要 MLflow Python 客户端，版本 1.11.0 或以上。

注意

此工作流是通过 MLflow 客户端中的逻辑实现的。确保运行客户端的环境有权对包含远程模型注册表的 Azure Databricks 工作区发出网络请求。注册表工作区的一个常见限制是 IP 允许列表，它可以禁止来自另一个工作区的群集中运行的 MLflow 客户端的连接。

设置远程注册表的 API 令牌

在模型注册表工作区中创建访问令牌。
在本地工作区中，创建机密来存储访问令牌和远程工作区信息：
1. 创建机密范围：databricks secrets create-scope <scope>。
2. 为目标工作区选取唯一名称，此处显示为 <prefix>。然后创建三个机密：
  - databricks secrets put-secret <scope> <prefix>-host：输入模型注册表工作区的主机名。例如，https://chinaeast2.databricks.azure.cn/ 或 https://adb-5555555555555555.19.databricks.azure.cn/。
  - databricks secrets put-secret <scope> <prefix>-token：输入模型注册表工作区的访问令牌。
  - databricks secrets put-secret <scope> <prefix>-workspace-id：输入模型注册表工作区的工作区 ID，其可以在任意页面的 URL 中查找。

注意

你可能想要与其他用户共享机密范围，因为每个工作区的机密范围数量有限。

指定远程注册表

根据为远程注册表工作区创建的机密范围和名称前缀，可以构造以下格式的注册表 URI：

registry_uri = f'databricks://<scope>:<prefix>'

你可以使用 URI 通过首次调用来指定 Fluent API 方法的远程注册表：

mlflow.set_registry_uri(registry_uri)

或者，你可以在实例化 MlflowClient 时显式指定它：

client = MlflowClient(registry_uri=registry_uri)

以下工作流演示了这两种方法的示例。

在远程注册表中注册模型

注册模型的一种方法是使用 mlflow.register_model API：

mlflow.set_registry_uri(registry_uri)
mlflow.register_model(model_uri=f'runs:/<run-id>/<artifact-path>', name=model_name)

其他模型注册方法的示例可以在本页末尾的笔记本中找到。

注意

在远程工作区中注册模型时，会在远程工作区中创建 DBFS 中的模型项目的临时副本。模型版本转为 READY 状态后，你可能需要删除此副本。可以在 /dbfs/databricks/mlflow/tmp-external-source/<run-id> 文件夹下找到临时文件。

你还可以指定 tracking_uri，以指向 registry_uri 的类似方式指向另一个工作区中的 MLflow 跟踪服务。这意味着可以在远程工作区上运行，并在当前或其他远程工作区中注册其模型。

使用远程注册表中的模型

可通过先设置注册表 URI，使用 mlflow.<flavor>.load_model 方法在远程注册表中加载和使用模型版本：

mlflow.set_registry_uri(registry_uri)
model = mlflow.pyfunc.load_model(f'models:/<model-name>/Staging')
model.predict(...)

或可以在 models:/ URI 中显式指定远程注册表：

model = mlflow.pyfunc.load_model(f'models://<scope>:<prefix>@databricks/<model-name>/Staging')
model.predict(...)

还支持用于访问模型文件的其他帮助器方法，例如：

client.get_latest_versions(model_name)
client.get_model_version_download_uri(model_name, version)

在远程注册表中管理模型

只要具有所需的权限，就可以对远程注册表中的模型执行任何操作。例如，如果对模型具有“可管理”权限，则可以使用 MlflowClient 方法转换模型版本阶段或删除模型：

client = MlflowClient(tracking_uri=None, registry_uri=registry_uri)
client.transition_model_version_stage(model_name, version, 'Archived')
client.delete_registered_model(model_name)

笔记本示例：远程模型注册表

以下笔记本适用于未启用 Unity Catalog 的工作区。它演示如何将模型从当前工作区记录到 MLflow 跟踪服务器，并将模型注册到其他工作区中的模型注册表中。

远程模型注册表示例笔记本

获取笔记本

要求