如何在多个Python脚本中向同一MLflow Run记录实验数据？

核心问题原因 #

你无法自行生成Run ID并直接使用，因为MLflow要求Run ID必须对应已在追踪系统中创建的Run实例。正确的思路是先通过MLflow创建Run并获取其合法ID，再将该ID传递给所有需要共享的脚本。

具体实现步骤 #

1. 编写Run初始化脚本

创建init_mlflow_run.py，负责生成MLflow Run并输出其ID，供后续脚本调用：

import mlflow

# 若使用远程MLflow服务器，需设置追踪URI（本地存储可注释）
# mlflow.set_tracking_uri("http://your-mlflow-server:5000")

# 创建指定实验下的新Run，无实验名则使用默认实验
with mlflow.start_run(experiment_name="your_experiment_name") as run:
    # 将Run ID输出到标准输出，方便PowerShell捕获
    print(run.info.run_id)

2. 修改PowerShell触发脚本

在PowerShell中先执行初始化脚本获取Run ID，再将ID作为参数传递给预处理、训练脚本：

# 执行初始化脚本，捕获生成的Run ID
$RUN_ID = python init_mlflow_run.py

# 传递Run ID到预处理脚本
python preprocess.py --mlflow-run-id $RUN_ID

# 传递Run ID到训练脚本
python train.py --mlflow-run-id $RUN_ID

3. 改造业务脚本支持共享Run

以预处理脚本preprocess.py为例，添加参数接收Run ID，并附加到目标Run：

import mlflow
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--mlflow-run-id", required=True, help="目标MLflow Run的ID")
args = parser.parse_args()

# 保持与初始化脚本一致的追踪URI设置
# mlflow.set_tracking_uri("http://your-mlflow-server:5000")

# 附加到已存在的Run进行日志记录
with mlflow.start_run(run_id=args.mlflow_run_id) as run:
    mlflow.log_param("预处理批次大小", 64)
    mlflow.log_metric("数据清洗完成率", 1.0)
    # mlflow.log_artifact("预处理后数据路径")

训练脚本train.py同理：

import mlflow
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--mlflow-run-id", required=True, help="目标MLflow Run的ID")
args = parser.parse_args()

# 统一追踪URI设置
# mlflow.set_tracking_uri("http://your-mlflow-server:5000")

with mlflow.start_run(run_id=args.mlflow_run_id) as run:
    mlflow.log_param("学习率", 0.01)
    mlflow.log_metric("训练准确率", 0.98)
    # mlflow.sklearn.log_model(model, "训练模型")

4. 保留单独运行脚本的能力

若需要单独运行某脚本（不通过PowerShell触发），可将Run ID参数设为可选，无参数时自动创建新Run：

parser.add_argument("--mlflow-run-id", required=False, help="目标MLflow Run的ID（可选）")
args = parser.parse_args()

if args.mlflow_run_id:
    # 附加到已有Run
    with mlflow.start_run(run_id=args.mlflow_run_id) as run:
        # 业务日志逻辑
        pass
else:
    # 创建新Run
    with mlflow.start_run(experiment_name="your_experiment_name") as run:
        # 业务日志逻辑
        pass

关键注意事项 #

• 所有脚本必须使用相同的MLflow追踪URI，确保读写同一MLflow存储实例；
• 不要自行生成Run ID，必须通过MLflow创建Run后获取合法ID；
• 即使Run被标记为FINISHED，后续脚本仍可通过ID重新打开并追加日志，MLflow会自动将Run状态切换为RUNNING，完成后再改为FINISHED。

内容的提问来源于stack exchange，提问作者gevra