591 lines
22 KiB
Python
591 lines
22 KiB
Python
|
"""DAG 管道测试。
|
||
|
|
|
||
|
|
分三层：
|
||
|
|
|
||
|
|
1. **定义期校验**——环、悬空依赖、未注册任务、quorum 都要在创建任何任务之前拒绝，
|
||
|
|
否则存储里会留下一堆看起来正常的半成品任务；
|
||
|
|
2. **推进**——屏障、扇出、扇入、三种汇合语义；
|
||
|
|
3. **收场**——fail_fast 连坐、下游跳过、整体状态推导。
|
||
|
|
|
||
|
|
约定：节点 id 一律用 ``node_id=`` 显式给出。第一个位置参数是**已注册的任务名**，
|
||
|
|
其后才是传给任务的位置参数——这两者很容易看混，显式写出来既清楚又自成文档。
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from chongming_task.controller import Controller
|
||
|
|
from chongming_task.errors import ConfigurationError
|
||
|
|
from chongming_task.pipeline import (
|
||
|
|
Pipeline,
|
||
|
|
PipelineStatus,
|
||
|
|
render_pipeline,
|
||
|
|
)
|
||
|
|
from chongming_task.policies import AggregationPolicy, JoinMode
|
||
|
|
from chongming_task.states import TaskStatus
|
||
|
|
from chongming_task.stores.memory import InMemoryTaskStore
|
||
|
|
|
||
|
|
_WINDOW = 15.0
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# 1. 定义期校验
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
|
||
|
|
def test_empty_pipeline_is_rejected() -> None:
|
||
|
|
with pytest.raises(ConfigurationError, match="没有任何节点"):
|
||
|
|
Pipeline("empty").validate()
|
||
|
|
|
||
|
|
|
||
|
|
def test_dangling_dependency_is_rejected() -> None:
|
||
|
|
pipeline = Pipeline("dangling")
|
||
|
|
pipeline.node("add", node_id="a", depends_on=["no-such-node"])
|
||
|
|
with pytest.raises(ConfigurationError, match="不存在的节点"):
|
||
|
|
pipeline.validate()
|
||
|
|
|
||
|
|
|
||
|
|
def test_self_dependency_is_rejected() -> None:
|
||
|
|
pipeline = Pipeline("self")
|
||
|
|
pipeline.node("add", node_id="a", depends_on=["a"])
|
||
|
|
with pytest.raises(ConfigurationError, match="依赖了自己"):
|
||
|
|
pipeline.validate()
|
||
|
|
|
||
|
|
|
||
|
|
def test_cycle_is_rejected() -> None:
|
||
|
|
"""有环时必须显式报错。
|
||
|
|
|
||
|
|
不检测的话调度器只会安静地什么都不做（所有节点都在等对方），表现为"管道卡住了"
|
||
|
|
——那种故障最难查。
|
||
|
|
"""
|
||
|
|
pipeline = Pipeline("cyclic")
|
||
|
|
first = pipeline.node("add", node_id="a")
|
||
|
|
second = pipeline.node("add", node_id="b", depends_on=[first])
|
||
|
|
# 人为制造环：a 反过来依赖 b。
|
||
|
|
object.__setattr__(pipeline.node_by_id("a"), "depends_on", ("b",))
|
||
|
|
assert second.depends_on == ("a",)
|
||
|
|
|
||
|
|
with pytest.raises(ConfigurationError, match="环形依赖"):
|
||
|
|
pipeline.validate()
|
||
|
|
|
||
|
|
|
||
|
|
def test_unregistered_task_is_rejected_before_creating_anything() -> None:
|
||
|
|
"""控制面只接受任务名。写错的 DAG 不能留下半成品任务。"""
|
||
|
|
store = InMemoryTaskStore()
|
||
|
|
controller = Controller(store=store)
|
||
|
|
controller.register("add", _add)
|
||
|
|
|
||
|
|
pipeline = Pipeline("bad")
|
||
|
|
pipeline.node("add", node_id="ok")
|
||
|
|
pipeline.node("nope", node_id="typo")
|
||
|
|
|
||
|
|
with pytest.raises(ConfigurationError):
|
||
|
|
controller.submit_pipeline(pipeline)
|
||
|
|
assert store.query() == [], "校验失败前不该创建任何任务"
|
||
|
|
|
||
|
|
|
||
|
|
def test_quorum_settlement_is_rejected_with_an_explanation() -> None:
|
||
|
|
"""诚实拒绝，而不是把 quorum 硬套到 DAG 上发明一个没人能预测的语义。"""
|
||
|
|
pipeline = Pipeline("q", aggregation=AggregationPolicy.of_quorum(2))
|
||
|
|
pipeline.node("add")
|
||
|
|
with pytest.raises(ConfigurationError, match="不支持 quorum"):
|
||
|
|
pipeline.validate()
|
||
|
|
|
||
|
|
|
||
|
|
def test_duplicate_node_id_is_rejected() -> None:
|
||
|
|
pipeline = Pipeline("dup")
|
||
|
|
pipeline.node("add", node_id="same")
|
||
|
|
with pytest.raises(ConfigurationError, match="节点 id 重复"):
|
||
|
|
pipeline.node("add", node_id="same")
|
||
|
|
|
||
|
|
|
||
|
|
def test_same_task_can_appear_many_times_with_auto_ids() -> None:
|
||
|
|
"""扇出靠"同名多次"表达，不该强迫调用方给每个节点起名。"""
|
||
|
|
pipeline = Pipeline("fanout")
|
||
|
|
pipeline.node("add")
|
||
|
|
pipeline.node("add")
|
||
|
|
pipeline.node("add")
|
||
|
|
assert [node.node_id for node in pipeline.nodes] == ["add", "add#2", "add#3"]
|
||
|
|
|
||
|
|
|
||
|
|
def test_positional_args_belong_to_the_task_not_the_node_id() -> None:
|
||
|
|
"""位置参数是传给任务的，不是节点名——这正是把参数叫 ``node_id`` 的原因。"""
|
||
|
|
pipeline = Pipeline("args")
|
||
|
|
node = pipeline.node("add", 2, 3, node_id="sum")
|
||
|
|
assert node.node_id == "sum"
|
||
|
|
assert node.args == (2, 3)
|
||
|
|
|
||
|
|
|
||
|
|
def test_topological_order_puts_upstream_first() -> None:
|
||
|
|
pipeline = Pipeline("order")
|
||
|
|
last = pipeline.node("add", node_id="last")
|
||
|
|
first = pipeline.node("add", node_id="first")
|
||
|
|
pipeline.node("add", node_id="middle", depends_on=[first, last])
|
||
|
|
|
||
|
|
order = [node.node_id for node in pipeline.topological_order()]
|
||
|
|
assert order.index("first") < order.index("middle")
|
||
|
|
assert order.index("last") < order.index("middle")
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# 2. 推进
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
|
||
|
|
def test_a_linear_pipeline_runs_in_dependency_order() -> None:
|
||
|
|
controller = Controller(max_parallel=4)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
|
||
|
|
pipeline = Pipeline("linear")
|
||
|
|
extract = pipeline.node("mark", node_id="extract")
|
||
|
|
transform = pipeline.node("mark", node_id="transform", depends_on=[extract])
|
||
|
|
pipeline.node("mark", node_id="load", depends_on=[transform])
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p1")
|
||
|
|
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
|
||
|
|
assert _finished_order(controller, "p1") == ["extract", "transform", "load"]
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_a_downstream_node_waits_in_blocked_until_its_upstream_finishes() -> None:
|
||
|
|
gate = threading.Event()
|
||
|
|
controller = Controller(max_parallel=4)
|
||
|
|
controller.register("wait_gate", _wait_gate)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
|
||
|
|
pipeline = Pipeline("blocking")
|
||
|
|
upstream = pipeline.node("wait_gate", gate, node_id="upstream")
|
||
|
|
pipeline.node("mark", node_id="downstream", depends_on=[upstream])
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p2")
|
||
|
|
try:
|
||
|
|
assert _wait_until(lambda: run.node_status("upstream") is TaskStatus.RUNNING)
|
||
|
|
time.sleep(0.1)
|
||
|
|
|
||
|
|
assert run.node_status("downstream") is TaskStatus.BLOCKED, "上游没完成前必须停在 blocked"
|
||
|
|
assert run.status is PipelineStatus.RUNNING
|
||
|
|
|
||
|
|
gate.set()
|
||
|
|
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
|
||
|
|
finally:
|
||
|
|
gate.set()
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_fan_out_runs_siblings() -> None:
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
|
||
|
|
pipeline = Pipeline("fanout")
|
||
|
|
root = pipeline.node("mark", node_id="root")
|
||
|
|
for index in range(4):
|
||
|
|
pipeline.node("mark", node_id=f"shard{index}", depends_on=[root])
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p3")
|
||
|
|
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
|
||
|
|
assert set(run.task_ids) == {"root", "shard0", "shard1", "shard2", "shard3"}
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_fan_in_with_all_succeeded_waits_for_every_upstream() -> None:
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
|
||
|
|
pipeline = Pipeline("fanin")
|
||
|
|
root = pipeline.node("mark", node_id="root")
|
||
|
|
left = pipeline.node("mark", node_id="left", depends_on=[root])
|
||
|
|
right = pipeline.node("mark", node_id="right", depends_on=[root])
|
||
|
|
pipeline.node("mark", node_id="join", depends_on=[left, right])
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p4")
|
||
|
|
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
|
||
|
|
|
||
|
|
order = _finished_order(controller, "p4")
|
||
|
|
assert order.index("join") > order.index("left")
|
||
|
|
assert order.index("join") > order.index("right")
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_any_succeeded_join_runs_after_the_first_success() -> None:
|
||
|
|
"""先到先得：任一上游成功即可放行，不必等另一个失败的上游。"""
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
controller.register("boom", _boom)
|
||
|
|
|
||
|
|
pipeline = Pipeline("any")
|
||
|
|
good = pipeline.node("mark", node_id="good")
|
||
|
|
bad = pipeline.node("boom", node_id="bad")
|
||
|
|
winner = pipeline.node(
|
||
|
|
"mark", node_id="winner", depends_on=[good, bad], join=JoinMode.ANY_SUCCEEDED
|
||
|
|
)
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p5")
|
||
|
|
outcome = run.wait(timeout=_WINDOW)
|
||
|
|
|
||
|
|
assert winner.join is JoinMode.ANY_SUCCEEDED
|
||
|
|
assert outcome.nodes["winner"] is TaskStatus.SUCCEEDED, "有一个上游成功就该放行"
|
||
|
|
assert outcome.status is PipelineStatus.FAILED, "另一个上游失败，整条管道仍失败"
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_any_succeeded_join_skips_when_nothing_succeeds() -> None:
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("boom", _boom)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
|
||
|
|
pipeline = Pipeline("any-none")
|
||
|
|
bad = pipeline.node("boom", node_id="bad")
|
||
|
|
pipeline.node("mark", node_id="winner", depends_on=[bad], join=JoinMode.ANY_SUCCEEDED)
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p6")
|
||
|
|
outcome = run.wait(timeout=_WINDOW)
|
||
|
|
assert outcome.nodes["winner"] is TaskStatus.SKIPPED
|
||
|
|
assert outcome.status is PipelineStatus.FAILED
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_all_settled_join_runs_even_after_an_upstream_failure() -> None:
|
||
|
|
"""清理/通知型节点：不管上游成败都要跑。"""
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("boom", _boom)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
|
||
|
|
pipeline = Pipeline("cleanup", aggregation=AggregationPolicy.all_settled())
|
||
|
|
failed = pipeline.node("boom", node_id="work")
|
||
|
|
cleanup = pipeline.node(
|
||
|
|
"mark", node_id="cleanup", depends_on=[failed], join=JoinMode.ALL_SETTLED
|
||
|
|
)
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p7")
|
||
|
|
outcome = run.wait(timeout=_WINDOW)
|
||
|
|
assert cleanup.join is JoinMode.ALL_SETTLED
|
||
|
|
assert outcome.nodes["cleanup"] is TaskStatus.SUCCEEDED, "ALL_SETTLED 下游必须照跑"
|
||
|
|
assert outcome.status is PipelineStatus.FAILED, "上游失败仍让整条管道失败"
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# 3. 收场
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
|
||
|
|
def test_failed_upstream_skips_downstream_with_a_reason() -> None:
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("boom", _boom)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
|
||
|
|
pipeline = Pipeline("skip")
|
||
|
|
bad = pipeline.node("boom", node_id="bad")
|
||
|
|
pipeline.node("mark", node_id="after", depends_on=[bad])
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p8")
|
||
|
|
outcome = run.wait(timeout=_WINDOW)
|
||
|
|
|
||
|
|
assert outcome.nodes["after"] is TaskStatus.SKIPPED
|
||
|
|
snapshot = controller.get("p8:after")
|
||
|
|
assert snapshot is not None
|
||
|
|
assert snapshot.status is TaskStatus.SKIPPED
|
||
|
|
assert snapshot.failure is None, "跳过不是失败，不该带 failure"
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_fail_fast_cancels_independent_running_nodes() -> None:
|
||
|
|
"""真正需要连坐的场景：另一个分支已经在跑，必须被取消。
|
||
|
|
|
||
|
|
被连坐的必须落 ``CANCELLED`` 而不是 ``STOPPED``——"上游挂了"与"值班同学手动停的"
|
||
|
|
是完全不同的事，混成一个状态就再也分不出来了。
|
||
|
|
"""
|
||
|
|
gate = threading.Event()
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("boom", _boom)
|
||
|
|
controller.register("wait_gate", _wait_gate)
|
||
|
|
|
||
|
|
pipeline = Pipeline("cascade", aggregation=AggregationPolicy.fail_fast())
|
||
|
|
pipeline.node("wait_gate", gate, node_id="independent")
|
||
|
|
pipeline.node("boom", node_id="bad")
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p9")
|
||
|
|
try:
|
||
|
|
assert _wait_until(lambda: run.node_status("independent") is TaskStatus.RUNNING)
|
||
|
|
outcome = run.wait(timeout=_WINDOW)
|
||
|
|
|
||
|
|
assert outcome.status is PipelineStatus.FAILED
|
||
|
|
assert outcome.nodes["independent"] is TaskStatus.CANCELLED, (
|
||
|
|
"被连坐的应当是 cancelled，而不是 stopped"
|
||
|
|
)
|
||
|
|
snapshot = controller.get("p9:independent")
|
||
|
|
assert snapshot is not None
|
||
|
|
assert snapshot.failure is not None
|
||
|
|
assert snapshot.failure.kind.value == "cancelled"
|
||
|
|
finally:
|
||
|
|
gate.set()
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_fail_fast_never_settles_while_a_cancelled_node_is_still_running() -> None:
|
||
|
|
"""``status`` 必须先看终态、再看失败标志。
|
||
|
|
|
||
|
|
顺序反了，等待者会被提前唤醒，看到一个"已经失败、但还有任务在跑"的矛盾状态。
|
||
|
|
"""
|
||
|
|
gate = threading.Event()
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("boom", _boom)
|
||
|
|
controller.register("wait_gate", _wait_gate)
|
||
|
|
|
||
|
|
pipeline = Pipeline("ordering", aggregation=AggregationPolicy.fail_fast())
|
||
|
|
pipeline.node("wait_gate", gate, node_id="independent")
|
||
|
|
pipeline.node("boom", node_id="bad")
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p10")
|
||
|
|
try:
|
||
|
|
assert _wait_until(lambda: run.node_status("independent") is TaskStatus.RUNNING)
|
||
|
|
outcome = run.wait(timeout=_WINDOW)
|
||
|
|
for node_id in outcome.nodes:
|
||
|
|
assert outcome.nodes[node_id] in {
|
||
|
|
TaskStatus.SUCCEEDED,
|
||
|
|
TaskStatus.FAILED,
|
||
|
|
TaskStatus.SKIPPED,
|
||
|
|
TaskStatus.CANCELLED,
|
||
|
|
TaskStatus.STOPPED,
|
||
|
|
TaskStatus.LOST,
|
||
|
|
}, f"{node_id} 在管道给出结论时还没落终态"
|
||
|
|
finally:
|
||
|
|
gate.set()
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_all_settled_does_not_cancel_siblings() -> None:
|
||
|
|
gate = threading.Event()
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("boom", _boom)
|
||
|
|
controller.register("wait_gate", _wait_gate)
|
||
|
|
|
||
|
|
pipeline = Pipeline("settled", aggregation=AggregationPolicy.all_settled())
|
||
|
|
pipeline.node("wait_gate", gate, node_id="independent")
|
||
|
|
pipeline.node("boom", node_id="bad")
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p11")
|
||
|
|
try:
|
||
|
|
assert _wait_until(lambda: run.node_status("independent") is TaskStatus.RUNNING)
|
||
|
|
time.sleep(0.2)
|
||
|
|
assert run.node_status("independent") is TaskStatus.RUNNING, "all_settled 不该连坐"
|
||
|
|
|
||
|
|
gate.set()
|
||
|
|
outcome = run.wait(timeout=_WINDOW)
|
||
|
|
assert outcome.nodes["independent"] is TaskStatus.SUCCEEDED
|
||
|
|
assert outcome.status is PipelineStatus.FAILED, "有人失败过，整体就是失败"
|
||
|
|
finally:
|
||
|
|
gate.set()
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_pipeline_tasks_carry_the_run_id() -> None:
|
||
|
|
controller = Controller(max_parallel=4)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
pipeline = Pipeline("tagged")
|
||
|
|
pipeline.node("mark", node_id="one")
|
||
|
|
pipeline.node("mark", node_id="two")
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="tagged-1")
|
||
|
|
run.wait(timeout=_WINDOW)
|
||
|
|
|
||
|
|
for node_id, task_id in run.task_ids.items():
|
||
|
|
snapshot = controller.get(task_id)
|
||
|
|
assert snapshot is not None
|
||
|
|
assert snapshot.pipeline_run_id == "tagged-1"
|
||
|
|
assert task_id == f"tagged-1:{node_id}"
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_cancelling_the_whole_pipeline() -> None:
|
||
|
|
gate = threading.Event()
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("wait_gate", _wait_gate)
|
||
|
|
|
||
|
|
pipeline = Pipeline("cancelme")
|
||
|
|
root = pipeline.node("wait_gate", gate, node_id="root")
|
||
|
|
pipeline.node("wait_gate", gate, node_id="child", depends_on=[root])
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p12")
|
||
|
|
try:
|
||
|
|
assert _wait_until(lambda: run.node_status("root") is TaskStatus.RUNNING)
|
||
|
|
assert run.cancel(reason="运维取消") >= 1
|
||
|
|
|
||
|
|
outcome = run.wait(timeout=_WINDOW)
|
||
|
|
assert outcome.status is PipelineStatus.CANCELLED
|
||
|
|
assert outcome.nodes["root"] is TaskStatus.CANCELLED
|
||
|
|
assert outcome.nodes["child"] is TaskStatus.CANCELLED
|
||
|
|
finally:
|
||
|
|
gate.set()
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_wait_timeout_does_not_change_anything() -> None:
|
||
|
|
gate = threading.Event()
|
||
|
|
controller = Controller(max_parallel=4)
|
||
|
|
controller.register("wait_gate", _wait_gate)
|
||
|
|
|
||
|
|
pipeline = Pipeline("slow")
|
||
|
|
pipeline.node("wait_gate", gate, node_id="root")
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p13")
|
||
|
|
try:
|
||
|
|
outcome = run.wait(timeout=0.2)
|
||
|
|
assert outcome.status is PipelineStatus.RUNNING
|
||
|
|
assert run.node_status("root") in {TaskStatus.STARTING, TaskStatus.RUNNING}
|
||
|
|
assert not run.settled
|
||
|
|
finally:
|
||
|
|
gate.set()
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_pipeline_run_is_registered_on_the_controller() -> None:
|
||
|
|
controller = Controller(max_parallel=4)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
pipeline = Pipeline("reg")
|
||
|
|
pipeline.node("mark", node_id="only")
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="reg-1")
|
||
|
|
run.wait(timeout=_WINDOW)
|
||
|
|
assert controller.get_pipeline("reg-1") is run
|
||
|
|
assert controller.get_pipeline("nope") is None
|
||
|
|
assert controller.pipelines() == (run,)
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_a_diamond_pipeline_works_end_to_end() -> None:
|
||
|
|
"""完整形状：一个源 → 两个分支 → 汇合。"""
|
||
|
|
controller = Controller(max_parallel=8)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
|
||
|
|
pipeline = Pipeline("diamond")
|
||
|
|
source = pipeline.node("mark", node_id="source")
|
||
|
|
left = pipeline.node("mark", node_id="left", depends_on=[source])
|
||
|
|
right = pipeline.node("mark", node_id="right", depends_on=[source])
|
||
|
|
pipeline.node("mark", node_id="sink", depends_on=[left, right])
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p14")
|
||
|
|
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
|
||
|
|
|
||
|
|
order = _finished_order(controller, "p14")
|
||
|
|
assert order[0] == "source"
|
||
|
|
assert order[-1] == "sink"
|
||
|
|
assert "sink" in render_pipeline(run)
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_two_runs_of_the_same_pipeline_are_independent() -> None:
|
||
|
|
controller = Controller(max_parallel=4)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
pipeline = Pipeline("reusable")
|
||
|
|
pipeline.node("mark", node_id="only")
|
||
|
|
|
||
|
|
first = controller.submit_pipeline(pipeline, pipeline_run_id="r1")
|
||
|
|
second = controller.submit_pipeline(pipeline, pipeline_run_id="r2")
|
||
|
|
assert first.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
|
||
|
|
assert second.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
|
||
|
|
assert set(first.task_ids.values()).isdisjoint(second.task_ids.values())
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
def test_pipeline_node_args_reach_the_task() -> None:
|
||
|
|
controller = Controller(max_parallel=4)
|
||
|
|
controller.register("add", _add)
|
||
|
|
pipeline = Pipeline("args")
|
||
|
|
pipeline.node("add", 20, 22, node_id="sum")
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p15")
|
||
|
|
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
|
||
|
|
snapshot = controller.get("p15:sum")
|
||
|
|
assert snapshot is not None
|
||
|
|
assert snapshot.outcome == {"value": 42}
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# 与并发额度协同
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
|
||
|
|
def test_deferred_nodes_do_not_consume_concurrency_slots() -> None:
|
||
|
|
"""等待上游的节点不该占额度——否则一条长管道会把额度全锁死。"""
|
||
|
|
gate = threading.Event()
|
||
|
|
controller = Controller(max_parallel=2)
|
||
|
|
controller.register("wait_gate", _wait_gate)
|
||
|
|
controller.register("mark", _mark)
|
||
|
|
|
||
|
|
pipeline = Pipeline("slots")
|
||
|
|
root = pipeline.node("wait_gate", gate, node_id="root")
|
||
|
|
for index in range(5):
|
||
|
|
pipeline.node("mark", node_id=f"child{index}", depends_on=[root])
|
||
|
|
|
||
|
|
run = controller.submit_pipeline(pipeline, pipeline_run_id="p16")
|
||
|
|
try:
|
||
|
|
assert _wait_until(lambda: run.node_status("root") is TaskStatus.RUNNING)
|
||
|
|
time.sleep(0.1)
|
||
|
|
assert controller.active_count == 1, "只有 root 该占额度"
|
||
|
|
assert controller.queued_count == 0, "下游还在 blocked，不是 queued"
|
||
|
|
|
||
|
|
gate.set()
|
||
|
|
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
|
||
|
|
finally:
|
||
|
|
gate.set()
|
||
|
|
controller.close(timeout=_WINDOW)
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# 工具
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
|
||
|
|
def _finished_order(controller: Controller, run_id: str) -> list[str]:
|
||
|
|
"""按结束时间排出某个管道 run 里各节点的完成顺序。"""
|
||
|
|
snapshots = [
|
||
|
|
item
|
||
|
|
for item in controller.list_tasks()
|
||
|
|
if item.pipeline_run_id == run_id and item.ended_at is not None
|
||
|
|
]
|
||
|
|
snapshots.sort(key=lambda item: item.ended_at or item.created_at) # type: ignore[arg-type,return-value]
|
||
|
|
return [item.name.split("/", 1)[-1] for item in snapshots]
|
||
|
|
|
||
|
|
|
||
|
|
def _wait_until(predicate: Any, timeout: float = _WINDOW) -> bool:
|
||
|
|
deadline = time.monotonic() + timeout
|
||
|
|
while time.monotonic() < deadline:
|
||
|
|
if predicate():
|
||
|
|
return True
|
||
|
|
time.sleep(0.01)
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# 任务体
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
|
||
|
|
def _mark(ctx: Any) -> str:
|
||
|
|
del ctx
|
||
|
|
return "ok"
|
||
|
|
|
||
|
|
|
||
|
|
def _add(ctx: Any, left: int, right: int) -> int:
|
||
|
|
del ctx
|
||
|
|
return left + right
|
||
|
|
|
||
|
|
|
||
|
|
def _boom(ctx: Any) -> str:
|
||
|
|
del ctx
|
||
|
|
raise RuntimeError("故意失败")
|
||
|
|
|
||
|
|
|
||
|
|
def _wait_gate(ctx: Any, gate: threading.Event) -> str:
|
||
|
|
while not gate.wait(0.01):
|
||
|
|
ctx.checkpoint()
|
||
|
|
ctx.checkpoint()
|
||
|
|
return "released"
|