420 lines
13 KiB
Python
420 lines
13 KiB
Python
"""PyFlowX 性能基准套件.
|
||
|
||
用法::
|
||
|
||
python -m benchmarks # 运行全部基准
|
||
python -m benchmarks graph # 仅图构建基准
|
||
python -m benchmarks execution # 仅执行基准
|
||
python -m benchmarks context # 仅上下文注入基准
|
||
python -m benchmarks storage # 仅状态后端基准
|
||
python -m benchmarks advanced # 仅高级基准
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import sys
|
||
import tempfile
|
||
from collections.abc import Callable
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from rich.console import Console
|
||
|
||
import pyflowx as px
|
||
from benchmarks import print_results, time_it
|
||
from benchmarks.bench_advanced import run_advanced
|
||
from pyflowx import Graph, GraphDefaults, RetryPolicy, TaskSpec
|
||
from pyflowx.context import build_call_args
|
||
from pyflowx.storage import JSONBackend, MemoryBackend, SQLiteBackend
|
||
|
||
# ============================================================================
|
||
# 图生成工具
|
||
# ============================================================================
|
||
|
||
|
||
def make_chain(n: int) -> list[TaskSpec]:
|
||
"""生成 n 个任务的链式 DAG。"""
|
||
specs = [TaskSpec(f"t{i}", cmd=["true"]) for i in range(n)]
|
||
for i in range(1, n):
|
||
specs[i] = TaskSpec(f"t{i}", cmd=["true"], depends_on=(f"t{i - 1}",))
|
||
return specs
|
||
|
||
|
||
def make_diamond(n: int) -> list[TaskSpec]:
|
||
"""生成 n 个任务的菱形 DAG（每层宽度约 sqrt(n)）。"""
|
||
width = max(1, int(math.sqrt(n)))
|
||
specs: list[TaskSpec] = []
|
||
prev_layer: list[str] = []
|
||
layer = 0
|
||
count = 0
|
||
while count < n:
|
||
cur_layer: list[str] = []
|
||
for j in range(width):
|
||
if count >= n:
|
||
break
|
||
name = f"l{layer}_t{j}"
|
||
deps = tuple(prev_layer) if prev_layer else ()
|
||
specs.append(TaskSpec(name, cmd=["true"], depends_on=deps))
|
||
cur_layer.append(name)
|
||
count += 1
|
||
prev_layer = cur_layer
|
||
layer += 1
|
||
return specs
|
||
|
||
|
||
def make_wide(n: int) -> list[TaskSpec]:
|
||
"""生成 n 个独立任务（无依赖，最大并行度）。"""
|
||
return [TaskSpec(f"t{i}", cmd=["true"]) for i in range(n)]
|
||
|
||
|
||
# ============================================================================
|
||
# 基准：图构建
|
||
# ============================================================================
|
||
|
||
|
||
def bench_construction() -> None:
|
||
"""图构建（from_specs + validate）基准。"""
|
||
results = []
|
||
for n in (10, 100, 500, 1000):
|
||
specs = make_chain(n)
|
||
ms, _ = time_it(lambda s=specs: Graph.from_specs(s), iterations=20)
|
||
results.append((f"chain({n})", 20, ms, 1.0 / (ms / 1000) if ms > 0 else 0))
|
||
|
||
for n in (10, 100, 500, 1000):
|
||
specs = make_diamond(n)
|
||
ms, _ = time_it(lambda s=specs: Graph.from_specs(s), iterations=20)
|
||
results.append((f"diamond({n})", 20, ms, 1.0 / (ms / 1000) if ms > 0 else 0))
|
||
|
||
print_results("图构建 (from_specs + validate)", results)
|
||
|
||
|
||
def bench_layers() -> None:
|
||
"""拓扑分层基准（冷启动 vs 缓存命中）。"""
|
||
results = []
|
||
for n in (100, 500, 1000):
|
||
specs = make_diamond(n)
|
||
graph = Graph.from_specs(specs)
|
||
|
||
def _cold(g: Graph = graph) -> None:
|
||
g._layers_cache = None # type: ignore[attr-defined]
|
||
g.layers()
|
||
|
||
ms_cold, ops_cold = time_it(_cold, iterations=50, warmup=5)
|
||
results.append((f"layers(cold,{n})", 50, ms_cold, ops_cold))
|
||
|
||
ms_hot, ops_hot = time_it(lambda g=graph: g.layers(), iterations=200, warmup=10)
|
||
results.append((f"layers(cached,{n})", 200, ms_hot, ops_hot))
|
||
|
||
print_results("拓扑分层 (layers)", results)
|
||
|
||
|
||
def bench_resolved_spec() -> None:
|
||
"""resolved_spec 缓存命中基准。"""
|
||
results = []
|
||
for n in (100, 500, 1000):
|
||
specs = make_chain(n)
|
||
defaults = GraphDefaults(retry=RetryPolicy(max_attempts=2))
|
||
graph = Graph.from_specs(specs, defaults=defaults)
|
||
name = f"t{n // 2}"
|
||
ms, ops = time_it(lambda g=graph, nm=name: g.resolved_spec(nm), iterations=500, warmup=20)
|
||
results.append((f"resolved_spec(cached,{n})", 500, ms, ops))
|
||
|
||
print_results("resolved_spec (缓存命中)", results)
|
||
|
||
|
||
def run_graph() -> None:
|
||
"""运行全部图基准。"""
|
||
bench_construction()
|
||
bench_layers()
|
||
bench_resolved_spec()
|
||
|
||
|
||
# ============================================================================
|
||
# 基准：任务执行
|
||
# ============================================================================
|
||
|
||
|
||
def bench_sequential() -> None:
|
||
"""sequential 策略执行基准。"""
|
||
results = []
|
||
|
||
def noop() -> None:
|
||
pass
|
||
|
||
for n in (50, 200, 500):
|
||
specs = [TaskSpec(f"t{i}", fn=noop) for i in range(n)]
|
||
graph = Graph.from_specs(specs)
|
||
ms, ops = time_it(lambda g=graph: px.run(g, strategy="sequential"), iterations=10, warmup=2)
|
||
results.append((f"sequential({n})", 10, ms, ops))
|
||
|
||
print_results("执行策略: sequential", results)
|
||
|
||
|
||
def bench_thread() -> None:
|
||
"""thread 策略执行基准。"""
|
||
results = []
|
||
|
||
def noop() -> None:
|
||
pass
|
||
|
||
for n in (50, 200, 500):
|
||
specs = [TaskSpec(f"t{i}", fn=noop) for i in range(n)]
|
||
graph = Graph.from_specs(specs)
|
||
ms, ops = time_it(lambda g=graph: px.run(g, strategy="thread"), iterations=10, warmup=2)
|
||
results.append((f"thread({n})", 10, ms, ops))
|
||
|
||
print_results("执行策略: thread", results)
|
||
|
||
|
||
def bench_async() -> None:
|
||
"""async 策略执行基准。"""
|
||
results = []
|
||
|
||
def noop() -> None:
|
||
pass
|
||
|
||
for n in (50, 200, 500):
|
||
specs = [TaskSpec(f"t{i}", fn=noop) for i in range(n)]
|
||
graph = Graph.from_specs(specs)
|
||
ms, ops = time_it(lambda g=graph: px.run(g, strategy="async"), iterations=10, warmup=2)
|
||
results.append((f"async({n})", 10, ms, ops))
|
||
|
||
print_results("执行策略: async", results)
|
||
|
||
|
||
def bench_dependency() -> None:
|
||
"""dependency 策略执行基准。"""
|
||
results = []
|
||
|
||
def noop() -> None:
|
||
pass
|
||
|
||
for n in (50, 200, 500):
|
||
specs = [TaskSpec(f"t{i}", fn=noop) for i in range(n)]
|
||
graph = Graph.from_specs(specs)
|
||
ms, ops = time_it(lambda g=graph: px.run(g, strategy="dependency"), iterations=10, warmup=2)
|
||
results.append((f"dependency({n})", 10, ms, ops))
|
||
|
||
print_results("执行策略: dependency", results)
|
||
|
||
|
||
def bench_cmd_execution() -> None:
|
||
"""cmd 任务执行基准（真实子进程）。"""
|
||
results = []
|
||
for n in (10, 50, 100):
|
||
specs = [TaskSpec(f"t{i}", cmd=["true"]) for i in range(n)]
|
||
graph = Graph.from_specs(specs)
|
||
ms, ops = time_it(lambda g=graph: px.run(g, strategy="sequential"), iterations=5, warmup=1)
|
||
results.append((f"cmd-sequential({n})", 5, ms, ops))
|
||
|
||
for n in (10, 50, 100):
|
||
specs = [TaskSpec(f"t{i}", cmd=["true"]) for i in range(n)]
|
||
graph = Graph.from_specs(specs)
|
||
ms, ops = time_it(lambda g=graph: px.run(g, strategy="thread", max_workers=8), iterations=5, warmup=1)
|
||
results.append((f"cmd-thread({n})", 5, ms, ops))
|
||
|
||
print_results("cmd 任务执行 (['true'])", results)
|
||
|
||
|
||
def run_execution() -> None:
|
||
"""运行全部执行基准。"""
|
||
bench_sequential()
|
||
bench_thread()
|
||
bench_async()
|
||
bench_dependency()
|
||
bench_cmd_execution()
|
||
|
||
|
||
# ============================================================================
|
||
# 基准：上下文注入
|
||
# ============================================================================
|
||
|
||
|
||
def bench_context_no_deps() -> None:
|
||
"""无依赖 fn 任务的上下文注入基准。"""
|
||
results = []
|
||
|
||
def noop() -> None:
|
||
pass
|
||
|
||
spec = TaskSpec("noop", fn=noop)
|
||
context: dict[str, Any] = {}
|
||
ms, ops = time_it(lambda s=spec, c=context: build_call_args(s, c), iterations=2000, warmup=100)
|
||
results.append(("fn(no-deps)", 2000, ms, ops))
|
||
|
||
# cmd 任务快速路径
|
||
spec_cmd = TaskSpec("cmd", cmd=["true"])
|
||
ms, ops = time_it(lambda s=spec_cmd, c=context: build_call_args(s, c), iterations=2000, warmup=100)
|
||
results.append(("cmd(fast-path)", 2000, ms, ops))
|
||
|
||
print_results("上下文注入 (build_call_args)", results)
|
||
|
||
|
||
def bench_context_with_deps() -> None:
|
||
"""有依赖 fn 任务的上下文注入基准。"""
|
||
results = []
|
||
|
||
def consumer(a: int, b: int) -> int:
|
||
return a + b
|
||
|
||
spec = TaskSpec("consumer", fn=consumer, depends_on=("a", "b"))
|
||
context = {"a": 1, "b": 2, "c": 3}
|
||
ms, ops = time_it(lambda s=spec, c=context: build_call_args(s, c), iterations=2000, warmup=100)
|
||
results.append(("fn(2-deps)", 2000, ms, ops))
|
||
|
||
# Context 标注
|
||
from pyflowx.task import Context
|
||
|
||
def ctx_fn(ctx: Context) -> int:
|
||
return sum(ctx.values())
|
||
|
||
spec_ctx = TaskSpec("ctx", fn=ctx_fn, depends_on=("a", "b"))
|
||
ms, ops = time_it(lambda s=spec_ctx, c=context: build_call_args(s, c), iterations=2000, warmup=100)
|
||
results.append(("fn(Context-annotated)", 2000, ms, ops))
|
||
|
||
# **kwargs
|
||
def kwargs_fn(**kwargs: int) -> int:
|
||
return sum(kwargs.values())
|
||
|
||
spec_kw = TaskSpec("kw", fn=kwargs_fn, depends_on=("a", "b"))
|
||
ms, ops = time_it(lambda s=spec_kw, c=context: build_call_args(s, c), iterations=2000, warmup=100)
|
||
results.append(("fn(**kwargs)", 2000, ms, ops))
|
||
|
||
print_results("上下文注入 (有依赖)", results)
|
||
|
||
|
||
def run_context() -> None:
|
||
"""运行全部上下文注入基准。"""
|
||
bench_context_no_deps()
|
||
bench_context_with_deps()
|
||
|
||
|
||
# ============================================================================
|
||
# 基准：状态后端
|
||
# ============================================================================
|
||
|
||
|
||
def bench_storage() -> None:
|
||
"""状态后端 save/load 基准。"""
|
||
results = []
|
||
|
||
# MemoryBackend
|
||
mem = MemoryBackend()
|
||
ms, ops = time_it(lambda: mem.save("key", "value"), iterations=1000, warmup=50)
|
||
results.append(("MemoryBackend.save", 1000, ms, ops))
|
||
|
||
ms, ops = time_it(mem.load, iterations=1000, warmup=50)
|
||
results.append(("MemoryBackend.load", 1000, ms, ops))
|
||
|
||
# JSONBackend（batch 模式）
|
||
tmp_dir = tempfile.mkdtemp()
|
||
json_path = str(Path(tmp_dir) / "state.json")
|
||
|
||
json_backend = JSONBackend(json_path)
|
||
with json_backend.batch():
|
||
for i in range(100):
|
||
json_backend.save(f"task_{i}", f"result_{i}")
|
||
|
||
def _json_save() -> None:
|
||
jb = JSONBackend(json_path)
|
||
with jb.batch():
|
||
for i in range(10):
|
||
jb.save(f"bench_{i}", f"val_{i}")
|
||
|
||
ms, ops = time_it(_json_save, iterations=50, warmup=5)
|
||
results.append(("JSONBackend.save(batch=10)", 50, ms, ops))
|
||
|
||
# 复杂 value（嵌套 dict）—— 展示 batch 模式延迟验证的优化效果
|
||
complex_value = {
|
||
"output": {"path": "/data/result.json", "size": 1024, "checksum": "abc123"},
|
||
"metrics": {"accuracy": 0.95, "latency_ms": 120, "samples": 1000},
|
||
"artifacts": [f"artifact_{j}.bin" for j in range(10)],
|
||
"metadata": {"version": "1.0", "tags": ["trained", "validated"], "created_at": "2026-07-07"},
|
||
}
|
||
|
||
def _json_save_complex() -> None:
|
||
jb = JSONBackend(json_path)
|
||
with jb.batch():
|
||
for i in range(10):
|
||
jb.save(f"bench_complex_{i}", complex_value)
|
||
|
||
ms, ops = time_it(_json_save_complex, iterations=50, warmup=5)
|
||
results.append(("JSONBackend.save(batch=10,complex)", 50, ms, ops))
|
||
|
||
ms, ops = time_it(json_backend.load, iterations=200, warmup=10)
|
||
results.append(("JSONBackend.load", 200, ms, ops))
|
||
|
||
# SQLiteBackend
|
||
db_path = str(Path(tmp_dir) / "state.db")
|
||
|
||
sqlite_backend = SQLiteBackend(db_path)
|
||
with sqlite_backend.batch():
|
||
for i in range(100):
|
||
sqlite_backend.save(f"task_{i}", f"result_{i}")
|
||
|
||
def _sqlite_save() -> None:
|
||
sb = SQLiteBackend(db_path)
|
||
with sb.batch():
|
||
for i in range(10):
|
||
sb.save(f"bench_{i}", f"val_{i}")
|
||
|
||
ms, ops = time_it(_sqlite_save, iterations=50, warmup=5)
|
||
results.append(("SQLiteBackend.save(batch=10)", 50, ms, ops))
|
||
|
||
ms, ops = time_it(sqlite_backend.load, iterations=200, warmup=10)
|
||
results.append(("SQLiteBackend.load", 200, ms, ops))
|
||
|
||
print_results("状态后端 (save/load)", results)
|
||
|
||
# 清理临时目录
|
||
import shutil
|
||
|
||
shutil.rmtree(tmp_dir, ignore_errors=True)
|
||
|
||
|
||
def run_storage() -> None:
|
||
"""运行全部状态后端基准。"""
|
||
bench_storage()
|
||
|
||
|
||
# ============================================================================
|
||
# 主入口
|
||
# ============================================================================
|
||
|
||
|
||
BENCH_MODULES: dict[str, Callable[[], None]] = {
|
||
"graph": run_graph,
|
||
"execution": run_execution,
|
||
"context": run_context,
|
||
"storage": run_storage,
|
||
"advanced": run_advanced,
|
||
}
|
||
|
||
|
||
def main(argv: list[str] | None = None) -> int:
|
||
"""CLI 入口。"""
|
||
args = argv if argv is not None else sys.argv[1:]
|
||
console = Console()
|
||
console.print("[bold cyan]PyFlowX 性能基准套件[/bold cyan]\n")
|
||
|
||
if not args or args[0] in ("--all", "-a"):
|
||
for name, fn in BENCH_MODULES.items():
|
||
console.print(f"[bold]运行: {name}[/bold]")
|
||
fn()
|
||
elif args[0] in BENCH_MODULES:
|
||
BENCH_MODULES[args[0]]()
|
||
elif args[0] in ("--help", "-h"):
|
||
console.print("用法: python -m benchmarks [graph|execution|context|storage]")
|
||
console.print(" 无参数 = 运行全部基准")
|
||
else:
|
||
console.print(f"[red]未知基准模块: {args[0]}[/red]")
|
||
console.print(f"可用: {', '.join(BENCH_MODULES)}")
|
||
return 1
|
||
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|