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

在本教程中，您将学习如何使用 MLflow Tracking Server 设置 MLflow 跟踪环境以进行团队开发。

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

协作：多个用户可以向同一端点记录运行，并查询其他用户记录的运行和模型。
共享结果：跟踪服务器还提供一个 Tracking UI 端点，团队成员可以在其中轻松查看彼此的结果。
集中访问：跟踪服务器可以作为远程访问元数据和 artifact 的代理运行，从而更容易保护和审计数据访问。

工作原理？

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

使用带有 PostgreSQL 和 S3 的 MLflow Tracking Server 的架构图

注意

您可以在 artifact 存储和后端存储文档指南中找到支持的数据存储列表。

当您开始向 MLflow Tracking Server 记录运行时，会发生以下情况

第 1a 和 1b 部分:
- MLflow 客户端创建一个 RestStore 实例，并发送 REST API 请求来记录 MLflow 实体
- 跟踪服务器创建一个 SQLAlchemyStore 实例，并连接到远程主机，用于在数据库中插入跟踪信息（即指标、参数、标签等）
第 1c 和 1d 部分:
- 客户端的检索请求从配置的 SQLAlchemyStore 表中返回信息
第 2a 和 2b 部分:
- 客户端使用 HttpArtifactRepository 记录 artifact 事件，将文件写入 MLflow Tracking Server
- 然后，跟踪服务器使用假定角色认证将这些文件写入配置的对象存储位置
第 2c 和 2d 部分:
- 从配置的后端存储为用户请求检索 artifact 使用与服务器启动时配置的相同授权认证完成
- Artifact 通过 Tracking Server 通过 HttpArtifactRepository 的接口传递给最终用户

入门

前言

在实际生产部署环境中，您将有多个远程主机运行跟踪服务器和数据库，如上图所示。然而，出于本教程的目的，我们将仅使用一台机器，上面运行多个 Docker 容器，模拟远程环境，从而实现一个更容易评估的教程设置。我们还将使用 MinIO（一个兼容 S3 的对象存储）作为 artifact 存储，这样您就不需要拥有 AWS 账户即可运行本教程。

步骤 1 - 获取 MLflow 和附加依赖项

MLflow 可在 PyPI 上获取。此外，需要 pyscopg2 和 boto3 才能使用 Python 访问 PostgreSQL 和 S3。如果您的系统上尚未安装它们，您可以使用以下命令安装：

pip install mlflow psycopg2 boto3

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

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

PostgreSQL 数据库作为后端存储。
MinIO 服务器作为 artifact 存储。

安装 docker 和 docker-compose

注意

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

按照官方说明安装 Docker 和 Docker Compose。然后，运行 docker --version 和 docker-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/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

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

启动跟踪服务器

要指定后端存储和 artifact 存储，您可以使用 --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 Tracking 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 RandomForest 模型的训练：

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：在 Tracking UI 中查看记录的运行

我们的伪远程 MLflow Tracking Server 也在同一端点上托管 Tracking UI。在具有远程跟踪服务器的实际部署环境中，情况也是如此。您可以通过在浏览器中访问 http://127.0.0.1:5000（在实际环境中请替换为远程主机名或 IP 地址）来访问 UI。

步骤 6：下载 artifact

MLflow Tracking Server 也充当 artifact 访问的代理主机。artifact 访问通过 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 Tracking Server 进行远程实验跟踪！还有一些更高级的主题可以探索：

跟踪服务器的其他配置：默认情况下，MLflow Tracking Server 同时提供后端存储和 artifact 存储。您还可以配置跟踪服务器仅提供后端存储或 artifact 存储，以处理不同的用例，例如高流量或安全问题。有关如何针对这些用例自定义跟踪服务器的信息，请参阅其他用例。
保护跟踪服务器安全：--host 选项在所有接口上暴露服务。如果在生产环境中运行服务器，我们建议不要广泛暴露内置服务器（因为它未经认证且未加密）。请阅读保护跟踪服务器安全，了解在生产环境中保护跟踪服务器安全的最佳实践。
新功能：MLflow 团队和众多社区贡献者不断开发新功能，以支持更广泛的用例。请参阅新功能，了解最新特性！

工作原理？​

入门​

前言​

步骤 1 - 获取 MLflow 和附加依赖项​

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

安装 docker 和 docker-compose​

创建 compose.yaml​

启动容器​

步骤 3 - 启动跟踪服务器​

配置访问权限​

启动跟踪服务器​

步骤 4：记录到跟踪服务器​

步骤 5：在 Tracking UI 中查看记录的运行​

步骤 6：下载 artifact​

下一步是什么？​