113 lines
3.4 KiB
Python
113 lines
3.4 KiB
Python
|
"""示例 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()
|