跳到主要内容

使用 MLflow 追踪服务器进行远程实验追踪

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

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

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

它是如何工作的?

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

带PostgreSQL和S3的MLflow Tracking Server的架构图

注意

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

当您开始将运行记录到MLflow Tracking Server时,会发生以下情况:

  • 第一部分a和b:

    • MLflow客户端创建RestStore的实例,并发送REST API请求来记录MLflow实体。
    • Tracking Server创建SQLAlchemyStore的实例,并连接到远程主机,将跟踪信息(例如,度量、参数、标签等)插入数据库。
  • 第一部分c和d:

    • 客户端的检索请求从配置的SQLAlchemyStore表中返回信息。
  • 第二部分a和b:

    • 客户端通过HttpArtifactRepository进行工件的记录事件,将文件写入MLflow Tracking Server。
    • 然后,Tracking Server使用假定的角色身份验证将这些文件写入配置的对象存储位置。
  • 第二部分c和d:

    • 从配置的后端存储中检索工件以响应用户请求,是通过在服务器启动时配置的相同授权身份验证完成的。
    • 通过Tracking Server,通过HttpArtifactRepository接口将工件传递给最终用户。

开始使用

前言

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

步骤1 - 获取MLflow和其他依赖项

MLflow可在PyPI上获得。此外,要使用Python访问PostgreSQL和S3,还需要pyscopg2boto3。如果您尚未在系统上安装它们,可以使用以下命令进行安装:

bash
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
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

# MinIO server
minio:
image: minio/minio
expose:
- "9000"
ports:
- "9000:9000"
# MinIO Console is available at https://: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”的新存储桶。

bash
docker compose up -d

步骤3 - 启动Tracking Server

注意

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

配置访问

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

bash
export MLFLOW_S3_ENDPOINT_URL=https://: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

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

启动跟踪服务器

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

bash
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:将运行记录到Tracking Server

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

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

然后像往常一样使用MLflow跟踪API运行您的代码。以下代码在diabetes数据集上运行scikit-learn RandomForest模型的训练。

python
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:下载工件

MLflow Tracking Server还充当工件访问的代理主机。工件访问通过代理URI(如models:/mlflow-artifacts:/)启用,使开发人员无需管理直接访问的凭据或权限即可访问此位置。

python
import mlflow

model_id = "YOUR_MODEL_ID" # You can find model ID in the Tracking UI

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

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

下一步是什么?

现在您已经学会了如何设置MLflow Tracking Server进行远程实验跟踪!还有一些更高级的主题可以探索:

  • Tracking Server的其他配置:默认情况下,MLflow Tracking Server同时提供后端存储和工件存储。您还可以将Tracking Server配置为仅提供后端存储或工件存储,以处理不同的用例,例如大流量或安全问题。请参阅other use cases了解如何为这些用例定制Tracking Server。
  • 保护Tracking Server--host选项会在所有接口上公开服务。如果在生产环境中运行服务器,我们建议不要广泛公开内置服务器(因为它没有身份验证且未加密)。请阅读Secure Tracking Server以了解在生产环境中保护Tracking Server的最佳实践。