将 MLflow 与 Tune 结合使用

MLflow 是一个用于管理机器学习生命周期的开源平台，包括实验、可复现性、部署和中心模型注册表。它目前提供了四个组件，其中包括用于记录和查询实验（包括代码、数据、配置和结果）的 MLflow Tracking。

Ray Tune 目前为 MLflow Tracking 提供了两个轻量级集成。一个是 MLflowLoggerCallback，它会自动将报告给 Tune 的指标记录到 MLflow Tracking API。

另一个是 setup_mlflow 函数，可以与函数 API 一起使用。它会自动使用 Tune 的训练信息初始化 MLflow API，并为每个 Tune 试验创建一个运行。然后，在您的训练函数中，您可以像平常一样使用 MLflow，例如使用 mlflow.log_metrics() 甚至 mlflow.autolog() 将日志记录到您的训练过程中。

运行 MLflow 示例

在下面的示例中，我们将使用上述两种方法，即 MLflowLoggerCallback 和 setup_mlflow 函数来记录指标。让我们从一些关键的导入开始

import os
import tempfile
import time

import mlflow

from ray import tune
from ray.air.integrations.mlflow import MLflowLoggerCallback, setup_mlflow

接下来，我们定义一个简单的训练函数（一个 Tune Trainable），它迭代计算步骤并评估中间得分，然后将得分报告给 Tune。

def evaluation_fn(step, width, height):
    return (0.1 + width * step / 100) ** (-1) + height * 0.1


def train_function(config):
    width, height = config["width"], config["height"]

    for step in range(config.get("steps", 100)):
        # Iterative training function - can be any arbitrary training procedure
        intermediate_score = evaluation_fn(step, width, height)
        # Feed the score back to Tune.
        tune.report({"iterations": step, "mean_loss": intermediate_score})
        time.sleep(0.1)

给定一个 MLFlow 跟踪 URI，您现在可以简单地将 MLflowLoggerCallback 用作 RunConfig() 的 callback 参数

def tune_with_callback(mlflow_tracking_uri, finish_fast=False):
    tuner = tune.Tuner(
        train_function,
        tune_config=tune.TuneConfig(num_samples=5),
        run_config=tune.RunConfig(
            name="mlflow",
            callbacks=[
                MLflowLoggerCallback(
                    tracking_uri=mlflow_tracking_uri,
                    experiment_name="mlflow_callback_example",
                    save_artifact=True,
                )
            ],
        ),
        param_space={
            "width": tune.randint(10, 100),
            "height": tune.randint(0, 100),
            "steps": 5 if finish_fast else 100,
        },
    )
    results = tuner.fit()

要使用 setup_mlflow 工具，只需在您的训练函数中调用此函数。请注意，我们还使用 mlflow.log_metrics(...) 将指标记录到 MLflow。除此之外，此版本的训练函数与其原始版本相同。

def train_function_mlflow(config):
    tracking_uri = config.pop("tracking_uri", None)
    setup_mlflow(
        config,
        experiment_name="setup_mlflow_example",
        tracking_uri=tracking_uri,
    )

    # Hyperparameters
    width, height = config["width"], config["height"]

    for step in range(config.get("steps", 100)):
        # Iterative training function - can be any arbitrary training procedure
        intermediate_score = evaluation_fn(step, width, height)
        # Log the metrics to mlflow
        mlflow.log_metrics(dict(mean_loss=intermediate_score), step=step)
        # Feed the score back to Tune.
        tune.report({"iterations": step, "mean_loss": intermediate_score})
        time.sleep(0.1)

有了这个新的目标函数，您现在可以按如下方式创建一个 Tune 运行

