From fbabbc98adaea39b2980d368441ed1916cf4f844 Mon Sep 17 00:00:00 2001 From: gooker_young Date: Wed, 8 Jul 2026 12:28:43 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E4=BB=BB=E5=8A=A1=E4=BC=98=E5=85=88?= =?UTF-8?q?=E7=BA=A7=E8=B0=83=E5=BA=A6=20+=20=E6=9D=A1=E4=BB=B6=E5=88=86?= =?UTF-8?q?=E6=94=AF=E6=9E=84=E9=80=A0=E5=99=A8=20+=20=E5=A4=A7=E5=9B=BE?= =?UTF-8?q?=E5=A2=9E=E9=87=8F=E5=B0=B1=E7=BB=AA=E9=9B=86=E4=BC=98=E5=8C=96?= =?UTF-8?q?=20=E2=80=94=20DependencyRunner=20=E9=87=8D=E5=86=99=E4=B8=BA?= =?UTF-8?q?=E4=BA=8B=E4=BB=B6=E9=A9=B1=E5=8A=A8=E4=BC=98=E5=85=88=E7=BA=A7?= =?UTF-8?q?=E8=B0=83=E5=BA=A6=EF=BC=9B=E6=96=B0=E5=A2=9E=20px.switch/px.br?= =?UTF-8?q?anch=20DAG=20=E6=9E=84=E9=80=A0=E5=99=A8=EF=BC=9B=5Fbuild=5Fdep?= =?UTF-8?q?endency=5Findex=20=E6=9B=BF=E4=BB=A3=20O(N)=20=E6=AF=8F?= =?UTF-8?q?=E8=BD=AE=E6=89=AB=E6=8F=8F=E4=BD=BF=E5=A4=A7=E5=9B=BE=E8=B0=83?= =?UTF-8?q?=E5=BA=A6=E4=BB=8E=20O(N=C2=B2)=20=E9=99=8D=E8=87=B3=20O(N)?= =?UTF-8?q?=EF=BC=8C10k=20=E4=BB=BB=E5=8A=A1=20<1s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .trae/docs/iter-27-functionality-perf.md | 84 +++++++ .trae/skills/pyflowx-development/SKILL.md | 42 ++++ benchmarks/__main__.py | 26 ++ src/pyflowx/__init__.py | 5 +- src/pyflowx/executors.py | 151 ++++++++---- src/pyflowx/pipelines.py | 126 +++++++++- tests/test_advanced_features.py | 49 ++++ tests/test_pipelines.py | 276 +++++++++++++++++++--- 8 files changed, 676 insertions(+), 83 deletions(-) create mode 100644 .trae/docs/iter-27-functionality-perf.md diff --git a/.trae/docs/iter-27-functionality-perf.md b/.trae/docs/iter-27-functionality-perf.md new file mode 100644 index 0000000..64d88e1 --- /dev/null +++ b/.trae/docs/iter-27-functionality-perf.md @@ -0,0 +1,84 @@ +# iter-27：功能扩展与性能优化 + +## 本轮目标 + +用户请求"继续扩展功能和优化性能"，选定 3 个方向： +1. **P1：任务优先级调度** —— DependencyRunner 支持按 `priority` 降序调度就绪任务 +2. **P2：条件分支/switch 模式** —— 新增 `px.switch()` 和 `px.branch()` DAG 构造器 +3. **P3：大图内存与并发优化** —— 消除 `_find_ready()` 的 O(N) 每轮扫描 + +## 改动文件清单 + +| 文件 | 变更 | +|------|------| +| `src/pyflowx/executors.py` | DependencyRunner 重写为事件驱动优先级调度；`_build_dependency_index` 模块级函数；`_on_complete`/`_register_dynamic` 闭包；模块 docstring 更新 | +| `src/pyflowx/pipelines.py` | 新增 `switch()` 和 `branch()` 函数 | +| `src/pyflowx/__init__.py` | 导出 `switch`/`branch` | +| `tests/test_advanced_features.py` | 新增 2 个优先级调度测试 | +| `tests/test_pipelines.py` | 新增 `TestSwitch`（5 测试）和 `TestBranch`（4 测试） | +| `benchmarks/__main__.py` | 新增 `bench_large_graph()` 大图基准（1k/5k/10k 任务） | +| `.trae/skills/pyflowx-development/SKILL.md` | 第三节新增"事件驱动优先级调度"和"增量就绪集优化"；第十七节新增 switch/branch | + +## 关键决策与依据 + +### P1：事件驱动优先级调度 + +**决策**：重写 DependencyRunner 从"一次性创建所有 asyncio Task"为"事件驱动按需创建"。 + +**依据**： +- 旧实现 `asyncio.gather(*[create_task(n) for n in all])` 一次性创建所有 Task， + 无法按优先级排序创建顺序（Task 一旦创建即进入事件循环调度队列）。 +- 新实现用 `completed`/`in_flight`/`remaining` 三集合跟踪生命周期， + 仅当任务依赖全部完成时才创建 asyncio Task，按 `priority` 降序创建。 +- 测试验证：`concurrency_key="g"` + `concurrency_limits={"g": 1}` 串行化无依赖 + 任务，使优先级顺序可观测（`order == ["high", "mid", "low"]`）。 + +### P2：条件分支构造器 + +**决策**：基于已有 `BuiltinConditions` 原语构建 switch/branch，不引入新条件类型。 + +**依据**： +- `DEP_EQUALS`/`DEP_MATCHES`/`NOT`/`AND` 已覆盖 switch/branch 所需语义。 +- `switch(selector, cases, *, default=None)`：每个 case 追加 + `DEP_EQUALS(selector.name, case_key)`；default 追加 + `AND(NOT(DEP_EQUALS(..., k)) for k in cases)`。 +- `branch(selector, predicate, if_true, *, if_false=None)`：if_true 追加 + `DEP_MATCHES`；if_false 追加 `NOT(DEP_MATCHES)`。 +- 用 `dataclasses.replace` 追加 `depends_on` 和 `conditions`，不修改原 spec。 + +### P3：增量就绪集优化 + +**决策**：用 `in_degree` 计数器 + `dependents` 反向邻接表替代每轮 O(N) 扫描。 + +**依据**： +- 旧 `_find_ready()` 每轮扫描全部 `remaining` 任务检查依赖是否满足， + O(N·D) per round、O(N²·D) total。10k 任务链式图（10k 轮）成为瓶颈。 +- 新方案预计算 `in_degree[name]`（尚未完成的硬+软依赖数）+ `dependents[d]` + （反向邻接表）。任务完成时 `_on_complete` 仅遍历 `dependents[name]`（O(D_out)）， + 递减 `in_degree`；降为 0 时加入 `ready` 集合。 +- 动态任务通过 `_register_dynamic` 接入增量结构：spawner 此时未标记 completed + （仍在 in_flight），故计入 unsatisfied；spawner 完成后 `_on_complete` 递减并触发就绪。 +- 提取 `_build_dependency_index` 为模块级函数，控制 `execute` 的 PLR0912 分支数 ≤12。 + +## 验证结果 + +### 基准（dependency 策略，noop 任务） + +| 图形状 | 1k | 5k | 10k | +|--------|-----|-----|------| +| chain | 87ms | 459ms | 948ms | +| diamond | 53ms | 320ms | 741ms | +| wide | 37ms | 247ms | 789ms | + +chain 线性扩展（1k→5k→10k ≈ 5.3x→10.9x），确认 O(N) 调度复杂度。 + +### 门禁 + +- pytest：1705 passed +- coverage：97.21%（branch，≥95% 门槛） +- ruff：All checks passed +- pyrefly：0 errors + +## 遗留事项 + +- 无。P1/P2/P3 全部交付完毕。 diff --git a/.trae/skills/pyflowx-development/SKILL.md b/.trae/skills/pyflowx-development/SKILL.md index 0167432..54314bb 100644 --- a/.trae/skills/pyflowx-development/SKILL.md +++ b/.trae/skills/pyflowx-development/SKILL.md @@ -243,6 +243,39 @@ sequential(500) 从 764 → 930 ops/s（+22%），thread(500) 从 108 → 119 op - **适用边界**：async/dependency 策略对 noop 任务的固有并发开销（事件循环 + 线程池提交）主导，per-task 优化相对影响小；真实 I/O 任务才能体现并行优势。 +### 事件驱动优先级调度（iter-27） + +DependencyRunner 从"一次性创建所有 asyncio Task"重写为"事件驱动按需创建"： +仅当任务依赖全部完成时才创建 asyncio Task，多个就绪任务按 `priority` 降序 +创建，高优先级任务优先获取事件循环与线程池资源。 + +- **`completed` / `in_flight` / `remaining` 三集合**：跟踪任务生命周期。 + `completed` = 已完成（含 SUCCESS/SKIPPED/FAILED）；`in_flight` = asyncio + Task 已创建未完成；`remaining` = 尚未调度。 +- **`concurrency_key=1` 串行化验证优先级**：测试中用 `concurrency_key="g"` + + `concurrency_limits={"g": 1}` 使无依赖任务串行执行，优先级顺序可观测。 + 无 concurrency_key 时任务真并行，优先级仅影响 asyncio Task 创建顺序。 +- **与层策略优先级的差异**：层策略（sequential/thread/async）仅在同层内 + 按优先级排序；dependency 策略跨层，按就绪时刻排序——依赖完成后立即就绪， + 不等同层其他任务。 + +### 增量就绪集优化（iter-27） + +DependencyRunner 主循环每轮调用 `_find_ready()` 扫描全部 `remaining` 任务 +检查依赖是否满足，O(N*D) per round、O(N²·D) total。10k 任务链式图 +（10k 轮调度）下成为瓶颈。 + +- **`_build_dependency_index` 模块级函数**：预计算 `in_degree[name]` + （尚未完成的硬+软依赖数）+ `dependents[d]`（依赖 d 的任务列表，反向邻接表） + + `ready`（in_degree=0 的初始就绪集）。复杂度从 O(N²·D) 降至 O(N·D)。 +- **`_on_complete(name)`**：任务完成后遍历 `dependents[name]`，逐个递减 + `in_degree`；降为 0 时加入 `ready` 集合。每完成一个任务仅 O(D_out) 操作。 +- **`_register_dynamic(raw_spec, spawner)`**：动态生成任务接入增量结构。 + spawner 此时未标记 completed（仍在 in_flight），故计入 unsatisfied； + spawner 完成后 `_on_complete` 会递减并触发就绪。 +- **基准验证**：chain(10000) ~950ms、diamond(10000) ~740ms、wide(10000) + ~790ms。chain 线性扩展（1k:87ms → 5k:459ms → 10k:948ms），确认 O(N) 调度。 + ### 踩坑总结 - **`lru_cache` 对签名内省有 dict lookup 开销**：即便 `functools.lru_cache` @@ -896,6 +929,15 @@ sequential(500) 从 764 → 930 ops/s（+22%），thread(500) 从 108 → 119 op - **闭包工厂模式绑定参数**：`_make_worker_fn(worker, item)` 与 `_make_reduce_fn(reduce, worker_names)` 参照 `imaging._make_step_fn` 模式， 闭包捕获函数参数（非循环变量），避免 Python 闭包"延迟绑定"陷阱。 +- **`switch`/`branch` 条件分支构造器（iter-27）**：基于已有条件原语 + （`DEP_EQUALS`/`DEP_MATCHES`/`NOT`/`AND`）构建，不引入新条件类型。 + `switch(selector, cases, *, default=None)` 为每个 case 追加 + `DEP_EQUALS(selector.name, case_key)` 条件；`default` 追加 + `AND(NOT(DEP_EQUALS(..., k)) for k in cases)` 条件。 + `branch(selector, predicate, if_true, *, if_false=None)` 为 `if_true` 追加 + `DEP_MATCHES(selector.name, predicate)`，为 `if_false` 追加 `NOT(...)`。 + 用 `dataclasses.replace` 在不修改原 spec 的前提下追加 `depends_on` 和 + `conditions`。 ### 设计决策 diff --git a/benchmarks/__main__.py b/benchmarks/__main__.py index 27034bd..0ee3501 100644 --- a/benchmarks/__main__.py +++ b/benchmarks/__main__.py @@ -202,6 +202,31 @@ def bench_dependency() -> None: print_results("执行策略: dependency", results) +def bench_large_graph() -> None: + """大图调度基准：验证增量就绪集优化效果（1k–10k 任务）。 + + 覆盖三种图形状： + * chain —— 深链（每轮仅 1 个就绪，调度轮数 = N） + * wide —— 完全并行（首轮全部就绪，仅 1 轮调度） + * diamond —— 菱形（多层，每层多任务） + """ + results = [] + + def noop() -> None: + pass + + for shape, maker in (("chain", make_chain), ("diamond", make_diamond), ("wide", make_wide)): + for n in (1000, 5000, 10000): + specs = maker(n) + # 替换为 noop fn（避免子进程开销，纯测调度性能） + specs = [TaskSpec(s.name, fn=noop, depends_on=s.depends_on) for s in specs] + graph = Graph.from_specs(specs) + ms, ops = time_it(lambda g=graph: px.run(g, strategy="dependency"), iterations=3, warmup=1) + results.append((f"dependency-{shape}({n})", 3, ms, ops)) + + print_results("大图调度 (dependency, 增量就绪集)", results) + + def bench_cmd_execution() -> None: """cmd 任务执行基准（真实子进程）。""" results = [] @@ -226,6 +251,7 @@ def run_execution() -> None: bench_thread() bench_async() bench_dependency() + bench_large_graph() bench_cmd_execution() diff --git a/src/pyflowx/__init__.py b/src/pyflowx/__init__.py index 889dc58..ae190ea 100644 --- a/src/pyflowx/__init__.py +++ b/src/pyflowx/__init__.py @@ -14,6 +14,7 @@ * :class:`TaskHooks` —— 任务生命周期钩子（pre_run/post_run/on_failure）。 * :class:`GraphDefaults` —— 图级默认值。 * :func:`compose` —— 编程式组合多图。 +* :func:`switch` / :func:`branch` —— 条件分支 DAG 构造器。 * :func:`task_template` —— 批量生成相似 TaskSpec 的工厂。 * 状态后端：:class:`StateBackend`、:class:`MemoryBackend`、:class:`JSONBackend`、:class:`SQLiteBackend`。 @@ -86,7 +87,7 @@ from .notification import ( WebhookNotifier, WeChatNotifier, ) -from .pipelines import command_chain, data_pipeline, fan_out_fan_in +from .pipelines import branch, command_chain, data_pipeline, fan_out_fan_in, switch from .profiling import ProfileReport, TaskProfile from .progress import ProgressCallback, RichProgressMonitor from .report import RunReport @@ -193,6 +194,7 @@ __all__ = [ "ToolSpec", "WeChatNotifier", "WebhookNotifier", + "branch", "build_call_args", "cmd", "command_chain", @@ -214,6 +216,7 @@ __all__ = [ "run_iter", "run_tool", "start_metrics_server", + "switch", "task", "task_template", "tool", diff --git a/src/pyflowx/executors.py b/src/pyflowx/executors.py index 9b1f269..10e709b 100644 --- a/src/pyflowx/executors.py +++ b/src/pyflowx/executors.py @@ -24,13 +24,15 @@ * :class:`SequentialLayerRunner` / :class:`ThreadedLayerRunner` / :class:`AsyncLayerRunner` —— 层级执行器，调用上述模块级辅助。 * :class:`DependencyRunner` —— 依赖驱动调度（非层模型），同样调用模块级辅助。 + 使用 **增量就绪集**（``in_degree`` 计数器 + ``dependents`` 反向邻接表）替代 + 每轮 O(N) 扫描，大图（10k+ 任务）调度开销从 O(N²) 降至 O(N)。 所有策略共享统一异步内核，支持： * :class:`RetryPolicy`（max_attempts/delay/backoff/jitter/retry_on） * 软依赖注入与默认值 * :class:`TaskHooks`（pre_run/post_run/on_failure） * 按任务策略覆盖 -* 优先级排序（同层内） +* 优先级排序（同层内 / 依赖驱动策略中按就绪顺序） * 并发限制（concurrency_key + concurrency_limits） * ``continue_on_error`` * ``cache_key`` 存储键 @@ -245,6 +247,33 @@ def _sort_by_priority(layer: list[str], specs: Mapping[str, TaskSpec[Any]]) -> l return sorted(layer, key=lambda n: -specs[n].priority) +def _build_dependency_index( + remaining: set[str], + all_specs: Mapping[str, TaskSpec[Any]], + completed: set[str], +) -> tuple[dict[str, int], dict[str, list[str]], set[str]]: + """构建增量就绪集索引：in_degree 计数器 + dependents 反向邻接表 + 初始 ready 集合。 + + 用于 :class:`DependencyRunner` 替代每轮 O(N) 扫描 ``remaining``。 + 每轮调度开销从 O(N*D) 降至 O(D_out)，大图（10k+ 任务）显著加速。 + """ + in_degree: dict[str, int] = {} + dependents: dict[str, list[str]] = {name: [] for name in all_specs} + ready: set[str] = set() + for name in remaining: + spec = all_specs[name] + deps = (*spec.depends_on, *spec.soft_depends_on) + unsatisfied = [d for d in deps if d not in completed] + in_degree[name] = len(unsatisfied) + for d in unsatisfied: + if d not in dependents: + dependents[d] = [] + dependents[d].append(name) + if in_degree[name] == 0: + ready.add(name) + return in_degree, dependents, ready + + # ---------------------------------------------------------------------- # # 任务级跳过 / 重试 / 成功处理：模块级函数 # ---------------------------------------------------------------------- # @@ -766,6 +795,9 @@ class DependencyRunner: 所有任务通过 asyncio 并发调度。同步任务卸载到线程池。 + 优先级调度：当多个任务同时就绪（依赖完成）时，按 ``priority`` 降序 + 创建 asyncio Task，高优先级任务优先获取事件循环与线程池资源。 + 本类不继承层 Mixin：依赖驱动调度不是层模型，直接调用模块级共享辅助 函数（:func:`_build_semaphores` / :func:`_get_sem` / :func:`_store_result`）， 职责更清晰。 @@ -782,23 +814,57 @@ class DependencyRunner: cancel_event: threading.Event | CancelToken | None = None, ) -> None: all_names = list(graph.all_specs().keys()) - all_specs = {name: graph.resolved_spec(name) for name in all_names} + all_specs: dict[str, TaskSpec[Any]] = {name: graph.resolved_spec(name) for name in all_names} semaphores = _build_semaphores(all_names, all_specs, asyncio.Semaphore, concurrency_limits) - futures: dict[str, asyncio.Task[TaskResult[Any]]] = {} + + # 事件驱动调度：跟踪 completed / in_flight / remaining。 + # 仅当任务依赖全部完成时才创建 asyncio Task，按优先级降序调度。 + completed: set[str] = set() + in_flight: dict[str, asyncio.Task[TaskResult[Any]]] = {} + remaining: set[str] = set(all_names) + + # 检查点恢复：已在 report 中的 SUCCESS 任务直接标记完成，跳过调度。 + for name in list(remaining): + if name in report.results and report.results[name].status == TaskStatus.SUCCESS: + completed.add(name) + remaining.discard(name) + context[name] = report.results[name].value + + # 增量就绪集：用 in_degree 计数器 + dependents 反向邻接表替代每轮 O(N) 扫描。 + # 每轮调度开销从 O(N*D) 降至 O(D_out)，大图（10k+ 任务）显著加速。 + in_degree, dependents, ready = _build_dependency_index(remaining, all_specs, completed) + + def _on_complete(name: str) -> None: + """任务完成后，递减其依赖者的 in_degree，新就绪的加入 ready 集合。""" + for dependent in dependents.get(name, ()): + in_degree[dependent] -= 1 + if in_degree[dependent] == 0: + ready.add(dependent) + + def _register_dynamic(raw_spec: TaskSpec[Any], spawner: str) -> None: + """注册动态生成的 spec，计算 in_degree 并接入 dependents/semaphores。""" + final_spec = raw_spec + if spawner not in raw_spec.depends_on and spawner not in raw_spec.soft_depends_on: + final_spec = dc_replace(raw_spec, depends_on=(*raw_spec.depends_on, spawner)) + graph._register_single(final_spec) + all_specs[final_spec.name] = graph.resolved_spec(final_spec.name) + remaining.add(final_spec.name) + if final_spec.concurrency_key and final_spec.concurrency_key not in semaphores: + limit = concurrency_limits.get(final_spec.concurrency_key, 1) + semaphores[final_spec.concurrency_key] = asyncio.Semaphore(limit) + # 计算新任务的 in_degree（spawner 可能尚未标记 completed，计入 unsatisfied） + deps = (*final_spec.depends_on, *final_spec.soft_depends_on) + unsatisfied = [d for d in deps if d not in completed] + in_degree[final_spec.name] = len(unsatisfied) + for d in unsatisfied: + if d not in dependents: + dependents[d] = [] + dependents[d].append(final_spec.name) + if in_degree[final_spec.name] == 0: + ready.add(final_spec.name) async def _run_task(name: str) -> TaskResult[Any]: - spec = graph.resolved_spec(name) - # 检查点恢复：已在 report 中的 SUCCESS 任务直接返回 - if name in report.results and report.results[name].status == TaskStatus.SUCCESS: - return report.results[name] - # 等待所有硬依赖完成 - for dep in spec.depends_on: - if dep in futures: - await futures[dep] - # 等待所有软依赖完成（但不检查其状态） - for dep in spec.soft_depends_on: - if dep in futures: - await futures[dep] + spec = all_specs[name] # 取消检查：依赖完成后、执行前检查；已取消则标记 SKIPPED if _is_cancelled(cancel_event) and name not in report.results: @@ -824,47 +890,48 @@ class DependencyRunner: if spawned: spawned_names = [s.name for s in spawned] result = dc_replace(result, value=spawned_names) - # 重新存储已修改的 result context[name] = result.value report.results[name] = result _emit(on_event, result) else: _store_result(name, result, spec, task_ctx, context, report, backend, on_event) - # 注册并调度动态生成的 specs + # 注册动态生成的 specs，接入增量调度结构 for raw_spec in spawned: - # 确保新 spec 依赖生成方（保证顺序与上下文可见性） - final_spec = raw_spec - if name not in raw_spec.depends_on and name not in raw_spec.soft_depends_on: - final_spec = dc_replace(raw_spec, depends_on=(*raw_spec.depends_on, name)) - graph._register_single(final_spec) - # 为新 spec 的 concurrency_key 创建信号量 - if final_spec.concurrency_key and final_spec.concurrency_key not in semaphores: - limit = concurrency_limits.get(final_spec.concurrency_key, 1) - semaphores[final_spec.concurrency_key] = asyncio.Semaphore(limit) - futures[final_spec.name] = loop.create_task(_run_task(final_spec.name)) + _register_dynamic(raw_spec, name) return result loop = asyncio.get_running_loop() - for name in all_names: - futures[name] = loop.create_task(_run_task(name)) - # 循环等待所有 futures（含动态生成的），直到全部完成。 - # asyncio.wait 不像 gather 自动传播异常，需手动检查并 re-raise， - # 匹配 gather 的 fail-fast 行为（首个异常即取消剩余任务并抛出）。 - while futures: - await asyncio.wait(futures.values(), return_when=asyncio.FIRST_COMPLETED) - completed_names = [n for n, f in futures.items() if f.done()] - for done_name in completed_names: - f = futures[done_name] - exc = f.exception() + # 主循环：调度就绪任务 → 等待完成 → 更新 completed → 重复。 + # fail-fast：首个异常即取消剩余任务并抛出（匹配 gather 语义）。 + while remaining or in_flight: + # 按优先级降序调度所有就绪任务 + if ready: + to_schedule = _sort_by_priority(list(ready), all_specs) + ready.clear() + for name in to_schedule: + remaining.discard(name) + in_flight[name] = loop.create_task(_run_task(name)) + + if not in_flight: + if remaining: # pragma: no cover - 图已校验无环，防御性处理 + raise RuntimeError(f"调度死锁：剩余任务 {remaining} 无法就绪") + break + + done, _ = await asyncio.wait(in_flight.values(), return_when=asyncio.FIRST_COMPLETED) + for task in done: + done_name = next(n for n, t in in_flight.items() if t is task) + del in_flight[done_name] + completed.add(done_name) + _on_complete(done_name) + exc = task.exception() if exc is not None: - for remaining in futures.values(): - if not remaining.done(): - remaining.cancel() + for t in in_flight.values(): + if not t.done(): + t.cancel() raise exc - del futures[done_name] # ---------------------------------------------------------------------- # diff --git a/src/pyflowx/pipelines.py b/src/pyflowx/pipelines.py index b58ac95..57a2f0e 100644 --- a/src/pyflowx/pipelines.py +++ b/src/pyflowx/pipelines.py @@ -1,10 +1,12 @@ """通用 DAG 构造器集合. -提供 3 个高层 DAG 构造器, 封装常见的任务编排模式: +提供 5 个高层 DAG 构造器, 封装常见的任务编排模式: * :func:`data_pipeline` —— 函数式数据流水线 (链式函数组合) * :func:`command_chain` —— 命令链 (串行命令执行) * :func:`fan_out_fan_in` —— map-reduce (扇出处理 + 扇入聚合) +* :func:`switch` —— 条件分支 (根据上游结果选择执行路径) +* :func:`branch` —— 二路分支 (if/else 模式) 与 :mod:`pyflowx.imaging` 的 :func:`image_pipeline` 同级, 复用 :meth:`Graph.chain` / :meth:`Graph.map` 的依赖链接模式, 但提供更高层次 @@ -21,13 +23,15 @@ from __future__ import annotations from collections.abc import Callable, Mapping, Sequence +from dataclasses import replace from pathlib import Path from typing import Any +from .conditions import BuiltinConditions from .graph import Graph from .task import Context, TaskSpec -__all__ = ["command_chain", "data_pipeline", "fan_out_fan_in"] +__all__ = ["branch", "command_chain", "data_pipeline", "fan_out_fan_in", "switch"] # ---------------------------------------------------------------------- # @@ -308,3 +312,121 @@ def _make_reduce_fn( return reduce(results) return _run + + +# ---------------------------------------------------------------------- # +# 条件分支 +# ---------------------------------------------------------------------- # +def switch( + selector: TaskSpec[Any], + cases: Mapping[str, TaskSpec[Any]], + *, + default: TaskSpec[Any] | None = None, +) -> Graph: + """条件分支 DAG：根据 ``selector`` 任务的结果值选择执行哪个 case. + + 每个 case 任务自动依赖 ``selector``，并附加 + :meth:`BuiltinConditions.DEP_EQUALS` 条件：仅当 ``selector`` 返回值 + 等于 case key 时才执行，其余 case 被跳过 (SKIPPED)。 + + Parameters + ---------- + selector : TaskSpec + 其返回值决定执行哪个分支的任务。 + cases : Mapping[str, TaskSpec] + ``结果值 → 任务`` 映射。每个 case 任务的 fn 参数名匹配 + ``selector.name`` 时会收到 selector 的结果。 + default : TaskSpec | None + 无 case 匹配时执行的任务（可选）。未提供时所有 case 不匹配则 + 全部跳过。 + + Returns + ------- + Graph + 包含 selector + 所有 case (+ default) 的 DAG。 + + Example + ------- + >>> @px.task + ... def classify(data: list[int]) -> str: ... + >>> @px.task + ... def handle_positive(classify: str) -> str: ... + >>> @px.task + ... def handle_negative(classify: str) -> str: ... + >>> graph = px.switch(classify, {"positive": handle_positive, "negative": handle_negative}) + >>> report = px.run(graph) + >>> # 仅 handle_positive 或 handle_negative 之一执行 + """ + specs: list[TaskSpec[Any] | str] = [selector] + sel_name = selector.name + + for case_key, case_spec in cases.items(): + cond = BuiltinConditions.DEP_EQUALS(sel_name, case_key) + deps = (*case_spec.depends_on, sel_name) if sel_name not in case_spec.depends_on else case_spec.depends_on + conds = (*case_spec.conditions, cond) if case_spec.conditions else (cond,) + specs.append(replace(case_spec, depends_on=deps, conditions=conds)) + + if default is not None: + # default 条件：所有 case 都不匹配 + not_match_conds = [BuiltinConditions.NOT(BuiltinConditions.DEP_EQUALS(sel_name, k)) for k in cases] + default_cond = BuiltinConditions.AND(*not_match_conds) if len(not_match_conds) > 1 else not_match_conds[0] + deps = (*default.depends_on, sel_name) if sel_name not in default.depends_on else default.depends_on + conds = (*default.conditions, default_cond) if default.conditions else (default_cond,) + specs.append(replace(default, depends_on=deps, conditions=conds)) + + return Graph.from_specs(specs) + + +def branch( + selector: TaskSpec[Any], + predicate: Callable[[Any], bool], + if_true: TaskSpec[Any], + *, + if_false: TaskSpec[Any] | None = None, +) -> Graph: + """二路分支 DAG (if/else)：根据 ``predicate`` 判断 ``selector`` 结果走哪条路径. + + ``if_true`` 任务附加 :meth:`BuiltinConditions.DEP_MATCHES` 条件 + (predicate 返回 True 时执行)；``if_false`` 任务附加取反条件 + (predicate 返回 False 时执行)。未匹配的分支被跳过 (SKIPPED)。 + + Parameters + ---------- + selector : TaskSpec + 其返回值由 ``predicate`` 判断。 + predicate : Callable[[Any], bool] + 判断函数，接收 ``selector`` 的返回值，返回布尔值。 + if_true : TaskSpec + ``predicate`` 返回 True 时执行的任务。 + if_false : TaskSpec | None + ``predicate`` 返回 False 时执行的任务（可选）。 + + Returns + ------- + Graph + 包含 selector + if_true (+ if_false) 的 DAG。 + + Example + ------- + >>> @px.task + ... def compute(data: list[int]) -> int: return sum(data) + >>> @px.task + ... def handle_positive(compute: int) -> str: ... + >>> @px.task + ... def handle_non_positive(compute: int) -> str: ... + >>> graph = px.branch(compute, lambda x: x > 0, handle_positive, if_false=handle_non_positive) + """ + sel_name = selector.name + match_cond = BuiltinConditions.DEP_MATCHES(sel_name, predicate) + + true_deps = (*if_true.depends_on, sel_name) if sel_name not in if_true.depends_on else if_true.depends_on + true_conds = (*if_true.conditions, match_cond) if if_true.conditions else (match_cond,) + specs: list[TaskSpec[Any]] = [selector, replace(if_true, depends_on=true_deps, conditions=true_conds)] + + if if_false is not None: + not_match_cond = BuiltinConditions.NOT(match_cond) + false_deps = (*if_false.depends_on, sel_name) if sel_name not in if_false.depends_on else if_false.depends_on + false_conds = (*if_false.conditions, not_match_cond) if if_false.conditions else (not_match_cond,) + specs.append(replace(if_false, depends_on=false_deps, conditions=false_conds)) + + return Graph.from_specs(specs) diff --git a/tests/test_advanced_features.py b/tests/test_advanced_features.py index 697a122..f17bcdb 100644 --- a/tests/test_advanced_features.py +++ b/tests/test_advanced_features.py @@ -722,6 +722,55 @@ class TestPriority: spec = px.TaskSpec("a", lambda: "ok") assert spec.priority == 0 + def test_priority_orders_in_dependency_strategy(self) -> None: + """dependency 策略也应按优先级降序调度就绪任务。 + + 用 concurrency_key=1 串行化无依赖任务，使优先级顺序可观测。 + """ + order: list[str] = [] + + def make_fn(name: str) -> Any: + def fn() -> str: + order.append(name) + return name + + return fn + + graph = px.Graph.from_specs( + [ + px.TaskSpec("low", make_fn("low"), priority=1, concurrency_key="g"), + px.TaskSpec("high", make_fn("high"), priority=10, concurrency_key="g"), + px.TaskSpec("mid", make_fn("mid"), priority=5, concurrency_key="g"), + ] + ) + report = px.run(graph, strategy="dependency", concurrency_limits={"g": 1}) + assert report.success + assert order == ["high", "mid", "low"] + + def test_priority_orders_ready_tasks_after_dep_completes(self) -> None: + """依赖完成后，多个就绪任务按优先级降序调度。""" + order: list[str] = [] + + def make_fn(name: str) -> Any: + def fn() -> str: + order.append(name) + return name + + return fn + + graph = px.Graph.from_specs( + [ + px.TaskSpec("root", make_fn("root")), + px.TaskSpec("low", make_fn("low"), depends_on=("root",), priority=1, concurrency_key="g"), + px.TaskSpec("high", make_fn("high"), depends_on=("root",), priority=10, concurrency_key="g"), + px.TaskSpec("mid", make_fn("mid"), depends_on=("root",), priority=5, concurrency_key="g"), + ] + ) + report = px.run(graph, strategy="dependency", concurrency_limits={"g": 1}) + assert report.success + # root 先执行，然后 high/mid/low 按优先级降序 + assert order == ["root", "high", "mid", "low"] + # ---------------------------------------------------------------------- # # continue_on_error 容错 diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py index 720787d..c6c2b80 100644 --- a/tests/test_pipelines.py +++ b/tests/test_pipelines.py @@ -1,13 +1,15 @@ """通用 DAG 构造器测试. 验证 :func:`data_pipeline` / :func:`command_chain` / :func:`fan_out_fan_in` -三个构造器的拓扑生成、执行结果、错误处理与 DAG 组合能力. +/ :func:`switch` / :func:`branch` 五个构造器的拓扑生成、执行结果、错误处理与 DAG 组合能力. 测试组织: * ``TestDataPipeline`` —— 数据流水线拓扑/执行/命名/错误 * ``TestCommandChain`` —— 命令链拓扑/执行/透传/错误 * ``TestFanOutFanIn`` —— map-reduce 拓扑/执行/顺序/错误 * ``TestPipelinesWithFileops`` —— 与 fileops 模块联用组合示例 +* ``TestSwitch`` —— 条件分支 (多路分发) 拓扑/执行/默认/条件共存 +* ``TestBranch`` —— 二路分支 (if/else) 拓扑/执行/predicate """ from __future__ import annotations @@ -91,12 +93,10 @@ class TestDataPipeline: def test_lambda_fallback_naming(self) -> None: """lambda 函数回退为 step 命名.""" - graph = data_pipeline( - [ - lambda: 1, - lambda step: step + 10, - ] - ) + graph = data_pipeline([ + lambda: 1, + lambda step: step + 10, + ]) specs = graph.all_specs() assert "step" in specs assert "step_1" in specs @@ -190,13 +190,11 @@ class TestCommandChain: def test_basic_chain_topology(self) -> None: """三命令链式依赖.""" - graph = command_chain( - [ - ["echo", "first"], - ["echo", "second"], - ["echo", "third"], - ] - ) + graph = command_chain([ + ["echo", "first"], + ["echo", "second"], + ["echo", "third"], + ]) specs = graph.all_specs() assert len(specs) == 3 @@ -207,24 +205,20 @@ class TestCommandChain: def test_auto_naming_with_first_arg(self) -> None: """自动命名用 cmd_{i:02d}_{first_arg}.""" - graph = command_chain( - [ - ["echo", "hello"], - ["echo", "world"], - ] - ) + graph = command_chain([ + ["echo", "hello"], + ["echo", "world"], + ]) specs = graph.all_specs() assert "cmd_00_echo" in specs assert "cmd_01_echo" in specs def test_string_command_auto_naming(self) -> None: """shell 字符串命令自动命名取首个 token.""" - graph = command_chain( - [ - "echo hello", - "ls -la", - ] - ) + graph = command_chain([ + "echo hello", + "ls -la", + ]) specs = graph.all_specs() assert "cmd_00_echo" in specs assert "cmd_01_ls" in specs @@ -243,12 +237,10 @@ class TestCommandChain: def test_execution_succeeds(self) -> None: """命令链实际执行成功.""" - graph = command_chain( - [ - [sys.executable, "-c", "print('first')"], - [sys.executable, "-c", "print('second')"], - ] - ) + graph = command_chain([ + [sys.executable, "-c", "print('first')"], + [sys.executable, "-c", "print('second')"], + ]) report = px.run(graph, strategy="sequential") assert report.success @@ -312,12 +304,10 @@ class TestCommandChain: def test_command_failure_aborts_chain(self) -> None: """命令失败时抛 TaskFailedError (continue_on_error=False).""" - graph = command_chain( - [ - [sys.executable, "-c", "import sys; sys.exit(1)"], - [sys.executable, "-c", "print('should not run')"], - ] - ) + graph = command_chain([ + [sys.executable, "-c", "import sys; sys.exit(1)"], + [sys.executable, "-c", "print('should not run')"], + ]) with pytest.raises(px.TaskFailedError): px.run(graph, strategy="sequential") @@ -578,3 +568,213 @@ class TestPipelinesWithFileops: report = px.run(graph, strategy="sequential") assert report.success assert report["third"] == "first\nsecond\n" + + +# ---------------------------------------------------------------------- # +# 条件分支: switch / branch +# ---------------------------------------------------------------------- # +class TestSwitch: + """测试 px.switch() 条件分支构造器.""" + + def test_switch_executes_matching_case(self) -> None: + """selector 结果匹配的 case 执行，其余跳过。""" + + def classify() -> str: + return "positive" + + def handle_positive(classify: str) -> str: + return f"pos:{classify}" + + def handle_negative(classify: str) -> str: + return f"neg:{classify}" + + graph = px.switch( + px.TaskSpec("classify", classify), + { + "positive": px.TaskSpec("handle_positive", handle_positive), + "negative": px.TaskSpec("handle_negative", handle_negative), + }, + ) + report = px.run(graph) + assert report.success + assert report["classify"] == "positive" + assert report["handle_positive"] == "pos:positive" + assert report.results["handle_negative"].status == px.TaskStatus.SKIPPED + + def test_switch_negative_case(self) -> None: + """selector 返回 negative 时执行 negative case。""" + + def classify() -> str: + return "negative" + + def handle_positive(classify: str) -> str: + return "pos" + + def handle_negative(classify: str) -> str: + return "neg" + + graph = px.switch( + px.TaskSpec("classify", classify), + { + "positive": px.TaskSpec("handle_positive", handle_positive), + "negative": px.TaskSpec("handle_negative", handle_negative), + }, + ) + report = px.run(graph) + assert report.success + assert report["handle_negative"] == "neg" + assert report.results["handle_positive"].status == px.TaskStatus.SKIPPED + + def test_switch_default_when_no_match(self) -> None: + """无 case 匹配时执行 default 任务。""" + + def classify() -> str: + return "unknown" + + def handle_a(classify: str) -> str: + return "a" + + def handle_default(classify: str) -> str: + return "default" + + graph = px.switch( + px.TaskSpec("classify", classify), + {"a": px.TaskSpec("handle_a", handle_a)}, + default=px.TaskSpec("handle_default", handle_default), + ) + report = px.run(graph) + assert report.success + assert report.results["handle_a"].status == px.TaskStatus.SKIPPED + assert report["handle_default"] == "default" + + def test_switch_all_skipped_without_default(self) -> None: + """无 case 匹配且无 default 时所有 case 跳过。""" + + def classify() -> str: + return "unknown" + + def handle_a(classify: str) -> str: + return "a" + + graph = px.switch( + px.TaskSpec("classify", classify), + {"a": px.TaskSpec("handle_a", handle_a)}, + ) + report = px.run(graph) + assert report.success + assert report.results["handle_a"].status == px.TaskStatus.SKIPPED + + def test_switch_preserves_existing_conditions(self) -> None: + """switch 附加的条件与已有条件共存（AND 语义由执行器保证）。""" + + def classify() -> str: + return "go" + + def handle_go(classify: str) -> str: + return "went" + + spec = px.TaskSpec("handle_go", handle_go, conditions=(px.IS_LINUX,)) + graph = px.switch( + px.TaskSpec("classify", classify), + {"go": spec}, + ) + report = px.run(graph) + # IS_LINUX + DEP_EQUALS("classify","go") 均满足时才执行 + if px.IS_LINUX({}): + assert report["handle_go"] == "went" + else: + assert report.results["handle_go"].status == px.TaskStatus.SKIPPED + + +class TestBranch: + """测试 px.branch() 二路分支构造器.""" + + def test_branch_true_path_executes(self) -> None: + """predicate 返回 True 时执行 if_true。""" + + def compute() -> int: + return 42 + + def handle_positive(compute: int) -> str: + return f"positive:{compute}" + + def handle_non_positive(compute: int) -> str: + return "non-positive" + + graph = px.branch( + px.TaskSpec("compute", compute), + lambda x: x > 0, + px.TaskSpec("handle_positive", handle_positive), + if_false=px.TaskSpec("handle_non_positive", handle_non_positive), + ) + report = px.run(graph) + assert report.success + assert report["handle_positive"] == "positive:42" + assert report.results["handle_non_positive"].status == px.TaskStatus.SKIPPED + + def test_branch_false_path_executes(self) -> None: + """predicate 返回 False 时执行 if_false。""" + + def compute() -> int: + return -5 + + def handle_positive(compute: int) -> str: + return "positive" + + def handle_non_positive(compute: int) -> str: + return f"non-positive:{compute}" + + graph = px.branch( + px.TaskSpec("compute", compute), + lambda x: x > 0, + px.TaskSpec("handle_positive", handle_positive), + if_false=px.TaskSpec("handle_non_positive", handle_non_positive), + ) + report = px.run(graph) + assert report.success + assert report["handle_non_positive"] == "non-positive:-5" + assert report.results["handle_positive"].status == px.TaskStatus.SKIPPED + + def test_branch_without_if_false(self) -> None: + """无 if_false 时，predicate 返回 False 则 if_true 被跳过。""" + + def compute() -> int: + return -1 + + def handle_positive(compute: int) -> str: + return "positive" + + graph = px.branch( + px.TaskSpec("compute", compute), + lambda x: x > 0, + px.TaskSpec("handle_positive", handle_positive), + ) + report = px.run(graph) + assert report.success + assert report.results["handle_positive"].status == px.TaskStatus.SKIPPED + + def test_branch_with_predicate_function(self) -> None: + """使用命名函数作为 predicate。""" + + def compute() -> list[int]: + return [1, 2, 3] + + def is_long(items: list[int]) -> bool: + return len(items) > 2 + + def handle_long(compute: list[int]) -> str: + return "long" + + def handle_short(compute: list[int]) -> str: + return "short" + + graph = px.branch( + px.TaskSpec("compute", compute), + is_long, + px.TaskSpec("handle_long", handle_long), + if_false=px.TaskSpec("handle_short", handle_short), + ) + report = px.run(graph) + assert report.success + assert report["handle_long"] == "long" + assert report.results["handle_short"].status == px.TaskStatus.SKIPPED