如何用MLflow跟踪复杂数据准备管道？UI关联问题求解

核心思路：利用MLflow的Run依赖与Dataset血缘功能 #

MLflow Tracking原生支持通过Run层级依赖和Dataset血缘关系关联多步骤流程的输入输出，只需结合官方API的正确用法即可实现UI层面的关联展示。

1. 放弃自定义RunsDatasetSource，改用官方推荐的血缘记录方式 #

自定义RunsDatasetSource无法触发UI关联，因为MLflow的界面关联逻辑依赖内置的源解析器，而它明确排除了runs:/协议。正确操作如下：

• 第一步Run：处理完数据后，用mlflow.log_artifact()保存输出制品，同时用mlflow.data.from_filesystem()基于本地输出路径生成Dataset，再通过mlflow.log_input()记录该Dataset（作为第一步的输出数据集）。
• 第二步Run：先获取第一步的Run ID，通过mlflow.get_run(run_id).info.artifact_uri拼接出第一步制品的完整存储URI（比如本地路径./mlruns/1/xxx/artifacts/output_data或远程存储路径），而非直接使用runs:/<run_id>/<artifact_path>。
• 第二步中，将上述完整URI传入mlflow.data.from_filesystem()生成Dataset，再调用mlflow.log_input()记录为第二步的输入。

2. 建立Run间的父子依赖关系 #

启动第二步Run时，通过mlflow.start_run()的parent_run_id参数指定第一步的Run ID，让MLflow UI明确展示流程的层级关系：

# 第一步Run示例
with mlflow.start_run() as first_run:
    # 数据处理逻辑
    processed_step1 = process_step1(raw_data_path)
    # 记录输出制品
    mlflow.log_artifact(processed_step1, "step1_output")
    # 记录第一步的输出数据集
    dataset_step1 = mlflow.data.from_filesystem(processed_step1, name="step1_output_data")
    mlflow.log_input(dataset_step1, context="data_prep_step1")

# 第二步Run示例
with mlflow.start_run(parent_run_id=first_run.info.run_id) as second_run:
    # 下载第一步的制品
    artifact_local_path = mlflow.artifacts.download_artifacts(
        run_id=first_run.info.run_id,
        artifact_path="step1_output"
    )
    # 第二步数据处理
    processed_step2 = process_step2(artifact_local_path)
    # 记录第二步的输入数据集（使用第一步制品的完整URI）
    dataset_step2_input = mlflow.data.from_filesystem(
        f"{first_run.info.artifact_uri}/step1_output",
        name="step2_input_data"
    )
    mlflow.log_input(dataset_step2_input, context="data_prep_step2")
    # 记录第二步的输出制品
    mlflow.log_artifact(processed_step2, "step2_output")

3. 启用Dataset血缘可视化 #

确保使用MLflow 2.0及以上版本，完成上述操作后，在MLflow UI的「Datasets」页面中，可直接查看每个Dataset的上下游关联：包括它来自哪个Run的制品，以及被哪个Run作为输入使用。同时，「Parent/Child Runs」面板会展示多步骤流程的层级结构。

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