def tune_with_setup(mlflow_tracking_uri, finish_fast=False):
    # Set the experiment, or create a new one if does not exist yet.
    mlflow.set_tracking_uri(mlflow_tracking_uri)
    mlflow.set_experiment(experiment_name="setup_mlflow_example")

    tuner = tune.Tuner(
        train_function_mlflow,
        tune_config=tune.TuneConfig(num_samples=5),
        run_config=tune.RunConfig(
            name="mlflow",
        ),
        param_space={
            "width": tune.randint(10, 100),
            "height": tune.randint(0, 100),
            "steps": 5 if finish_fast else 100,
            "tracking_uri": mlflow.get_tracking_uri(),
        },
    )
    results = tuner.fit()

如果您恰好有一个 MLFlow 跟踪 URI，您可以在下面的 mlflow_tracking_uri 变量中设置它，并将 smoke_test=False。否则，您可以直接运行 tune_function 和 tune_decorated 函数的快速测试，而不使用 MLflow。

smoke_test = True

if smoke_test:
    mlflow_tracking_uri = os.path.join(tempfile.gettempdir(), "mlruns")
else:
    mlflow_tracking_uri = "<MLFLOW_TRACKING_URI>"

tune_with_callback(mlflow_tracking_uri, finish_fast=smoke_test)
if not smoke_test:
    df = mlflow.search_runs(
        [mlflow.get_experiment_by_name("mlflow_callback_example").experiment_id]
    )
    print(df)

tune_with_setup(mlflow_tracking_uri, finish_fast=smoke_test)
if not smoke_test:
    df = mlflow.search_runs(
        [mlflow.get_experiment_by_name("setup_mlflow_example").experiment_id]
    )
    print(df)

2022-12-22 10:37:53,580	INFO worker.py:1542 -- Started a local Ray instance. View the dashboard at http://127.0.0.1:8265 

Tune 状态

当前时间	2022-12-22 10:38:04
已运行时间	00:00:06.73
内存	10.4/16.0 GiB

系统信息

正在使用 FIFO 调度算法。
请求资源：0/16 CPU，0/0 GPU，0.0/4.03 GiB 堆内存，0.0/2.0 GiB 对象

试验状态

试验名称	状态	位置	高度	宽度	损失	迭代	总时间 (s)	迭代次数	neg_mean_loss
train_function_b275b_00000	已终止	127.0.0.1:801	66	36	7.24935	5	0.587302	4	-7.24935
train_function_b275b_00001	已终止	127.0.0.1:813	33	35	3.96667	5	0.507423	4	-3.96667
train_function_b275b_00002	已终止	127.0.0.1:814	75	29	8.29365	5	0.518995	4	-8.29365
train_function_b275b_00003	已终止	127.0.0.1:815	28	63	3.18168	5	0.567739	4	-3.18168
train_function_b275b_00004	已终止	127.0.0.1:816	20	18	3.21951	5	0.526536	4	-3.21951

试验进度

试验名称	日期	完成	总 Episode 数	实验 ID	实验标签	主机名	迭代次数	恢复后迭代次数	平均损失	neg_mean_loss	节点 IP	PID	恢复后时间	本次迭代时间 (s)	总时间 (s)	时间戳	恢复后时间步数	总时间步数	训练迭代	试验 ID	预热时间
train_function_b275b_00000	2022-12-22_10-38-01	真		28feaa4dd8ab4edab810e8109e77502e	0_height=66,width=36	kais-macbook-pro.anyscale.com.beta.tailscale.net	4	5	7.24935	-7.24935	127.0.0.1	801	0.587302	0.126818	0.587302	1671705481	0		5	b275b_00000	0.00293493
train_function_b275b_00001	2022-12-22_10-38-04	真		245010d0c3d0439ebfb664764ae9db3c	1_height=33,width=35	kais-macbook-pro.anyscale.com.beta.tailscale.net	4	5	3.96667	-3.96667	127.0.0.1	813	0.507423	0.122086	0.507423	1671705484	0		5	b275b_00001	0.00553799
train_function_b275b_00002	2022-12-22_10-38-04	真		898afbf9b906448c980f399c72a2324c	2_height=75,width=29	kais-macbook-pro.anyscale.com.beta.tailscale.net	4	5	8.29365	-8.29365	127.0.0.1	814	0.518995	0.123554	0.518995	1671705484	0		5	b275b_00002	0.0040431
train_function_b275b_00003	2022-12-22_10-38-04	真		03a4476f82734642b6ab0a5040ca58f8	3_height=28,width=63	kais-macbook-pro.anyscale.com.beta.tailscale.net	4	5	3.18168	-3.18168	127.0.0.1	815	0.567739	0.125471	0.567739	1671705484	0		5	b275b_00003	0.00406194
train_function_b275b_00004	2022-12-22_10-38-04	真		ff8c7c55ce6e404f9b0552c17f7a0c40	4_height=20,width=18	kais-macbook-pro.anyscale.com.beta.tailscale.net	4	5	3.21951	-3.21951	127.0.0.1	816	0.526536	0.123327	0.526536	1671705484	0		5	b275b_00004	0.00332022

