From 1e9a372dd7d4be0f638790c80840b4b0f6ab01bf Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 10:11:31 -0800 Subject: [PATCH] Fix multiple interrupt/task test --- libs/langgraph/langgraph/pregel/runner.py | 59 ++++++++++------------- libs/langgraph/langgraph/types.py | 3 ++ libs/langgraph/tests/test_pregel.py | 1 - libs/langgraph/tests/test_pregel_async.py | 38 +++++++++++++++ 4 files changed, 66 insertions(+), 35 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 87e73850c..790d22a06 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -1,6 +1,5 @@ import asyncio import concurrent.futures -import threading import time from functools import partial from typing import ( @@ -22,6 +21,7 @@ from langchain_core.callbacks import Callbacks from langgraph.constants import ( CONF, CONFIG_KEY_CALL, + CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_SEND, ERROR, INTERRUPT, @@ -71,8 +71,6 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: - locks: dict[str, threading.Lock] = {} - def writer( task: PregelExecutableTask, writes: Sequence[tuple[str, Any]], @@ -82,24 +80,19 @@ class PregelRunner: if all(w[0] != PUSH for w in writes): return task.config[CONF][CONFIG_KEY_SEND](writes) - if task.id not in locks: - locks[task.id] = threading.Lock() - with locks[task.id]: - prev_length = len(task.writes) - # delegate to the underlying writer - task.config[CONF][CONFIG_KEY_SEND](writes) - # confirm no other concurrent writes were added - assert len(task.writes) == prev_length + len(writes) # schedule PUSH tasks, collect futures + scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] + scratchpad.setdefault("call_counter", 0) rtn: dict[int, Optional[concurrent.futures.Future]] = {} - for idx, w in enumerate(writes, start=prev_length): + for idx, w in enumerate(writes): # bail if not a PUSH write if w[0] != PUSH: continue # schedule the next task, if the callback returns one - if next_task := self.schedule_task( - task, idx, calls[idx - prev_length] if calls else None - ): + wcall = calls[idx] if calls else None + cnt = scratchpad["call_counter"] + scratchpad["call_counter"] += 1 + if next_task := self.schedule_task(task, cnt, wcall): if fut := next( ( f @@ -110,7 +103,7 @@ class PregelRunner: ): # if the parent task was retried, # the next task might already be running - rtn[idx - prev_length] = fut + rtn[idx] = fut elif next_task.writes: # if it already ran, return the result fut = concurrent.futures.Future() @@ -130,7 +123,7 @@ class PregelRunner: ) else: fut.set_result(None) - rtn[idx - prev_length] = fut + rtn[idx] = fut else: # schedule the next task fut = self.submit( @@ -148,7 +141,7 @@ class PregelRunner: ) fut.add_done_callback(partial(self.commit, next_task)) futures[fut] = next_task - rtn[idx - prev_length] = fut + rtn[idx] = fut return [rtn.get(i) for i in range(len(writes))] def call( @@ -196,6 +189,8 @@ class PregelRunner: raise if not futures: # maybe `t` schuduled another task return + else: + tasks = () # don't reschedule this task # add waiter task if requested if get_waiter is not None: futures[get_waiter()] = None @@ -262,8 +257,6 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, ) -> AsyncIterator[None]: - locks: dict[str, threading.Lock] = {} - def writer( task: PregelExecutableTask, writes: Sequence[tuple[str, Any]], @@ -273,23 +266,19 @@ class PregelRunner: if all(w[0] != PUSH for w in writes): return task.config[CONF][CONFIG_KEY_SEND](writes) - if task.id not in locks: - locks[task.id] = threading.Lock() - with locks[task.id]: - prev_length = len(task.writes) - # delegate to the underlying writer - task.config[CONF][CONFIG_KEY_SEND](writes) - # confirm no other concurrent writes were added - assert len(task.writes) == prev_length + len(writes) # schedule PUSH tasks, collect futures + scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] + scratchpad.setdefault("call_counter", 0) rtn: dict[int, Optional[asyncio.Future]] = {} - for idx, w in enumerate(writes, start=prev_length): + for idx, w in enumerate(writes): # bail if not a PUSH write if w[0] != PUSH: continue # schedule the next task, if the callback returns one - wcall = calls[idx - prev_length] if calls is not None else None - if next_task := self.schedule_task(task, idx, wcall): + wcall = calls[idx] if calls is not None else None + cnt = scratchpad["call_counter"] + scratchpad["call_counter"] += 1 + if next_task := self.schedule_task(task, cnt, wcall): # if the parent task was retried, # the next task might already be running if fut := next( @@ -302,7 +291,7 @@ class PregelRunner: ): # if the parent task was retried, # the next task might already be running - rtn[idx - prev_length] = fut + rtn[idx] = fut elif next_task.writes: # if it already ran, return the result fut = asyncio.Future() @@ -322,7 +311,7 @@ class PregelRunner: ) else: fut.set_result(None) - rtn[idx - prev_length] = fut + rtn[idx] = fut else: # schedule the next task fut = cast( @@ -346,7 +335,7 @@ class PregelRunner: ) fut.add_done_callback(partial(self.commit, next_task)) futures[fut] = next_task - rtn[idx - prev_length] = fut + rtn[idx] = fut return [rtn.get(i) for i in range(len(writes))] def call( @@ -401,6 +390,8 @@ class PregelRunner: raise if not futures: # maybe `t` schuduled another task return + else: + tasks = () # don't reschedule this task # add waiter task if requested if get_waiter is not None: futures[get_waiter()] = None diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 066d7d106..0b9fb9b1b 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -342,9 +342,12 @@ class LoopProtocol: class PregelScratchpad(TypedDict, total=False): + # interrupt interrupt_counter: int used_null_resume: bool resume: list[Any] + # call + call_counter: int def interrupt(value: Any) -> Any: diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index c5fb33d71..12b7c6863 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5309,7 +5309,6 @@ def test_multiple_interrupts_imperative() -> None: configurable = {"configurable": {"thread_id": uuid.uuid4()}} graph.invoke({}, configurable) - # Currently fails when double accepts a variable graph.invoke(Command(resume="a"), configurable) graph.invoke(Command(resume="b"), configurable) result = graph.invoke(Command(resume="c"), configurable) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index e4ab4b0c2..d826daf99 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6712,3 +6712,41 @@ async def test_falsy_return_from_task() -> None: configurable = {"configurable": {"thread_id": uuid.uuid4()}} await graph.ainvoke({"a": 5}, configurable) await graph.ainvoke(Command(resume="123"), configurable) + + +async def test_multiple_interrupts_imperative() -> None: + """Test multiple interrupts with an imperative API.""" + from langgraph.checkpoint.memory import MemorySaver + from langgraph.func import entrypoint, task + + checkpointer = MemorySaver() + counter = 0 + + @task + async def double(x: int) -> int: + """Increment the counter.""" + nonlocal counter + counter += 1 + return 2 * x + + @entrypoint(checkpointer=checkpointer) + async def graph(state: dict) -> dict: + """React tool.""" + + values = [] + + for idx in [1, 2, 3]: + values.extend([await double(idx), interrupt({"a": "boo"})]) + + return {"values": values} + + configurable = {"configurable": {"thread_id": uuid.uuid4()}} + await graph.ainvoke({}, configurable) + await graph.ainvoke(Command(resume="a"), configurable) + await graph.ainvoke(Command(resume="b"), configurable) + result = await graph.ainvoke(Command(resume="c"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 4, "b", 6, "c"], + } + assert counter == 3