Files
chongming-utils-task/examples/07_pipeline.py
T

113 lines
3.4 KiB
Python
Raw Normal View History

"""示例 7:DAG 管道 —— 一个源 → 两个分支 → 汇合。
管道把多个任务按依赖关系串起来,由一个调度器推进:节点先停在 ``blocked`` 等上游,
上游满足后才放行;注定跑不了的下游会被 ``skipped``。
运行::
uv run python examples/07_pipeline.py
"""
from __future__ import annotations
import time
from typing import Any
from chongming_task import Controller, Pipeline, render_pipeline
from chongming_task.policies import AggregationPolicy, JoinMode
def fetch(ctx: Any, day: str) -> str:
"""取数:模拟一次耗时操作。"""
for _ in range(10):
ctx.checkpoint()
time.sleep(0.03)
return f"{day} 的数据"
def transform(ctx: Any, day: str, shard: int) -> str:
"""分片转换。"""
for _ in range(8):
ctx.checkpoint()
time.sleep(0.02)
return f"{day}/shard{shard}"
def load(ctx: Any, day: str) -> str:
"""入库:依赖全部分片完成。"""
del ctx
return f"{day} 已入库"
def audit(ctx: Any, day: str) -> str:
"""审计:**不管上游成败都要跑**(ALL_SETTLED)。"""
del ctx
return f"{day} 审计完成"
def main() -> None:
controller = Controller(max_parallel=4)
# ---- 定义 ----
pipeline = Pipeline("nightly-etl", aggregation=AggregationPolicy.all_settled())
controller.register("fetch", fetch, tags={"io"})
controller.register("transform", transform)
controller.register("load", load, tags={"io"})
controller.register("audit", audit)
fetch_node = pipeline.node("fetch", "2026-10-02", node_id="fetch")
shards = [
pipeline.node(
"transform",
"2026-10-02",
index,
node_id=f"shard{index}",
depends_on=[fetch_node],
)
for index in range(3)
]
# 入库:等所有分片成功(默认 ALL_SUCCEEDED 屏障)
load_node = pipeline.node("load", "2026-10-02", node_id="load", depends_on=shards)
# 审计:等 load 落终态即可,不管成败
pipeline.node(
"audit",
"2026-10-02",
node_id="audit",
depends_on=[load_node],
join=JoinMode.ALL_SETTLED,
)
print(f"管道 {pipeline.name!r} 共 {len(pipeline)} 个节点,拓扑序:")
print(" " + " -> ".join(node.node_id for node in pipeline.topological_order()))
print(" 注意:位置参数是传给任务的,节点名由 node_id= 指定\n")
# ---- 跑起来 ----
run = controller.submit_pipeline(pipeline)
print(f"已提交,pipeline_run_id = {run.run_id}\n")
time.sleep(0.05)
print("--- 刚提交时:下游都停在 blocked ---")
print(render_pipeline(run))
print("\n--- 等待收场 ---")
outcome = run.wait(timeout=60)
print(render_pipeline(run))
print(f"\n管道结论:{outcome.status.value}")
for node_id, status in outcome.nodes.items():
print(f" {node_id:<8} {status.value}")
print(f"\n每个任务都带 pipeline_run_id:{_run_id_of(controller, run, 'load')}")
controller.close(timeout=10)
print("\n要点:依赖由状态机表达(blocked -> queued),调度器只负责放行。")
def _run_id_of(controller: Controller, run: Any, node_id: str) -> str:
snapshot = controller.get(run.task_id_of(node_id))
return "?" if snapshot is None else str(snapshot.pipeline_run_id)
if __name__ == "__main__":
main()