Merge pull request #81 from langchain-ai/nc/resume-or-restart

Only execute tasks leftover from previous checkpoint when resuming with no new input

langgraph/pregel/__init__.py

@@ -252,6 +252,7 @@ class Pregel(
         *,
         input_keys: Optional[Union[str, Sequence[str]]] = None,
         output_keys: Optional[Union[str, Sequence[str]]] = None,
+        interrupt: Optional[Sequence[str]] = None,
     ) -> Iterator[Union[dict[str, Any], Any]]:
         try:
             if config["recursion_limit"] < 1:
@@ -267,6 +268,7 @@ class Pregel(
                 input_keys = self.input
             else:
                 validate_keys(input_keys, self.channels)
+            interrupt = interrupt or self.interrupt
             # copy nodes to ignore mutations during execution
             processes = {**self.nodes}
             # get checkpoint from saver, or create an empty one
@@ -277,13 +279,19 @@ class Pregel(
                 self.channels, checkpoint
             ) as channels, get_executor_for_config(config) as executor:
                 # map inputs to channel updates
-                _apply_writes(
-                    checkpoint,
-                    channels,
-                    deque(w for c in input for w in map_input(input_keys, c)),
-                    config,
-                    0,
-                )
+                if input_writes := deque(
+                    w for c in input for w in map_input(input_keys, c)
+                ):
+                    # discard any unfinished tasks from previous checkpoint
+                    _prepare_next_tasks(checkpoint, processes, channels)
+                    # apply input writes
+                    _apply_writes(
+                        checkpoint,
+                        channels,
+                        input_writes,
+                        config,
+                        0,
+                    )
 
                 read = partial(_read_channel, channels)
 
@@ -371,9 +379,7 @@ class Pregel(
                         self.checkpointer.put(config, checkpoint)
 
                     # interrupt if any channel written to is in interrupt list
-                    if any(
-                        chan for chan, _ in pending_writes if chan in self.interrupt
-                    ):
+                    if any(chan for chan, _ in pending_writes if chan in interrupt):
                         break
 
                 # save end of run checkpoint
@@ -400,6 +406,7 @@ class Pregel(
         *,
         input_keys: Optional[Union[str, Sequence[str]]] = None,
         output_keys: Optional[Union[str, Sequence[str]]] = None,
+        interrupt: Optional[Sequence[str]] = None,
     ) -> AsyncIterator[Union[dict[str, Any], Any]]:
         try:
             if config["recursion_limit"] < 1:
@@ -424,6 +431,7 @@ class Pregel(
                 input_keys = self.input
             else:
                 validate_keys(input_keys, self.channels)
+            interrupt = interrupt or self.interrupt
             # copy nodes to ignore mutations during execution
             processes = {**self.nodes}
             # get checkpoint from saver, or create an empty one
@@ -434,13 +442,19 @@ class Pregel(
             # create channels from checkpoint
             async with AsyncChannelsManager(self.channels, checkpoint) as channels:
                 # map inputs to channel updates
-                _apply_writes(
-                    checkpoint,
-                    channels,
-                    deque([w async for c in input for w in map_input(input_keys, c)]),
-                    config,
-                    0,
-                )
+                if input_writes := deque(
+                    [w async for c in input for w in map_input(input_keys, c)]
+                ):
+                    # discard any unfinished tasks from previous checkpoint
+                    _prepare_next_tasks(checkpoint, processes, channels)
+                    # apply input writes
+                    _apply_writes(
+                        checkpoint,
+                        channels,
+                        input_writes,
+                        config,
+                        0,
+                    )
 
                 read = partial(_read_channel, channels)
 
@@ -535,9 +549,7 @@ class Pregel(
                         await self.checkpointer.aput(config, checkpoint)
 
                     # interrupt if any channel written to is in interrupt list
-                    if any(
-                        chan for chan, _ in pending_writes if chan in self.interrupt
-                    ):
+                    if any(chan for chan, _ in pending_writes if chan in interrupt):
                         break
 
                 # save end of run checkpoint

tests/test_pregel.py

@@ -245,6 +245,40 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
     assert step == 3
 
 
+def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
+
+    memory = MemorySaver()
+    app = Pregel(
+        nodes={"one": one, "two": two}, checkpointer=memory, interrupt=["inbox"]
+    )
+
+    # start execution, stop at inbox
+    assert app.invoke(2, {"configurable": {"thread_id": 1}}) is None
+
+    # inbox == 3
+    checkpoint = memory.get({"configurable": {"thread_id": 1}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"]["inbox"] == 3
+
+    # resume execution, finish
+    assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 4
+
+    # start execution again, stop at inbox
+    assert app.invoke(20, {"configurable": {"thread_id": 1}}) is None
+
+    # inbox == 21
+    checkpoint = memory.get({"configurable": {"thread_id": 1}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"]["inbox"] == 21
+
+    # send a new value in, interrupting the previous execution
+    assert app.invoke(3, {"configurable": {"thread_id": 1}}) is None
+    assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 5
+
+
 def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")

tests/test_pregel_async.py

@@ -252,6 +252,40 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
     assert step == 3
 
 
+async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
+    two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
+
+    memory = MemorySaver()
+    app = Pregel(
+        nodes={"one": one, "two": two}, checkpointer=memory, interrupt=["inbox"]
+    )
+
+    # start execution, stop at inbox
+    assert await app.ainvoke(2, {"configurable": {"thread_id": 1}}) is None
+
+    # inbox == 3
+    checkpoint = await memory.aget({"configurable": {"thread_id": 1}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"]["inbox"] == 3
+
+    # resume execution, finish
+    assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 4
+
+    # start execution again, stop at inbox
+    assert await app.ainvoke(20, {"configurable": {"thread_id": 1}}) is None
+
+    # inbox == 21
+    checkpoint = await memory.aget({"configurable": {"thread_id": 1}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"]["inbox"] == 21
+
+    # send a new value in, interrupting the previous execution
+    assert await app.ainvoke(3, {"configurable": {"thread_id": 1}}) is None
+    assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 5
+
+
 async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")