From 42648c88ddbe06a25ceaf9f69f798b5fb31fbf18 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 19 Oct 2024 12:45:01 -0700 Subject: [PATCH 1/2] lib: Add max_concurrency for async executions --- libs/langgraph/langgraph/pregel/executor.py | 14 +++++- libs/langgraph/langgraph/pregel/loop.py | 4 +- libs/langgraph/tests/test_pregel_async.py | 47 +++++++++++++++++++ .../langgraph/scheduler/kafka/executor.py | 2 +- 4 files changed, 64 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 488f12661..2b627df2e 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -118,11 +118,15 @@ 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 = asyncio.Semaphore(max_concurrency) + else: + self.semaphore = None def submit( # type: ignore[valid-type] self, @@ -134,6 +138,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 +189,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 f25fa2300..bc8ed4a3e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -839,7 +839,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"], From 7d3f2ca3ed13863afaf1af59d8c07d4ae766a6fe Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 19 Oct 2024 15:11:18 -0700 Subject: [PATCH 2/2] Lint --- libs/langgraph/langgraph/pregel/executor.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 2b627df2e..246510fb4 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -124,7 +124,9 @@ class AsyncBackgroundExecutor(AsyncContextManager): self.sentinel = object() self.loop = asyncio.get_running_loop() if max_concurrency := config.get("max_concurrency"): - self.semaphore = asyncio.Semaphore(max_concurrency) + self.semaphore: Optional[asyncio.Semaphore] = asyncio.Semaphore( + max_concurrency + ) else: self.semaphore = None