亚马逊AWS官方博客

Amazon SageMaker 上完全托管式 MLflow 正式上线

今天，我们很高兴地宣布，Amazon SageMaker 上完全托管式 MLflow 功能正式上线。 MLflow 是一种广泛使用的开源工具，可有效帮助机器学习（ML）团队管理整个 ML 生命周期。借助该功能，客户现在只需几个步骤即可轻松设置和管理 MLflow 跟踪服务器，从而简化流程并提高生产力。

数据科学家和 ML 开发人员可以利用 MLflow 来跟踪训练模型的多次尝试，将其作为实验中的运行，以可视化的方式比较这些运行，评估模型，并将最佳模型注册到模型注册表。Amazon SageMaker 消除了设置和管理 MLflow 的繁重工作，为 ML 管理员提供了一种在 AWS 上建立安全且可扩展的 MLflow 环境的快速、有效的方法。

SageMaker 上托管式 MLflow 的核心组件

SageMaker 上完全托管式 MLflow 功能围绕三个核心组件构建：

  • MLflow 跟踪服务器：只需几个步骤，您就可以通过 SageMaker Studio UI 创建 MLflow 跟踪服务器。这个独立的 HTTP 服务器为跟踪运行和实验提供多个 REST API 端点，使您能够有效监控 ML 实验。为了实现更精细的安全自定义，您还可以使用 AWS 命令行界面（AWS CLI）。
  • MLflow 后端元数据存储：元数据存储是 MLflow 跟踪服务器的关键组成部分，与实验、运行和构件相关的所有元数据都保存在其中。这包括实验名称、运行 ID、参数值、指标、标签和构件位置，从而让您能够全面跟踪和管理 ML 实验。
  • MLflow 构件存储：此组件为 ML 实验期间生成的所有构件（例如训练模型、数据集、日志和图表）提供存储空间。该组件使用 Amazon Simple Storage Service（Amazon S3）存储桶，提供一个客户管理的 AWS 账户，实现构件的安全高效存储。

Amazon SageMaker with MLflow 的优势

使用 Amazon SageMaker with MLflow 可以简化并增强机器学习工作流程：

  • 全面的实验跟踪：跟踪在本地集成式开发环境（IDE）、SageMaker Studio 中的托管 IDE、SageMaker 训练作业、SageMaker 处理作业和 SageMaker Pipelines 中 MLflow 中的实验。
  • 完整的 MLflow 功能： 使用各项 MLflow 实验功能，包括 MLflow 跟踪、MLflow 评估和 MLflow 模型注册表，可轻松比较和评估训练迭代的结果。
  • 统一模型治理：在 MLflow 中注册的模型会自动出现在 SageMaker 模型注册表中，实现统一的模型治理体验，帮助您将 MLflow 模型部署到 SageMaker 推理，而无需构建自定义容器。
  • 高效的服务器管理：使用 SageMaker API 或 SageMaker Studio UI 根据需要预调配、删除和升级 MLflow 跟踪服务器。SageMaker 管理跟踪服务器的扩展、修补和持续维护，而无需客户管理底层基础设施。
  • 增强的安全性：使用 AWS Identity and Access Management（IAM）管理对 MLflow 跟踪服务器的安全访问。编写 IAM 策略来授予或拒绝对特定 MLflow API 的访问，确保 ML 环境安全。
  • 有效的监控和治理：使用 Amazon EventBridge 和 AWS CloudTrail 监控 MLflow 跟踪服务器上的活动，有效治理跟踪服务器。

MLflow 跟踪服务器先决条件（环境搭建）

  1. 创建 SageMaker Studio 域
    您可以使用全新 SageMaker Studio 体验创建 SageMaker Studio 域。
  2. 配置 IAM 执行角色
    MLflow 跟踪服务器需要 IAM 执行角色来读取和写入构件到 Amazon S3，并在 SageMaker 中注册模型。您可以使用 Studio 域执行角色作为跟踪服务器执行角色，也可以为跟踪服务器执行角色创建单独的角色。如果您选择创建新角色，请参阅《SageMaker 开发人员指南》，了解有关 IAM 角色的更多信息。如果您选择更新 Studio 域执行角色，请参阅《SageMaker 开发人员指南》，了解该角色所需的 IAM 策略的详细信息。

创建 MLflow 跟踪服务器
在此演练中，我使用默认设置创建 MLflow 跟踪服务器，这些设置包括跟踪服务器版本（2.13.2）、跟踪服务器大小（小）和跟踪服务器执行角色（Studio 域执行角色）。跟踪服务器的大小决定跟踪服务器可支持的使用量，我们建议，对于不超过 25 个用户的团队，应使用小型跟踪服务器。有关跟踪服务器配置的更多信息，请阅读《SageMaker 开发人员指南》。

