From 6036b4a7d4c4b0b9ccf0e7ced9377b5e0b47ddae Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 30 Aug 2024 10:11:32 -0700 Subject: [PATCH] Remove the need to fork a thread to replay a past state --- libs/langgraph/langgraph/pregel/algo.py | 12 ++--- libs/langgraph/langgraph/pregel/loop.py | 13 ++---- libs/langgraph/langgraph/pregel/types.py | 3 -- libs/langgraph/tests/test_pregel.py | 51 ++++----------------- libs/langgraph/tests/test_pregel_async.py | 55 ++++------------------- 5 files changed, 24 insertions(+), 110 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index e4e4b45eb..f3d82c835 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -50,7 +50,7 @@ from langgraph.pregel.io import read_channel, read_channels from langgraph.pregel.log import logger from langgraph.pregel.manager import ChannelsManager from langgraph.pregel.read import PregelNode -from langgraph.pregel.types import EXACT_MATCH, All, PregelExecutableTask, PregelTask +from langgraph.pregel.types import All, PregelExecutableTask, PregelTask class WritesProtocol(Protocol): @@ -357,10 +357,7 @@ def prepare_next_tasks( ), triggers, proc.retry_policy, - None - if task_checkpoint_ns - in configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}) - else EXACT_MATCH, + None, task_id, ) ) @@ -466,10 +463,7 @@ def prepare_next_tasks( ), triggers, proc.retry_policy, - None - if task_checkpoint_ns - in configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}) - else EXACT_MATCH, + None, task_id, ) ) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 13fe4c58b..d791a5539 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -122,6 +122,7 @@ class PregelLoop: output_keys: Union[str, Sequence[str]] stream_keys: Union[str, Sequence[str]] is_nested: bool + skip_done_tasks: bool checkpointer_get_next_version: Callable[[Optional[V]], V] checkpointer_put_writes: Optional[ @@ -180,6 +181,7 @@ class PregelLoop: self.output_keys = output_keys self.stream_keys = stream_keys self.is_nested = CONFIG_KEY_TASK_ID in self.config.get("configurable", {}) + self.skip_done_tasks = "checkpoint_id" not in config["configurable"] if CONFIG_KEY_STREAM in config["configurable"]: self.stream = DuplexStream( self.stream, config["configurable"][CONFIG_KEY_STREAM] @@ -340,18 +342,11 @@ class PregelLoop: return False # if there are pending writes from a previous loop, apply them - if self.checkpoint_pending_writes: + if self.skip_done_tasks and self.checkpoint_pending_writes: for tid, k, v in self.checkpoint_pending_writes: if k in (ERROR, INTERRUPT): continue - if task := next( - ( - t - for t in self.tasks - if t.id == tid and t.cache_policy is not None - ), - None, - ): + if task := next((t for t in self.tasks if t.id == tid), None): task.writes.append((k, v)) # print output for any tasks we applied previous writes to for task in self.tasks: diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 8a5a37b37..6da463286 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -63,9 +63,6 @@ class CachePolicy(NamedTuple): pass -EXACT_MATCH = CachePolicy() - - class PregelTask(NamedTuple): id: str name: str diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0ee645dbe..704824fe5 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -798,27 +798,13 @@ def test_invoke_two_processes_in_out_interrupt( ), ] - # re-running from any previous checkpoint w/out forking should do nothing + # re-running from any previous checkpoint should re-run nodes assert [c for c in app.stream(None, history[0].config, stream_mode="updates")] == [] assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == [ - {"two": {"output": 5}, "__metadata__": {"cached": True}}, + {"two": {"output": 5}}, ] assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == [ - {"one": {"inbox": 4}, "__metadata__": {"cached": True}}, - ] - - # forking and re-running from any prev checkpoint should re-run nodes - fork_config = app.update_state(history[0].config, None) - assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [] - - fork_config = app.update_state(history[1].config, None) - assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [ - {"two": {"output": 5}} - ] - - fork_config = app.update_state(history[2].config, None) - assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [ - {"one": {"inbox": 4}} + {"one": {"inbox": 4}}, ] @@ -987,25 +973,18 @@ def test_fork_always_re_runs_nodes( ), ] - # forking from any previous checkpoint w/out forking should do nothing + # forking from any previous checkpoint should re-run nodes assert [ c for c in graph.stream(None, history[0].config, stream_mode="updates") ] == [] assert [ c for c in graph.stream(None, history[1].config, stream_mode="updates") - ] == [{"add_one": 1, "__metadata__": {"cached": True}}] - - # forking and re-running from any prev checkpoint should re-run nodes - fork_config = graph.update_state(history[0].config, None) - assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [] - - fork_config = graph.update_state(history[1].config, None) - assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [ - {"add_one": 1} + ] == [ + {"add_one": 1}, ] - - fork_config = graph.update_state(history[2].config, None) - assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [ + assert [ + c for c in graph.stream(None, history[2].config, stream_mode="updates") + ] == [ {"add_one": 1}, {"add_one": 1}, ] @@ -9813,18 +9792,6 @@ def test_doubly_nested_graph_state( assert [ c for c in app.stream(None, grandchild_history[2].config, subgraphs=True) ] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - { - "grandchild_1": {"my_key": "hi my value here"}, - "__metadata__": {"cached": True}, - }, - ) - ] - - # fork and replay - fork = app.update_state(grandchild_history[2].config, None) - assert [c for c in app.stream(None, fork, subgraphs=True)] == [ ( (AnyStr("child:"), AnyStr("child_1:")), {"grandchild_1": {"my_key": "hi my value here"}}, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 8083ddd77..e5cda8e11 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1020,35 +1020,19 @@ async def test_invoke_two_processes_in_out_interrupt( ), ] - # forking from any previous checkpoint w/out forking should do nothing + # forking from any previous checkpoint should re-run nodes assert [ c async for c in app.astream(None, history[0].config, stream_mode="updates") ] == [] assert [ c async for c in app.astream(None, history[1].config, stream_mode="updates") ] == [ - {"two": {"output": 5}, "__metadata__": {"cached": True}}, + {"two": {"output": 5}}, ] assert [ c async for c in app.astream(None, history[2].config, stream_mode="updates") ] == [ - {"one": {"inbox": 4}, "__metadata__": {"cached": True}}, - ] - - # forking and re-running from any prev checkpoint should re-run nodes - fork_config = await app.aupdate_state(history[0].config, None) - assert [ - c async for c in app.astream(None, fork_config, stream_mode="updates") - ] == [] - - fork_config = await app.aupdate_state(history[1].config, None) - assert [c async for c in app.astream(None, fork_config, stream_mode="updates")] == [ - {"two": {"output": 5}} - ] - - fork_config = await app.aupdate_state(history[2].config, None) - assert [c async for c in app.astream(None, fork_config, stream_mode="updates")] == [ - {"one": {"inbox": 4}} + {"one": {"inbox": 4}}, ] @@ -1219,28 +1203,17 @@ async def test_fork_always_re_runs_nodes( ), ] - # forking from any previous checkpoint w/out forking should do nothing + # forking from any previous checkpoint should re-run nodes assert [ c async for c in graph.astream(None, history[0].config, stream_mode="updates") ] == [] assert [ c async for c in graph.astream(None, history[1].config, stream_mode="updates") - ] == [{"add_one": 1, "__metadata__": {"cached": True}}] - - # forking and re-running from any prev checkpoint should re-run nodes - fork_config = await graph.aupdate_state(history[0].config, None) + ] == [ + {"add_one": 1}, + ] assert [ - c async for c in graph.astream(None, fork_config, stream_mode="updates") - ] == [] - - fork_config = await graph.aupdate_state(history[1].config, None) - assert [ - c async for c in graph.astream(None, fork_config, stream_mode="updates") - ] == [{"add_one": 1}] - - fork_config = await graph.aupdate_state(history[2].config, None) - assert [ - c async for c in graph.astream(None, fork_config, stream_mode="updates") + c async for c in graph.astream(None, history[2].config, stream_mode="updates") ] == [ {"add_one": 1}, {"add_one": 1}, @@ -8340,18 +8313,6 @@ async def test_doubly_nested_graph_state( assert [ c async for c in app.astream(None, grandchild_history[2].config, subgraphs=True) ] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - { - "grandchild_1": {"my_key": "hi my value here"}, - "__metadata__": {"cached": True}, - }, - ) - ] - - # fork and replay - fork = await app.aupdate_state(grandchild_history[2].config, None) - assert [c async for c in app.astream(None, fork, subgraphs=True)] == [ ( (AnyStr("child:"), AnyStr("child_1:")), {"grandchild_1": {"my_key": "hi my value here"}},