120 lines
4.1 KiB
Python
120 lines
4.1 KiB
Python
|
"""示例 8：``@task`` 装饰器与 ``TaskContext``。
|
|||
|
|
|
|||
|
|
演示四件事：
|
|||
|
|
|
|||
|
|
* ``@task`` 把元数据写在函数旁边，``register_module`` 一次登记整个模块；
|
|||
|
|
* ``ctx.progress(...)`` 上报进度——它会进快照、进存储，并出现在 ``ps`` 的 PROGRESS 列；
|
|||
|
|
* ``ctx.log(...)`` 走标准库 ``logging``，logger 名字是 ``chongming_task.task.<task_id>``；
|
|||
|
|
* ``ctx.input(节点id)`` 取上游节点的返回值——管道之间真正能传值。
|
|||
|
|
|
|||
|
|
运行::
|
|||
|
|
|
|||
|
|
uv run python examples/08_task_context.py
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
import sys
|
|||
|
|
import time
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
from chongming_task import Controller, Pipeline, RetryPolicy, render_table, task
|
|||
|
|
|
|||
|
|
|
|||
|
|
@task(tags={"io"}, retry=RetryPolicy(max_attempts=2))
|
|||
|
|
def extract(ctx: Any, day: str) -> list[str]:
|
|||
|
|
"""按天抽取分区列表。"""
|
|||
|
|
ctx.log(f"开始抽取 {day}")
|
|||
|
|
partitions = [f"{day}/part{i}" for i in range(6)]
|
|||
|
|
for index, part in enumerate(partitions, start=1):
|
|||
|
|
ctx.checkpoint()
|
|||
|
|
ctx.progress(index, total=len(partitions), unit="partitions", message=part)
|
|||
|
|
time.sleep(0.05)
|
|||
|
|
return partitions
|
|||
|
|
|
|||
|
|
|
|||
|
|
@task
|
|||
|
|
def transform(ctx: Any, day: str) -> int:
|
|||
|
|
"""把上游抽出的分区逐个转换，返回处理条数。"""
|
|||
|
|
partitions = ctx.input("extract", [])
|
|||
|
|
ctx.log(f"待转换分区：{len(partitions)} 个")
|
|||
|
|
rows = 0
|
|||
|
|
for index, part in enumerate(partitions, start=1):
|
|||
|
|
ctx.checkpoint()
|
|||
|
|
rows += 100
|
|||
|
|
ctx.progress(index, total=len(partitions), unit="partitions", message=part)
|
|||
|
|
time.sleep(0.05)
|
|||
|
|
return rows
|
|||
|
|
|
|||
|
|
|
|||
|
|
@task
|
|||
|
|
def load(ctx: Any, day: str) -> str:
|
|||
|
|
"""入库：条数来自上游的返回值，而不是自己再算一遍。"""
|
|||
|
|
rows = ctx.input("transform", 0)
|
|||
|
|
ctx.log(f"{day} 共 {rows} 行待入库")
|
|||
|
|
for index in range(1, 5):
|
|||
|
|
ctx.checkpoint()
|
|||
|
|
ctx.progress(index, total=4, unit="batch")
|
|||
|
|
time.sleep(0.04)
|
|||
|
|
return f"{day}: {rows} 行已入库"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main() -> None:
|
|||
|
|
logging.basicConfig(
|
|||
|
|
level=logging.INFO,
|
|||
|
|
format="%(levelname)-7s %(name)s %(message)s",
|
|||
|
|
stream=sys.stdout,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
controller = Controller(max_parallel=2)
|
|||
|
|
|
|||
|
|
# ---- 一次登记整个模块 ----
|
|||
|
|
registered = controller.register_module(sys.modules[__name__])
|
|||
|
|
print("从 @task 元数据登记的：")
|
|||
|
|
for item in registered:
|
|||
|
|
print(f" {item.name:<10} tags={sorted(item.tags) or '-'} {item.description}")
|
|||
|
|
|
|||
|
|
# ---- 串成管道：数据沿依赖流动 ----
|
|||
|
|
pipeline = Pipeline("daily-etl")
|
|||
|
|
extract_node = pipeline.node("extract", "2026-10-02", node_id="extract")
|
|||
|
|
transform_node = pipeline.node(
|
|||
|
|
"transform", "2026-10-02", node_id="transform", depends_on=[extract_node]
|
|||
|
|
)
|
|||
|
|
pipeline.node("load", "2026-10-02", node_id="load", depends_on=[transform_node])
|
|||
|
|
|
|||
|
|
run = controller.submit_pipeline(pipeline)
|
|||
|
|
print(f"\n已提交管道 {run.run_id}\n")
|
|||
|
|
|
|||
|
|
# ---- 跑的过程中看进度 ----
|
|||
|
|
deadline = time.monotonic() + 30
|
|||
|
|
while time.monotonic() < deadline and not run.settled:
|
|||
|
|
snapshots = controller.list_tasks(pipeline_run_id=run.run_id)
|
|||
|
|
cells: list[str] = []
|
|||
|
|
for snapshot in snapshots:
|
|||
|
|
node = snapshot.name.split("/", 1)[-1]
|
|||
|
|
progress = "-"
|
|||
|
|
if snapshot.progress is not None:
|
|||
|
|
progress = f"{snapshot.progress.done}/{snapshot.progress.total}"
|
|||
|
|
cells.append(f"{node:<10} {snapshot.status.value:<10} {progress:<8}")
|
|||
|
|
sys.stdout.write("\r" + " | ".join(cells).ljust(120))
|
|||
|
|
sys.stdout.flush()
|
|||
|
|
time.sleep(0.1)
|
|||
|
|
|
|||
|
|
outcome = run.wait(timeout=30)
|
|||
|
|
print("\n\n--- 结束 ---")
|
|||
|
|
print(render_table(controller.list_tasks(pipeline_run_id=run.run_id)))
|
|||
|
|
print(f"\n管道结论：{outcome.status.value}")
|
|||
|
|
|
|||
|
|
final = controller.get(run.task_id_of("load"))
|
|||
|
|
print(f"load 的返回值：{final.outcome if final else None}") # type: ignore[union-attr]
|
|||
|
|
|
|||
|
|
controller.close(timeout=10)
|
|||
|
|
print(
|
|||
|
|
"\n要点：ctx 是任务与运行时之间唯一的协作入口——停止、暂停、进度、日志、上游数据都在这里。"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|