使用 MLflow Tracking Server 进行远程实验跟踪

在本教程中,您将学习如何使用 MLflow Tracking Server 为团队开发设置 MLflow Tracking 环境。

使用 MLflow Tracking Server 进行远程实验跟踪有许多好处:

  • 协作:多个用户可以将运行记录到同一个端点,并查询其他用户记录的运行和模型。

  • 分享结果: 跟踪服务器还提供了一个 跟踪用户界面 端点,团队成员可以轻松浏览彼此的结果。

  • 集中访问:跟踪服务器可以作为元数据和工件远程访问的代理运行,从而更容易保护和审计对数据的访问。

它是如何工作的?

下图描绘了使用带有 PostgreSQL 和 S3 的远程 MLflow Tracking Server 的架构

../../../_images/scenario_5.png

MLflow Tracking Server 与 PostgreSQL 和 S3 的架构图

备注

你可以在 artifact storesbackend stores 文档指南中找到支持的数据存储列表。

当你开始将运行记录到 MLflow 跟踪服务器时,会发生以下情况:

  • 第1a和1b部分

  • MLflow 客户端创建一个 RestStore 实例,并向日志 MLflow 实体发送 REST API 请求。

  • 跟踪服务器创建一个 SQLAlchemyStore 实例,并连接到远程主机以在数据库中插入跟踪信息(例如,指标、参数、标签等)。

  • 第一部分 1c 和 1d:

  • 客户端的检索请求从配置的 SQLAlchemyStore 表中返回信息。

  • 第二部分a和b:

  • 通过客户端使用 HttpArtifactRepository 将文件写入 MLflow Tracking Server 来记录工件事件。

  • 跟踪服务器随后将这些文件写入配置的对象存储位置,使用假定角色身份验证。

  • 第二部分c和d:

  • 从配置的后端存储中检索用户请求的工件是通过在服务器启动时配置的相同授权身份验证完成的。

  • 工件通过 HttpArtifactRepository 的接口通过跟踪服务器传递给最终用户。

入门指南

前言

在实际的生产部署环境中,您将拥有多个远程主机来运行跟踪服务器和数据库,如上图所示。然而,在本教程中,我们将仅使用一台机器,并在不同的端口上运行多个Docker容器,以模拟远程环境,从而实现更简单的评估教程设置。我们还将使用 MinIO ,一个兼容S3的对象存储,作为工件存储,这样您就不需要拥有AWS账户来运行本教程。

步骤 1 - 获取 MLflow 和额外依赖项

MLflow 可以在 PyPI 上获取。此外,访问 PostgreSQL 和 S3 需要 pyscopg2boto3。如果你还没有在系统上安装它们,你可以通过以下命令安装:

pip install mlflow psycopg2 boto3

步骤 2 - 设置远程数据存储

MLflow Tracking Server 可以与多种数据存储交互,以存储实验和运行数据以及工件。在本教程中,我们将使用 Docker Compose 启动两个容器,每个容器模拟实际环境中的远程服务器。

  1. PostgreSQL 数据库作为后端存储。

  2. MinIO 服务器作为制品存储。

安装 docker 和 docker-compose

备注

这些 Docker 步骤仅用于教程目的。MLflow 本身完全不依赖于 Docker。

按照官方说明安装 DockerDocker Compose。然后,运行 docker --versiondocker-compose --version 以确保它们已正确安装。

创建 compose.yaml

创建一个名为 compose.yaml 的文件,内容如下:

compose.yaml
version: '3.7'
services:
  # PostgreSQL database
  postgres:
    image: postgres:latest
    environment:
      POSTGRES_USER: user
      POSTGRES_PASSWORD: password
      POSTGRES_DB: mlflowdb
    ports:
      - 5432:5432
    volumes:
      - ./postgres-data:/var/lib/postgresql/data
  # MinIO server
  minio:
    image: minio/minio
    expose:
      - "9000"
    ports:
      - "9000:9000"
      # MinIO Console is available at http://localhost:9001
      - "9001:9001"
    environment:
      MINIO_ROOT_USER: "minio_user"
      MINIO_ROOT_PASSWORD: "minio_password"
    healthcheck:
      test: timeout 5s bash -c ':> /dev/tcp/127.0.0.1/9000' || exit 1
      interval: 1s
      timeout: 10s
      retries: 5
    command: server /data --console-address ":9001"
  # Create a bucket named "bucket" if it doesn't exist
  minio-create-bucket:
    image: minio/mc
    depends_on:
      minio:
        condition: service_healthy
    entrypoint: >
      bash -c "
      mc alias set minio http://minio:9000 minio_user minio_password &&
      if ! mc ls minio | grep --quiet bucket; then
        mc mb minio/bucket
      else
        echo 'bucket already exists'
      fi
      "