2022-12-22 10:38:04,477	INFO tune.py:772 -- Total run time: 7.99 seconds (6.71 seconds for the tuning loop).

Tune 状态

当前时间	2022-12-22 10:38:11
已运行时间	00:00:07.00
内存	10.7/16.0 GiB

系统信息

正在使用 FIFO 调度算法。
请求资源：0/16 CPU，0/0 GPU，0.0/4.03 GiB 堆内存，0.0/2.0 GiB 对象

试验状态

试验名称	状态	位置	高度	宽度	损失	迭代	总时间 (s)	迭代次数	neg_mean_loss
train_function_mlflow_b73bd_00000	已终止	127.0.0.1:842	37	68	4.05461	5	0.750435	4	-4.05461
train_function_mlflow_b73bd_00001	已终止	127.0.0.1:853	50	20	6.11111	5	0.652748	4	-6.11111
train_function_mlflow_b73bd_00002	已终止	127.0.0.1:854	38	83	4.0924	5	0.6513	4	-4.0924
train_function_mlflow_b73bd_00003	已终止	127.0.0.1:855	15	93	1.76178	5	0.650586	4	-1.76178
train_function_mlflow_b73bd_00004	已终止	127.0.0.1:856	75	43	8.04945	5	0.656046	4	-8.04945

试验进度

试验名称	日期	完成	总 Episode 数	实验 ID	实验标签	主机名	迭代次数	恢复后迭代次数	平均损失	neg_mean_loss	节点 IP	PID	恢复后时间	本次迭代时间 (s)	总时间 (s)	时间戳	恢复后时间步数	总时间步数	训练迭代	试验 ID	预热时间
train_function_mlflow_b73bd_00000	2022-12-22_10-38-08	真		62703cfe82e54d74972377fbb525b000	0_height=37,width=68	kais-macbook-pro.anyscale.com.beta.tailscale.net	4	5	4.05461	-4.05461	127.0.0.1	842	0.750435	0.108625	0.750435	1671705488	0		5	b73bd_00000	0.0030272
train_function_mlflow_b73bd_00001	2022-12-22_10-38-11	真		03ea89852115465392ed318db8021614	1_height=50,width=20	kais-macbook-pro.anyscale.com.beta.tailscale.net	4	5	6.11111	-6.11111	127.0.0.1	853	0.652748	0.110796	0.652748	1671705491	0		5	b73bd_00001	0.00303078
train_function_mlflow_b73bd_00002	2022-12-22_10-38-11	真		3731fc2966f9453ba58c650d89035ab4	2_height=38,width=83	kais-macbook-pro.anyscale.com.beta.tailscale.net	4	5	4.0924	-4.0924	127.0.0.1	854	0.6513	0.108578	0.6513	1671705491	0		5	b73bd_00002	0.00310016
train_function_mlflow_b73bd_00003	2022-12-22_10-38-11	真		fb35841742b348b9912d10203c730f1e	3_height=15,width=93	kais-macbook-pro.anyscale.com.beta.tailscale.net	4	5	1.76178	-1.76178	127.0.0.1	855	0.650586	0.109097	0.650586	1671705491	0		5	b73bd_00003	0.0576491
train_function_mlflow_b73bd_00004	2022-12-22_10-38-11	真		6d3cbf9ecc3446369e607ff78c67bc29	4_height=75,width=43	kais-macbook-pro.anyscale.com.beta.tailscale.net	4	5	8.04945	-8.04945	127.0.0.1	856	0.656046	0.109869	0.656046	1671705491	0		5	b73bd_00004	0.00265694

