Merge pull request #2143 from langchain-ai/nc/19oct/async-max-concurrency

lib: Add max_concurrency for async executions

libs/langgraph/langgraph/pregel/executor.py

@@ -118,11 +118,17 @@ class AsyncBackgroundExecutor(AsyncContextManager):
     - re-raises the first exception from tasks with `__reraise_on_exit__=True`
       ignoring CancelledError"""
 
-    def __init__(self) -> None:
+    def __init__(self, config: RunnableConfig) -> None:
         self.context_not_supported = sys.version_info < (3, 11)
         self.tasks: dict[asyncio.Task, tuple[bool, bool]] = {}
         self.sentinel = object()
         self.loop = asyncio.get_running_loop()
+        if max_concurrency := config.get("max_concurrency"):
+            self.semaphore: Optional[asyncio.Semaphore] = asyncio.Semaphore(
+                max_concurrency
+            )
+        else:
+            self.semaphore = None
 
     def submit(  # type: ignore[valid-type]
         self,
@@ -134,6 +140,8 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         **kwargs: P.kwargs,
     ) -> asyncio.Task[T]:
         coro = cast(Coroutine[None, None, T], fn(*args, **kwargs))
+        if self.semaphore:
+            coro = gated(self.semaphore, coro)
         if self.context_not_supported:
             task = self.loop.create_task(coro, name=__name__)
         else:
@@ -183,3 +191,9 @@ class AsyncBackgroundExecutor(AsyncContextManager):
                         raise exc
                 except asyncio.CancelledError:
                     pass
+
+
+async def gated(semaphore: asyncio.Semaphore, coro: Coroutine[None, None, T]) -> T:
+    """A coroutine that waits for a semaphore before running another coroutine."""
+    async with semaphore:
+        return await coro

libs/langgraph/langgraph/pregel/loop.py

@@ -847,7 +847,9 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             else []
         )
 
-        self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
+        self.submit = await self.stack.enter_async_context(
+            AsyncBackgroundExecutor(self.config)
+        )
         self.channels, self.managed = await self.stack.enter_async_context(
             AsyncChannelsManager(self.specs, self.checkpoint, self)
         )

libs/langgraph/tests/test_pregel_async.py

@@ -1922,6 +1922,53 @@ async def test_cond_edge_after_send() -> None:
     assert await graph.ainvoke(["0"]) == ["0", "1", "2", "2", "3"]
 
 
+async def test_max_concurrency() -> None:
+    class Node:
+        def __init__(self, name: str):
+            self.name = name
+            setattr(self, "__name__", name)
+            self.currently = 0
+            self.max_currently = 0
+
+        async def __call__(self, state):
+            self.currently += 1
+            if self.currently > self.max_currently:
+                self.max_currently = self.currently
+            await asyncio.sleep(0.1)
+            self.currently -= 1
+            return [self.name]
+
+    async def send_to_many(state):
+        return [Send("2", state)] * 100
+
+    async def route_to_three(state) -> Literal["3"]:
+        return "3"
+
+    node2 = Node("2")
+    builder = StateGraph(Annotated[list, operator.add])
+    builder.add_node(Node("1"))
+    builder.add_node(node2)
+    builder.add_node(Node("3"))
+    builder.add_edge(START, "1")
+    builder.add_conditional_edges("1", send_to_many)
+    builder.add_conditional_edges("2", route_to_three)
+    graph = builder.compile()
+
+    assert await graph.ainvoke(["0"]) == ["0", "1", *(["2"] * 100), "3"]
+    assert node2.max_currently == 100
+    assert node2.currently == 0
+    node2.max_currently = 0
+
+    assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [
+        "0",
+        "1",
+        *(["2"] * 100),
+        "3",
+    ]
+    assert node2.max_currently == 10
+    assert node2.currently == 0
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_invoke_checkpoint_three(
     mocker: MockerFixture, checkpointer_name: str

libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py

@@ -191,7 +191,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
                 step=saved.metadata["step"] + 1,
                 stop=saved.metadata["step"] + 2,
             ),
-        ) as (channels, managed), AsyncBackgroundExecutor() as submit:
+        ) as (channels, managed), AsyncBackgroundExecutor(msg["config"]) as submit:
             if task := await asyncio.to_thread(
                 prepare_single_task,
                 msg["task"]["path"],