diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 38db8ff67..eff62e0fb 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -1,3 +1,4 @@ +import functools import sys from collections import defaultdict, deque from functools import partial @@ -46,7 +47,6 @@ from langgraph.constants import ( EMPTY_SEQ, ERROR, INTERRUPT, - MISSING, NO_WRITES, NS_END, NS_SEP, @@ -324,7 +324,7 @@ def apply_writes( @overload def prepare_next_tasks( checkpoint: Checkpoint, - pending_writes: Sequence[PendingWrite], + pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -341,7 +341,7 @@ def prepare_next_tasks( @overload def prepare_next_tasks( checkpoint: Checkpoint, - pending_writes: Sequence[PendingWrite], + pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -357,7 +357,7 @@ def prepare_next_tasks( def prepare_next_tasks( checkpoint: Checkpoint, - pending_writes: Sequence[PendingWrite], + pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -418,7 +418,7 @@ def prepare_single_task( task_id_checksum: Optional[str], *, checkpoint: Checkpoint, - pending_writes: Sequence[PendingWrite], + pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -761,9 +761,13 @@ def prepare_single_task( def _scratchpad( - pending_writes: Sequence[PendingWrite], + pending_writes: list[PendingWrite], task_id: str, ) -> PregelScratchpad: + null_resume_write = next( + (w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None + ) + return PregelScratchpad( # call call_counter=0, @@ -772,10 +776,10 @@ def _scratchpad( resume=next( (w[2] for w in pending_writes if w[0] == task_id and w[1] == RESUME), [] ), - null_resume=next( - (w[2] for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), - MISSING, - ), + null_resume=null_resume_write[2] if null_resume_write is not None else None, + _consume_null_resume=functools.partial(pending_writes.remove, null_resume_write) + if null_resume_write is not None + else lambda: None, # subgraph subgraph_counter=0, ) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 363ff375b..0f064f518 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -89,7 +89,7 @@ def map_command( raise TypeError( f"In Command.goto, expected Send/str, got {type(send).__name__}" ) - if cmd.resume: + if cmd.resume is not None: if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume): for tid, resume in cmd.resume.items(): existing: list[Any] = next( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index b329aa3d5..0559b02ea 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -54,7 +54,6 @@ from langgraph.constants import ( ERROR, INPUT, INTERRUPT, - MISSING, NS_SEP, NULL_TASK_ID, PUSH, @@ -229,20 +228,22 @@ class PregelLoop(LoopProtocol): if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD) - if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and scratchpad is not None: - if scratchpad["subgraph_counter"]: + if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and isinstance( + scratchpad, PregelScratchpad + ): + if scratchpad.subgraph_counter: self.config = patch_configurable( self.config, { CONFIG_KEY_CHECKPOINT_NS: NS_SEP.join( ( config[CONF][CONFIG_KEY_CHECKPOINT_NS], - str(scratchpad["subgraph_counter"]), + str(scratchpad.subgraph_counter), ) ) }, ) - scratchpad["subgraph_counter"] += 1 + scratchpad.subgraph_counter += 1 if not self.is_nested and config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): self.config = patch_configurable( self.config, @@ -563,9 +564,14 @@ class PregelLoop(LoopProtocol): ) # take resume value from parent - if scratchpad := configurable.get(CONFIG_KEY_SCRATCHPAD): - if scratchpad["null_resume"] is not MISSING: - self.put_writes(NULL_TASK_ID, [(RESUME, scratchpad["null_resume"])]) + if scratchpad := cast( + Optional[PregelScratchpad], configurable.get(CONFIG_KEY_SCRATCHPAD) + ): + if ( + isinstance(scratchpad, PregelScratchpad) + and scratchpad.null_resume is not None + ): + self.put_writes(NULL_TASK_ID, [(RESUME, scratchpad.null_resume)]) # map command to writes if isinstance(self.input, Command): if self.input.resume is not None and not self.checkpointer: diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 3c899c985..47ab5dfe6 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -39,7 +39,7 @@ from langgraph.errors import GraphBubbleUp, GraphInterrupt from langgraph.pregel.algo import Call from langgraph.pregel.executor import Submit from langgraph.pregel.retry import arun_with_retry, run_with_retry -from langgraph.types import PregelExecutableTask, RetryPolicy +from langgraph.types import PregelExecutableTask, PregelScratchpad, RetryPolicy from langgraph.utils.future import chain_future F = TypeVar("F", concurrent.futures.Future, asyncio.Future) @@ -135,8 +135,7 @@ class PregelRunner: return task.config[CONF][CONFIG_KEY_SEND](writes) # schedule PUSH tasks, collect futures - scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] - scratchpad.setdefault("call_counter", 0) + scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] rtn: dict[int, Optional[concurrent.futures.Future]] = {} for idx, w in enumerate(writes): # bail if not a PUSH write @@ -144,8 +143,8 @@ class PregelRunner: continue # schedule the next task, if the callback returns one wcall = calls[idx] if calls else None - cnt = scratchpad["call_counter"] - scratchpad["call_counter"] += 1 + cnt = scratchpad.call_counter + scratchpad.call_counter += 1 if next_task := self.schedule_task(task, cnt, wcall): if fut := next( ( @@ -324,8 +323,7 @@ class PregelRunner: return task.config[CONF][CONFIG_KEY_SEND](writes) # schedule PUSH tasks, collect futures - scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] - scratchpad.setdefault("call_counter", 0) + scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] rtn: dict[int, Optional[asyncio.Future]] = {} for idx, w in enumerate(writes): # bail if not a PUSH write @@ -333,8 +331,8 @@ class PregelRunner: continue # schedule the next task, if the callback returns one wcall = calls[idx] if calls is not None else None - cnt = scratchpad["call_counter"] - scratchpad["call_counter"] += 1 + 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 diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 09f777d01..806335f7e 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -19,7 +19,7 @@ from typing import ( ) from langchain_core.runnables import Runnable, RunnableConfig -from typing_extensions import Self, TypedDict +from typing_extensions import Self from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata @@ -339,16 +339,26 @@ class LoopProtocol: self.stop = stop -class PregelScratchpad(TypedDict): +@dataclasses.dataclass(**{**_DC_KWARGS, "frozen": False}) +class PregelScratchpad: # call call_counter: int # interrupt interrupt_counter: int resume: list[Any] - null_resume: Any + null_resume: Optional[Any] + _consume_null_resume: Callable[[], None] # subgraph subgraph_counter: int + def consume_null_resume(self) -> Any: + if self.null_resume is not None: + value = self.null_resume + self._consume_null_resume() + self.null_resume = None + return value + raise ValueError("No null resume to consume") + def interrupt(value: Any) -> Any: """Interrupt the graph with a resumable exception from within a node. @@ -449,7 +459,6 @@ def interrupt(value: Any) -> Any: CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_SEND, - MISSING, NS_SEP, RESUME, ) @@ -459,19 +468,19 @@ def interrupt(value: Any) -> Any: conf = get_config()["configurable"] # track interrupt index scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD] - scratchpad["interrupt_counter"] += 1 - idx = scratchpad["interrupt_counter"] + scratchpad.interrupt_counter += 1 + idx = scratchpad.interrupt_counter # find previous resume values - if scratchpad["resume"]: - if idx < len(scratchpad["resume"]): - return scratchpad["resume"][idx] + if scratchpad.resume: + if idx < len(scratchpad.resume): + return scratchpad.resume[idx] # find current resume value - if scratchpad["null_resume"] is not MISSING: - assert len(scratchpad["resume"]) == idx, (scratchpad["resume"], idx) - v = scratchpad["null_resume"] - scratchpad["null_resume"] = MISSING - scratchpad["resume"].append(v) - conf[CONFIG_KEY_SEND]([(RESUME, scratchpad["resume"])]) + if scratchpad.null_resume is not None: + assert len(scratchpad.resume) == idx, (scratchpad.resume, idx) + print("consume null resume", scratchpad.null_resume) + v = scratchpad.consume_null_resume() + scratchpad.resume.append(v) + conf[CONFIG_KEY_SEND]([(RESUME, scratchpad.resume)]) return v # no resume value found raise GraphInterrupt( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 319ceed1c..b48ce9d8d 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5082,11 +5082,27 @@ def test_interrupt_task_functional( config = {"configurable": {"thread_id": "1"}} # First run, interrupted at bar - graph.invoke({"a": ""}, config) + assert not graph.invoke({"a": ""}, config) # Resume with an answer res = graph.invoke(Command(resume="bar"), config) assert res == {"a": "foobar"} + # Test that we can interrupt the same task multiple times + config = {"configurable": {"thread_id": "2"}} + + @entrypoint(checkpointer=checkpointer) + def graph(inputs: dict) -> dict: + foo_result = foo(inputs).result() + bar_result = bar(foo_result).result() + baz_result = bar(bar_result).result() + return baz_result + + # First run, interrupted at bar + assert not graph.invoke({"a": ""}, config) + # Provide resumes + assert not graph.invoke(Command(resume="bar"), config) + assert graph.invoke(Command(resume="baz"), config) == {"a": "foobarbaz"} + def test_root_mixed_return() -> None: def my_node(state: list[str]): diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 35d6143dd..166f8720a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6408,6 +6408,22 @@ async def test_interrupt_task_functional(checkpointer_name: str) -> None: res = await graph.ainvoke(Command(resume="bar"), config) assert res == {"a": "foobar"} + # Test that we can interrupt the same task multiple times + config = {"configurable": {"thread_id": "2"}} + + @entrypoint(checkpointer=checkpointer) + async def graph(inputs: dict) -> dict: + foo_result = await foo(inputs) + bar_result = await bar(foo_result) + baz_result = await bar(bar_result) + return baz_result + + # First run, interrupted at bar + assert not await graph.ainvoke({"a": ""}, config) + # Provide resumes + assert not await graph.ainvoke(Command(resume="bar"), config) + assert await graph.ainvoke(Command(resume="baz"), config) == {"a": "foobarbaz"} + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_command_with_static_breakpoints(checkpointer_name: str) -> None: diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 9ed72cd02..3e4499266 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -13,10 +13,8 @@ from typing_extensions import Self import langgraph.scheduler.kafka.serde as serde from langgraph.constants import ( - CONF, CONFIG_KEY_DEDUPE_TASKS, CONFIG_KEY_ENSURE_LATEST, - CONFIG_KEY_SCRATCHPAD, INTERRUPT, SCHEDULED, ) @@ -178,8 +176,6 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): CONFIG_KEY_ENSURE_LATEST: True, }, ) - if CONFIG_KEY_SCRATCHPAD in config[CONF]: - config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0 # send messages to executor futures = await asyncio.gather( *( @@ -366,8 +362,6 @@ class KafkaOrchestrator(AbstractContextManager): CONFIG_KEY_ENSURE_LATEST: True, }, ) - if CONFIG_KEY_SCRATCHPAD in config[CONF]: - config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0 # send messages to executor futures = [ self.producer.send(