如何在Dagster中收集并行深度优先操作结果并执行后续Op？

感谢@zyd在Stack Overflow回答中提供的Dagster并行深度优先执行方案协助后，我希望实现一个能在所有graph运行完成后执行的@op，或至少等待所有graph执行完毕（它们本身无强依赖）。现有工作代码如下：

from typing import List
from dagster import (
    op, DynamicOut, DynamicOutput, In, Nothing, graph, job
)

@op
def get_csv_filenames(context) -> List[str]:
    # 获取CSV文件名的逻辑
    pass

@op(out=DynamicOut())
def generate_subtasks(context, csv_list:List[str]):
  for csv_filename in csv_list:
    yield DynamicOutput(csv_filename, mapping_key=csv_filename)

@op # 首个任务无依赖
def load_csv_into_duckdb(context, csv_filename):
    pass

@op(ins={"start":In(Nothing)})
def transform_dates(context, csv_filename):
    pass

@op(ins={"start":In(Nothing)})
def from_code_2_categories(context, csv_filename):
    pass

@op(ins={"start":In(Nothing)})
def export_2_parquet(context, csv_filename):
    pass

@op(ins={"start":In(Nothing)})
def profile_dataset(context, csv_filename):
    pass

@graph
def process(context, csv_filename:str):
  task1 = load_csv_into_duckdb(context=context, csv_filename=csv_filename)
  task2 = transform_dates(start=task1, context=context, csv_filename=csv_filename)
  task3 = from_code_2_categories(start=task2, context=context, csv_filename=csv_filename)
  task4 = export_2_parquet(start=task3, context=context, csv_filename=csv_filename)
  profile_dataset(start=task4, context=context, csv_filename=csv_filename)

@job
def pipeline():
  csv_filename_list = get_csv_filenames()
  generate_subtasks(csv_filename_list).map(process)

我尝试使用.map(process).collect()方法，但Dagster报错“Nonetype has no attribute collect”，尽管我看到过不少示例使用该方法。我还尝试让@graph返回单个任务的结果列表，但DagsterUI报错，要求graph装饰的函数返回带映射键的字典。我觉得应该从Dagster的执行上下文中获取相关信息，但不知如何在graph函数内访问。请问有解决思路吗？

解决思路 #

1. 修复.collect()调用的错误 #

报错的核心原因是你的process Graph没有返回任何输出，导致.map(process)生成的动态集合没有可收集的结果。需要修改Graph，让它返回明确的输出：

@graph
def process(csv_filename: str):  # 移除不必要的context参数，Op会自动获取上下文
    task1 = load_csv_into_duckdb(csv_filename=csv_filename)
    task2 = transform_dates(start=task1, csv_filename=csv_filename)
    task3 = from_code_2_categories(start=task2, csv_filename=csv_filename)
    task4 = export_2_parquet(start=task3, csv_filename=csv_filename)
    final_task = profile_dataset(start=task4, csv_filename=csv_filename)
    return final_task  # 返回Graph的最终输出

之后在Job中即可正常使用.collect()，并将收集到的结果传递给后续Op：

@op
def post_process_all(context, all_results: List):
    # 所有Graph执行完成后触发的逻辑
    context.log.info(f"Processed {len(all_results)} CSV files")

@job
def pipeline():
    csv_filename_list = get_csv_filenames()
    all_process_results = generate_subtasks(csv_filename_list).map(process).collect()
    post_process_all(all_process_results)

2. 处理Graph返回多结果的要求 #

如果需要从process Graph返回多个结果，不能直接返回列表，必须返回带映射键的字典，示例如下：

@graph
def process(csv_filename: str):
    task1 = load_csv_into_duckdb(csv_filename=csv_filename)
    task2 = transform_dates(start=task1, csv_filename=csv_filename)
    task3 = from_code_2_categories(start=task2, csv_filename=csv_filename)
    task4 = export_2_parquet(start=task3, csv_filename=csv_filename)
    profile_task = profile_dataset(start=task4, csv_filename=csv_filename)
    
    # 返回带映射键的字典
    return {
        "profile_output": profile_task,
        "parquet_path": task4
    }

此时.collect()会返回包含所有映射结果的列表，后续Op可以遍历这个列表获取每个Graph的输出。

3. 无需在Graph中传入执行上下文 #

Graph本身不需要接收context参数，每个Op在定义时声明context后，Dagster会自动将执行上下文传递给Op。你之前在Graph中手动传递context=context是多余的，直接移除即可。

4. 无结果收集的等待方案 #

如果不需要收集任何结果，只是想等待所有并行Graph执行完毕后触发操作，可以使用Nothing类型的输入来实现依赖：

@op(ins={"wait_for_all": In(Nothing)})
def final_cleanup(context):
    context.log.info("All CSV processing tasks completed!")

@job
def pipeline():
    csv_filename_list = get_csv_filenames()
    process_tasks = generate_subtasks(csv_filename_list).map(process)
    # 让final_cleanup等待所有process_tasks完成
    final_cleanup(wait_for_all=process_tasks)

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