diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 488f12661..246510fb4 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/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 diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index c24134517..29a2d3ccb 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/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) ) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 169311035..82368be78 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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 diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index a7c99900d..0055e89b3 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/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"],