more

libs/langgraph/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py

@@ -7,7 +7,7 @@ from typing import Any
 
 # Configure advanced-graph runtime pools for this benchmark run.
 os.environ["LANGGRAPH_RUN_POOL_SIZE"] = "10"
-os.environ["LANGGRAPH_NODE_POOL_SIZE"] = "100"
+os.environ["LANGGRAPH_NODE_POOL_SIZE"] = "1000"
 
 from langgraph.advanced_graph import (
     AdvancedStateGraph,
@@ -20,6 +20,7 @@ from langgraph.types import Command, Send
 RUNS = 100
 MIDDLE_COUNT = 10
 SLEEP_SECONDS = 2.0
+BLOCKING_SECONDS = 0.1
 STATE_BYTES = 10 * 1024
 
 
@@ -144,6 +145,133 @@ def build_stategraph_sequential() -> Any:
     return graph.compile()
 
 
+def build_advanced_parallel_blocking() -> Any:
+    graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict)
+    done_channel = "__bench_done_channel_blocking"
+    graph.add_async_channel(done_channel, str)
+
+    async def start_node(state: dict[str, Any]) -> Command:
+        _ = state
+        sends = [Send(f"middle_blocking_{i}", None) for i in range(MIDDLE_COUNT)]
+        sends.append(Send("end_node_blocking", None))
+        return Command(goto=sends)
+
+    async def end_node_blocking(ctx: Any, state: dict[str, Any]) -> dict[str, Any]:
+        await ctx.wait_for(channel_condition(done_channel, n=MIDDLE_COUNT))
+        out = dict(state)
+        out["done"] = True
+        return out
+
+    graph.add_entry_node(start_node)
+    for i in range(MIDDLE_COUNT):
+        async def middle_blocking(
+            ctx: Any, state: dict[str, Any], idx: int = i
+        ) -> None:
+            _ = idx
+            _ = state
+            time.sleep(BLOCKING_SECONDS)
+            ctx.publish_to_channel(done_channel, "done")
+
+        graph.add_node(f"middle_blocking_{i}", middle_blocking)
+    graph.add_finish_node(end_node_blocking)
+    return graph.compile()
+
+
+def build_advanced_sequential_blocking() -> Any:
+    graph: AdvancedStateGraph[dict[str, Any]] = AdvancedStateGraph(dict)
+
+    async def start_node(state: dict[str, Any]) -> Command:
+        _ = state
+        return Command(goto=Send("middle_blocking_seq_0", None))
+
+    async def end_node_blocking_seq(state: dict[str, Any]) -> dict[str, Any]:
+        out = dict(state)
+        out["done"] = True
+        return out
+
+    graph.add_entry_node(start_node)
+
+    def make_middle(target: str):
+        async def middle_blocking_seq(state: dict[str, Any]) -> Command:
+            _ = state
+            time.sleep(BLOCKING_SECONDS)
+            return Command(goto=Send(target, None))
+
+        return middle_blocking_seq
+
+    for i in range(MIDDLE_COUNT):
+        next_name = (
+            "end_node_blocking_seq"
+            if i == MIDDLE_COUNT - 1
+            else f"middle_blocking_seq_{i+1}"
+        )
+        graph.add_node(f"middle_blocking_seq_{i}", make_middle(next_name))
+    graph.add_finish_node(end_node_blocking_seq)
+    return graph.compile()
+
+
+def build_stategraph_parallel_blocking() -> Any:
+    graph = StateGraph(dict)
+
+    async def start_node(state: dict[str, Any]) -> None:
+        _ = state
+
+    async def end_node(state: dict[str, Any]) -> dict[str, Any]:
+        out = dict(state)
+        out["done"] = True
+        return out
+
+    graph.add_node("start_node", start_node)
+    for i in range(MIDDLE_COUNT):
+        async def middle_blocking(state: dict[str, Any], idx: int = i) -> None:
+            _ = idx
+            _ = state
+            time.sleep(BLOCKING_SECONDS)
+
+        graph.add_node(f"middle_blocking_{i}", middle_blocking)
+    graph.add_node("end_node", end_node)
+
+    graph.add_edge(START, "start_node")
+    for i in range(MIDDLE_COUNT):
+        graph.add_edge("start_node", f"middle_blocking_{i}")
+        graph.add_edge(f"middle_blocking_{i}", "end_node")
+    graph.add_edge("end_node", END)
+    return graph.compile()
+
+
+def build_stategraph_sequential_blocking() -> Any:
+    graph = StateGraph(dict)
+
+    async def start_node(state: dict[str, Any]) -> None:
+        _ = state
+
+    async def end_node(state: dict[str, Any]) -> dict[str, Any]:
+        out = dict(state)
+        out["done"] = True
+        return out
+
+    graph.add_node("start_node", start_node)
+    for i in range(MIDDLE_COUNT):
+        async def middle_blocking_seq(state: dict[str, Any], idx: int = i) -> None:
+            _ = idx
+            _ = state
+            time.sleep(BLOCKING_SECONDS)
+
+        graph.add_node(f"middle_blocking_seq_{i}", middle_blocking_seq)
+    graph.add_node("end_node", end_node)
+
+    graph.add_edge(START, "start_node")
+    graph.add_edge("start_node", "middle_blocking_seq_0")
+    for i in range(MIDDLE_COUNT - 1):
+        graph.add_edge(
+            f"middle_blocking_seq_{i}",
+            f"middle_blocking_seq_{i+1}",
+        )
+    graph.add_edge(f"middle_blocking_seq_{MIDDLE_COUNT - 1}", "end_node")
+    graph.add_edge("end_node", END)
+    return graph.compile()
+
+
 async def run_benchmark(name: str, compiled: Any) -> float:
     started = time.perf_counter()
     tasks = [asyncio.create_task(compiled.ainvoke(make_initial_state())) for _ in range(RUNS)]
@@ -160,10 +288,14 @@ async def main() -> None:
         ("advanced-graph-sequential", build_advanced_sequential()),
         ("state-graph-parallel", build_stategraph_parallel()),
         ("state-graph-sequential", build_stategraph_sequential()),
+        ("advanced-graph-parallel-blocking", build_advanced_parallel_blocking()),
+        ("advanced-graph-sequential-blocking", build_advanced_sequential_blocking()),
+        ("state-graph-parallel-blocking", build_stategraph_parallel_blocking()),
+        ("state-graph-sequential-blocking", build_stategraph_sequential_blocking()),
     ]
     print(
         f"runs={RUNS}, middle_nodes={MIDDLE_COUNT}, sleep={SLEEP_SECONDS}s, "
-        f"state_bytes={STATE_BYTES}"
+        f"blocking_sleep={BLOCKING_SECONDS}s, state_bytes={STATE_BYTES}"
     )
     for name, compiled in suites:
         elapsed = await run_benchmark(name, compiled)