使用 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 对象

试验状态

试验名称	状态	位置	height	width	损失	迭代	总时间 (秒)	iterations	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

试验进度

试验名称	date	done	episodes_total	experiment_id	experiment_tag	hostname	iterations	iterations_since_restore	mean_loss	neg_mean_loss	node_ip	进程 ID	time_since_restore	time_this_iter_s	time_total_s	时间戳	timesteps_since_restore	timesteps_total	training_iteration	trial_id	warmup_time
train_function_b275b_00000	2022-12-22_10-38-01	True		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	True		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	True		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	True		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	True		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 对象

试验状态

试验名称	状态	位置	height	width	损失	迭代	总时间 (秒)	iterations	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

试验进度

试验名称	date	done	episodes_total	experiment_id	experiment_tag	hostname	iterations	iterations_since_restore	mean_loss	neg_mean_loss	node_ip	进程 ID	time_since_restore	time_this_iter_s	time_total_s	时间戳	timesteps_since_restore	timesteps_total	training_iteration	trial_id	warmup_time
train_function_mlflow_b73bd_00000	2022-12-22_10-38-08	True		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	True		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	True		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	True		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	True		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, log_params_on_trial_end: bool = False)

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

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

请注意，回调将在驱动程序上打开 MLflow 会话，而不是在可训练项上。因此，无法在可训练项内部调用 MLflow 函数，例如 mlflow.log_figure()，因为在可训练项上没有 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 的全部内容作为工件保存到 MlFlow 中相应的运行。

• log_params_on_trial_end – 如果设置为 True，则在试验结束时将参数记录到 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,
        log_params_on_trial_end=True)])

MLflow 设置 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 会话。会话将在可训练项上创建。

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

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

在 Ray Train 的分布式训练中，只有零秩工作进程会初始化 mlflow。所有其他工作进程将返回一个 noop 客户端，以避免在分布式运行中重复记录。这可以通过传递 rank_zero_only=False 来禁用，这将为每个训练工作进程初始化 mlflow。注意：对于 Ray Tune，没有工作进程秩的概念，因此 rank_zero_only 会被忽略。

此函数将返回 mlflow 模块或非零秩工作进程的 noop 模块 if rank_zero_only=True。通过使用 mlflow = setup_mlflow(config)，您可以确保只有零秩工作进程调用 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 – 存储运行工件的位置。如果未提供，MLFlow 将选择适当的默认值。如果实验已存在，则忽略。

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

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

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

• rank_zero_only – 如果为 True，则仅为分布式训练中的秩 0 工作进程返回已初始化的会话。如果为 False，则为所有工作进程初始化会话。默认为 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 的返回值。这将确保它仅在分布式训练运行的第一个工作进程上调用。

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.cn/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 示例：使用 MLflow 和 Pytorch Lightning 与 Ray Tune 的示例。