启动容器

compose.yaml 文件所在的同一目录运行以下命令以启动容器。这将启动后台的 PostgreSQL 和 Minio 服务器容器,并在 Minio 中创建一个名为 “bucket” 的新桶。

docker compose up -d

步骤 3 - 启动跟踪服务器

备注

在实际环境中,您将有一个远程主机来运行跟踪服务器,但在本教程中,我们将仅使用本地机器作为远程机器的模拟代理。

配置访问

为了使跟踪服务器能够访问远程存储,需要使用必要的凭证对其进行配置。

export MLFLOW_S3_ENDPOINT_URL=http://localhost:9000 # Replace this with remote storage endpoint e.g. s3://my-bucket in real use cases
export AWS_ACCESS_KEY_ID=minio_user
export AWS_SECRET_ACCESS_KEY=minio_password

你可以在 支持的存储 中找到如何为其他存储配置凭据的说明。

启动跟踪服务器

要指定后端存储和工件存储,可以分别使用 --backend-store-uri--artifacts-store-uri 选项。

mlflow server \
  --backend-store-uri postgresql://user:password@localhost:5432/mlflowdb \
  --artifacts-destination s3://bucket \
  --host 0.0.0.0 \
  --port 5000

在实际环境中,将 localhost 替换为数据库服务器的远程主机名或IP地址。

步骤 4:记录到跟踪服务器

一旦跟踪服务器运行,你可以通过将 MLflow 跟踪 URI 设置为跟踪服务器的 URI 来记录运行。或者,你可以使用 mlflow.set_tracking_uri() API 来设置跟踪 URI。

export MLFLOW_TRACKING_URI=http://127.0.0.1:5000  # Replace with remote host name or IP address in an actual environment

然后像往常一样使用 MLflow 跟踪 API 运行你的代码。以下代码在糖尿病数据集上运行 scikit-learn 随机森林模型的训练:

import mlflow

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_diabetes
from sklearn.ensemble import RandomForestRegressor

mlflow.autolog()

db = load_diabetes()
X_train, X_test, y_train, y_test = train_test_split(db.data, db.target)

# Create and train models.
rf = RandomForestRegressor(n_estimators=100, max_depth=6, max_features=3)
rf.fit(X_train, y_train)

# Use the model to make predictions on the test dataset.
predictions = rf.predict(X_test)

步骤 5:在跟踪界面中查看已记录的运行

我们的伪远程 MLflow 跟踪服务器也在同一端点上托管跟踪 UI。在实际部署环境中使用远程跟踪服务器时,情况也是如此。您可以通过在浏览器中导航到 `http://127.0.0.1:5000 <http://127.0.0.1:5000>`_(在实际环境中替换为远程主机名或 IP 地址)来访问 UI。

步骤 6:下载工件

MLflow 跟踪服务器还充当工件访问的代理主机。通过 runs:/mlflow-artifacts:/ 等代理 URI 启用工件访问,使用户无需管理直接访问的凭证或权限即可访问此位置。

import mlflow

run_id = "YOUR_RUN_ID"  # You can find run ID in the Tracking UI
artifact_path = "model"

# Download artifact via the tracking server
mlflow_artifact_uri = f"runs://{run_id}/{artifact_path}"
local_path = mlflow.artifacts.download_artifacts(mlflow_artifact_uri)

# Load the model
model = mlflow.sklearn.load_model(local_path)

下一步是什么?

现在你已经学会了如何为远程实验跟踪设置MLflow跟踪服务器!你还可以探索一些更高级的主题:

  • 跟踪服务器的其他配置:默认情况下,MLflow 跟踪服务器同时提供后端存储和工件存储。您还可以配置跟踪服务器仅提供后端存储或工件存储,以处理不同的使用场景,如高流量或安全问题。请参阅 其他使用场景 了解如何为这些使用场景自定义跟踪服务器。

  • 保护跟踪服务器--host 选项将服务暴露在所有接口上。如果在生产环境中运行服务器,我们建议不要广泛暴露内置服务器(因为它未经身份验证且未加密)。阅读 保护跟踪服务器 以了解在生产环境中保护跟踪服务器的最佳实践。

  • 新功能: MLflow 团队和众多社区贡献者不断开发新功能以支持更广泛的使用场景。查看 新功能 以了解最新功能!