首先，在上述环境搭建期间创建的 SageMaker Studio 域中，选择“应用程序”下的 MLflow，然后选择创建。

接下来，给跟踪服务器命名和构件存储位置（S3 URI）。

创建 MLflow 跟踪服务器最多可能需要 25 分钟。


跟踪和比较训练运行
要开始将指标、参数和构件记录到 MLflow，您需要一个 Jupyter Notebook 以及在创建步骤中分配到的跟踪服务器 ARN。您可以使用 MLflow 开发工具包来跟踪训练运行，并使用 MLflow UI 比较不同的运行。


要将 MLflow 模型注册表中的模型注册到 SageMaker 模型注册表，您需要使用 sagemaker-mlflow 插件，对 MLflow 开发工具包使用 AWS Signature V4 发出的所有 MLflow API 请求进行身份验证。

  1. 安装 MLflow 开发工具包和 sagemaker-mlflow 插件
    在您的 Notebook 中，首先安装 MLflow 开发工具包和 sagemaker-mlflow Python 插件。
    pip install mlflow==2.13.2 sagemaker-mlflow==0.1.0
  2. 跟踪实验中的运行
    要跟踪实验中的运行，请将以下代码复制到您的 Jupyter Notebook 中。

    import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

# 将其替换为你刚刚创建的跟踪服务器的 ARN
arn = 'YOUR-TRACKING-SERVER-ARN'

mlflow.set_tracking_uri(arn)

# 加载 Iris 数据集
iris = load_iris()
X, y = iris.data, iris.target

# 将数据拆分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 训练随机森林分类器
rf_model = RandomForestClassifier(n_estimators=100, random_state=42)
rf_model.fit(X_train, y_train)

# 使用测试集进行预测
y_pred = rf_model.predict(X_test)

# 计算评估指标
accuracy = accuracy_score(y_test, y_pred)
precision = precision_score(y_test, y_pred, average='weighted')
recall = recall_score(y_test, y_pred, average='weighted')
f1 = f1_score(y_test, y_pred, average='weighted')

# 开始 MLflow 运行
with mlflow.start_run():
# 记录模型
mlflow.sklearn.log_model(rf_model, "random_forest_model")

# 记录评估指标
mlflow.log_metric("accuracy", accuracy)
mlflow.log_metric("precision", precision)
mlflow.log_metric("recall", recall)
mlflow.log_metric("f1_score", f1)
  3. 在 MLflow UI 中查看运行
    运行步骤 2 中显示的 Notebook，您将在 MLflow UI 中看到新的运行。
  4. 比较运行
    您可以通过更改 random_state 多次运行此 Notebook，以便为每次训练运行生成不同的指标值。

注册候选模型
在根据步骤 4 中详述的方法比较了多次运行之后，您就可以在 MLflow 模型注册表中注册指标最符合您需求的模型。注册某个模型表明该模型可能适合用于生产部署，并且将进行进一步的测试以验证这种适用性。一旦模型在 MLflow 中注册，它就会自动出现在 SageMaker 模型注册表中，实现统一的模型治理体验，让您可以将 MLflow 模型部署到 SageMaker 推理中。这使得主要使用 MLflow 进行实验的数据科学家能够将他们的模型交给 ML 工程师，后者使用 SageMaker 模型注册表来管理模型的生产部署。

以下是在 MLflow 模型注册表中注册的模型。


以下是在 SageMaker 模型注册表中注册的模型。

清理
一经创建，MLflow 跟踪服务器便会开始产生费用，直至删除或停用。跟踪服务器的计费基于服务器运行的时间、大小以及记录到跟踪服务器的数据量。您可以在不使用时停用跟踪服务器，或者使用 API 或 SageMaker Studio UI 将其删除，以节省成本。有关定价的更多信息，请参阅 Amazon SageMaker 定价。

现已推出
SageMaker with MLflow 现已在所有提供 SageMaker Studio 的 AWS 区域正式推出，中国和美国 GovCloud 区域除外。我们邀请您探索这一新功能，让它帮助您提高机器学习项目的效率和控制力。要了解更多信息，请访问 SageMaker with MLflow 产品详细信息页面。

有关更多信息，请参阅《SageMaker 开发人员指南》，向 AWS re:Post for SageMaker 发送反馈，或通过您常用的 AWS Support 联系人发送反馈。

– Veliswa