2022-12-22 10:38:11,514	INFO tune.py:772 -- Total run time: 7.01 seconds (6.98 seconds for the tuning loop).

这结束了我们的 Tune 和 MLflow 演练。在以下部分中，您可以找到有关 Tune-MLflow 集成的 API 的更多详细信息。

MLflow 自动日志记录

您还可以查看此处，了解如何利用 MLflow 自动日志记录的示例，在本例中是与 Pytorch Lightning 结合使用

MLflow Logger API

class ray.air.integrations.mlflow.MLflowLoggerCallback(tracking_uri: str | None = None, *, registry_uri: str | None = None, experiment_name: str | None = None, tags: Dict | None = None, tracking_token: str | None = None, save_artifact: bool = False)

MLflow Logger 用于自动将 Tune 结果和配置记录到 MLflow。

MLflow (https://mlflow.org) Tracking 是一个用于记录和查询实验的开源库。此 Ray Tune LoggerCallback 将信息（配置参数、训练结果和指标以及 artifacts）发送到 MLflow 进行自动实验跟踪。

请记住，此回调将在驱动程序上打开 MLflow 会话，而不是在 trainable 上。因此，无法在 trainable 中调用 mlflow.log_figure() 等 MLflow 函数，因为 trainable 上没有 MLflow 会话。为了更精细的控制，请使用 ray.air.integrations.mlflow.setup_mlflow()。

参数:

• tracking_uri – 用于管理实验和运行的跟踪 URI。这可以是本地文件路径或远程服务器。此参数直接传递给 mlflow 初始化。在多节点设置中使用 Tune 时，请确保将其设置为远程服务器而不是本地文件路径。

• registry_uri – 直接传递给 mlflow 初始化的注册表 URI。

• experiment_name – 用于此 Tune 运行的实验名称。如果 MLflow 中已存在同名实验，则会重复使用。否则，将使用此名称创建一个新实验。

• tags – 可选的字符串键值字典，用作运行的标签

• tracking_token – 用于与 MLflow 进行身份验证的跟踪令牌。

• save_artifact – 如果设置为 True，则自动将 Tune 的 local_dir 的全部内容作为 artifact 保存到 MlFlow 中对应的运行。

示例

from ray.air.integrations.mlflow import MLflowLoggerCallback

tags = { "user_name" : "John",
         "git_commit_hash" : "abc123"}

tune.run(
    train_fn,
    config={
        # define search space here
        "parameter_1": tune.choice([1, 2, 3]),
        "parameter_2": tune.choice([4, 5, 6]),
    },
    callbacks=[MLflowLoggerCallback(
        experiment_name="experiment1",
        tags=tags,
        save_artifact=True)])

MLflow setup API

ray.air.integrations.mlflow.setup_mlflow(config: Dict | None = None, tracking_uri: str | None = None, registry_uri: str | None = None, experiment_id: str | None = None, experiment_name: str | None = None, tracking_token: str | None = None, artifact_location: str | None = None, run_name: str | None = None, create_experiment_if_not_exists: bool = False, tags: Dict | None = None, rank_zero_only: bool = True) → ModuleType | _NoopModule

设置 MLflow 会话。

此函数可用于在（分布式）训练或调优运行中初始化 MLflow 会话。会话将在 trainable 上创建。

默认情况下，MLflow 实验 ID 是 Ray 试验 ID，MLflow 实验名称是 Ray 试验名称。可以通过传递相应的关键字参数来覆盖这些设置。

的 config 字典会自动作为运行参数记录（不包括 mlflow 设置）。

在使用 Ray Train 进行分布式训练时，只有 rank 为零的 worker 会初始化 mlflow。所有其他 worker 将返回一个 noop 客户端，以便在分布式运行中不会重复记录日志。可以通过传递 rank_zero_only=False 来禁用此行为，此时每个训练 worker 都将初始化 mlflow。

如果 rank_zero_only=True，此函数将返回 mlflow 模块或非 rank 零 worker 的无操作模块。通过使用 mlflow = setup_mlflow(config)，可以确保只有 rank 零 worker 调用 mlflow API。

参数:

• config – 要作为参数记录到 mlflow 的配置字典。

• tracking_uri – MLflow 跟踪的跟踪 URI。如果在多节点设置中使用 Tune，请确保使用远程服务器进行跟踪。

• registry_uri – MLflow 模型注册表的注册表 URI。

• experiment_id – 已创建的 MLflow 实验的 ID。来自 tune.Tuner() 中所有试验的所有日志都将报告给此实验。如果未提供此参数或此 ID 的实验不存在，则必须提供 ``experiment_name``。此参数优先于 experiment_name。

• experiment_name – 已存在 MLflow 实验的名称。来自 tune.Tuner() 中所有试验的所有日志都将报告给此实验。如果未提供此参数，则必须提供有效的 experiment_id。

• tracking_token – 记录到远程跟踪服务器时用于 HTTP 身份验证的令牌。当您想记录到 Databricks 服务器时，此功能非常有用。此值将用于在所有远程训练进程上设置 MLFLOW_TRACKING_TOKEN 环境变量。

• artifact_location – 存储运行 artifacts 的位置。如果未提供，MLFlow 会选择一个合适的默认值。如果实验已存在，则忽略此参数。

• run_name – 将创建的新 MLflow 运行的名称。如果未设置，则默认为 experiment_name。

• create_experiment_if_not_exists – 如果提供的名称的实验不存在，是否创建它。默认为 False。

• tags – 为新运行设置的标签。

• rank_zero_only – 如果为 True，则仅为分布式训练中 rank 0 的 worker 返回已初始化的会话。如果为 False，则为所有 worker 初始化会话。默认为 True。

示例

默认情况下，您只需调用 setup_mlflow 并像平常一样继续使用 MLflow

from ray.air.integrations.mlflow import setup_mlflow

def training_loop(config):
    mlflow = setup_mlflow(config)
    # ...
    mlflow.log_metric(key="loss", val=0.123, step=0)

在分布式数据并行训练中，您可以利用 setup_mlflow 的返回值。这将确保它仅在分布式训练运行中的第一个 worker 上调用。

from ray.air.integrations.mlflow import setup_mlflow

def training_loop(config):
    mlflow = setup_mlflow(config)
    # ...
    mlflow.log_metric(key="loss", val=0.123, step=0)

如果您使用 Pytorch Lightning、XGBoost 等训练框架，也可以使用 MlFlow 的自动日志记录功能。更多信息可在此处找到 (https://mlflow.org/docs/latest/tracking.html#automatic-logging)。

from ray.air.integrations.mlflow import setup_mlflow

def train_fn(config):
    mlflow = setup_mlflow(config)
    mlflow.autolog()
    xgboost_results = xgb.train(config, ...)

PublicAPI (alpha): 此 API 处于 alpha 阶段，在稳定之前可能会发生变化。

更多 MLflow 示例

• MLflow PyTorch Lightning 示例：在 Ray Tune 中使用 MLflow 和 Pytorch Lightning 的示例。