From e1d4b5552d5c1f1c234f6c087f146096d736cf41 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 3 Apr 2025 16:51:53 -0700 Subject: [PATCH 01/13] Add checkpoint_during arg - This provides a new mode of execution where only the last checkpoint is saved - We save the last checkpoint no matter how the agent run is terminated (success, error, interrupt, etc) - This cuts down on cpu time spent on checkpointing, while not losing any resilience benefits, given individual task writes are still saved - If an error occurs and the run is retried, any tasks that completed successfully before will be skipped (as currently) - checkpoint_during=True is useful when you want to time-travel to inner steps of a run - The default value will remain the current behavior, ie. checkpoint_during=True --- libs/langgraph/bench/__main__.py | 4 +++ libs/langgraph/langgraph/pregel/__init__.py | 10 +++++++ libs/langgraph/langgraph/pregel/loop.py | 33 ++++++++++++++------- libs/langgraph/tests/test_interruption.py | 21 ++++++------- 4 files changed, 47 insertions(+), 21 deletions(-) diff --git a/libs/langgraph/bench/__main__.py b/libs/langgraph/bench/__main__.py index 033b5034b..8be3cf5c8 100644 --- a/libs/langgraph/bench/__main__.py +++ b/libs/langgraph/bench/__main__.py @@ -25,6 +25,7 @@ async def arun(graph: Pregel, input: dict): "configurable": {"thread_id": str(uuid4())}, "recursion_limit": 1000000000, }, + checkpoint_during=False, ) ] ) @@ -41,6 +42,7 @@ async def arun_first_event_latency(graph: Pregel, input: dict) -> None: "configurable": {"thread_id": str(uuid4())}, "recursion_limit": 1000000000, }, + checkpoint_during=False, ) try: @@ -60,6 +62,7 @@ def run(graph: Pregel, input: dict): "configurable": {"thread_id": str(uuid4())}, "recursion_limit": 1000000000, }, + checkpoint_during=False, ) ] ) @@ -76,6 +79,7 @@ def run_first_event_latency(graph: Pregel, input: dict) -> None: "configurable": {"thread_id": str(uuid4())}, "recursion_limit": 1000000000, }, + checkpoint_during=False, ) try: diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 545c8faf1..cbfbd12b9 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2094,6 +2094,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, + checkpoint_during: bool = True, debug: Optional[bool] = None, subgraphs: bool = False, ) -> Iterator[Union[dict[str, Any], Any]]: @@ -2115,6 +2116,7 @@ class Pregel(PregelProtocol): output_keys: The keys to stream, defaults to all non-context channels. interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. + checkpoint_during: Whether to checkpoint intermediate steps, defaults to True. If False, only the final checkpoint is saved. debug: Whether to print debug information during execution, defaults to False. subgraphs: Whether to stream subgraphs, defaults to False. @@ -2291,6 +2293,7 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, + checkpoint_during=checkpoint_during, trigger_to_nodes=self.trigger_to_nodes, migrate_checkpoint=self._migrate_checkpoint, ) as loop: @@ -2373,6 +2376,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, + checkpoint_during: bool = True, debug: Optional[bool] = None, subgraphs: bool = False, ) -> AsyncIterator[Union[dict[str, Any], Any]]: @@ -2394,6 +2398,7 @@ class Pregel(PregelProtocol): output_keys: The keys to stream, defaults to all non-context channels. interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. + checkpoint_during: Whether to checkpoint intermediate steps, defaults to True. If False, only the final checkpoint is saved. debug: Whether to print debug information during execution, defaults to False. subgraphs: Whether to stream subgraphs, defaults to False. @@ -2586,6 +2591,7 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, + checkpoint_during=checkpoint_during, trigger_to_nodes=self.trigger_to_nodes, migrate_checkpoint=self._migrate_checkpoint, ) as loop: @@ -2661,6 +2667,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, + checkpoint_during: bool = True, debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: @@ -2692,6 +2699,7 @@ class Pregel(PregelProtocol): output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, + checkpoint_during=checkpoint_during, debug=debug, **kwargs, ): @@ -2713,6 +2721,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, + checkpoint_during: bool = True, debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: @@ -2745,6 +2754,7 @@ class Pregel(PregelProtocol): output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, + checkpoint_during=checkpoint_during, debug=debug, **kwargs, ): diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 4c4ed6dde..00469fbbf 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -155,7 +155,7 @@ class PregelLoop(LoopProtocol): manager: Union[None, AsyncParentRunManager, ParentRunManager] interrupt_after: Union[All, Sequence[str]] interrupt_before: Union[All, Sequence[str]] - checkpoint_every_step: bool + checkpoint_during: bool debug: bool checkpointer_get_next_version: GetNextVersion @@ -215,7 +215,7 @@ class PregelLoop(LoopProtocol): debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, - checkpoint_every_step: bool = True, + checkpoint_during: bool = True, ) -> None: super().__init__( step=0, @@ -241,7 +241,7 @@ class PregelLoop(LoopProtocol): ) self._migrate_checkpoint = migrate_checkpoint self.trigger_to_nodes = trigger_to_nodes - self.checkpoint_every_step = checkpoint_every_step + self.checkpoint_during = checkpoint_during self.debug = debug if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) @@ -709,10 +709,13 @@ class PregelLoop(LoopProtocol): ) return updated_channels - def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: + def _put_checkpoint( + self, metadata: CheckpointMetadata, force: bool = False + ) -> None: # assign step and parents metadata["step"] = self.step metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) + self.checkpoint_metadata = metadata # debug flag if self.debug: print_step_checkpoint( @@ -724,19 +727,21 @@ class PregelLoop(LoopProtocol): else self.stream_keys ), ) + # do checkpoint? + do_checkpoint = self._checkpointer_put_after_previous is not None and ( + force or self.checkpoint_during + ) + # create new checkpoint + self.checkpoint = create_checkpoint( + self.checkpoint, self.channels if do_checkpoint else None, self.step + ) # bail if no checkpointer - if self._checkpointer_put_after_previous is not None: + if do_checkpoint: for k, v in self.config["metadata"].items(): if k in EXCLUDED_METADATA_KEYS: continue metadata.setdefault(k, v) # type: ignore - # create new checkpoint - self.checkpoint = create_checkpoint( - self.checkpoint, self.channels, self.step - ) - self.checkpoint_metadata = metadata - self.prev_checkpoint_config = ( self.checkpoint_config if CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF] @@ -789,6 +794,8 @@ class PregelLoop(LoopProtocol): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: + if not self.checkpoint_during: + self._put_checkpoint(self.checkpoint_metadata, force=True) # suppress interrupt suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested if suppress: @@ -907,6 +914,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + checkpoint_during: bool = True, ) -> None: super().__init__( input, @@ -925,6 +933,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): debug=debug, migrate_checkpoint=migrate_checkpoint, trigger_to_nodes=trigger_to_nodes, + checkpoint_during=checkpoint_during, ) self.stack = ExitStack() if checkpointer: @@ -1054,6 +1063,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + checkpoint_during: bool = True, ) -> None: super().__init__( input, @@ -1072,6 +1082,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): debug=debug, migrate_checkpoint=migrate_checkpoint, trigger_to_nodes=trigger_to_nodes, + checkpoint_during=checkpoint_during, ) self.stack = AsyncExitStack() if checkpointer: diff --git a/libs/langgraph/tests/test_interruption.py b/libs/langgraph/tests/test_interruption.py index aa543cf00..9d262bffc 100644 --- a/libs/langgraph/tests/test_interruption.py +++ b/libs/langgraph/tests/test_interruption.py @@ -1,5 +1,4 @@ import pytest -from pytest_mock import MockerFixture from typing_extensions import TypedDict from langgraph.graph import END, START, StateGraph @@ -12,9 +11,10 @@ from tests.conftest import ( pytestmark = pytest.mark.anyio +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_interruption_without_state_updates( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: """Test interruption without state updates. This test confirms that interrupting doesn't require a state key having been updated in the prev step""" @@ -40,20 +40,21 @@ def test_interruption_without_state_updates( initial_input = {"input": "hello world"} thread = {"configurable": {"thread_id": "1"}} - graph.invoke(initial_input, thread, debug=True) + graph.invoke(initial_input, thread, checkpoint_during=checkpoint_during) assert graph.get_state(thread).next == ("step_2",) - graph.invoke(None, thread, debug=True) + graph.invoke(None, thread, checkpoint_during=checkpoint_during) assert graph.get_state(thread).next == ("step_3",) - graph.invoke(None, thread, debug=True) + graph.invoke(None, thread, checkpoint_during=checkpoint_during) assert graph.get_state(thread).next == () +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interruption_without_state_updates_async( - checkpointer_name: str, mocker: MockerFixture -): + checkpointer_name: str, checkpoint_during: bool +) -> None: """Test interruption without state updates. This test confirms that interrupting doesn't require a state key having been updated in the prev step""" @@ -78,11 +79,11 @@ async def test_interruption_without_state_updates_async( initial_input = {"input": "hello world"} thread = {"configurable": {"thread_id": "1"}} - await graph.ainvoke(initial_input, thread, debug=True) + await graph.ainvoke(initial_input, thread, checkpoint_during=checkpoint_during) assert (await graph.aget_state(thread)).next == ("step_2",) - await graph.ainvoke(None, thread, debug=True) + await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during) assert (await graph.aget_state(thread)).next == ("step_3",) - await graph.ainvoke(None, thread, debug=True) + await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during) assert (await graph.aget_state(thread)).next == () From 7e08339335b46f4cb820b5691f223de0ffea2178 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 3 Apr 2025 16:55:23 -0700 Subject: [PATCH 02/13] mypy is dumb --- libs/langgraph/langgraph/pregel/loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 00469fbbf..302bb4744 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -736,7 +736,7 @@ class PregelLoop(LoopProtocol): self.checkpoint, self.channels if do_checkpoint else None, self.step ) # bail if no checkpointer - if do_checkpoint: + if do_checkpoint and self._checkpointer_put_after_previous is not None: for k, v in self.config["metadata"].items(): if k in EXCLUDED_METADATA_KEYS: continue From 0a1dd7a01aba32443338f1d455f8bb3505632860 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 3 Apr 2025 17:31:34 -0700 Subject: [PATCH 03/13] Do same thing for writes --- libs/langgraph/langgraph/pregel/loop.py | 44 ++++++++++++++++++++++++- 1 file changed, 43 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 302bb4744..4a42f95a5 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -316,7 +316,7 @@ class PregelLoop(LoopProtocol): self.checkpoint_pending_writes[idx] = (task_id, c, v) else: self.checkpoint_pending_writes.append((task_id, c, v)) - if self.checkpointer_put_writes is not None: + if self.checkpoint_during and self.checkpointer_put_writes is not None: config = patch_configurable( self.checkpoint_config, { @@ -349,6 +349,46 @@ class PregelLoop(LoopProtocol): if hasattr(self, "tasks"): self._output_writes(task_id, writes) + def _put_pending_writes(self) -> None: + if self.checkpointer_put_writes is None: + return + if not self.checkpoint_pending_writes: + return + # group by task id + by_task = defaultdict(list) + for task_id, channel, value in self.checkpoint_pending_writes: + by_task[task_id].append((channel, value)) + # patch config with checkpoint id + config = patch_configurable( + self.checkpoint_config, + { + CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( + CONFIG_KEY_CHECKPOINT_NS, "" + ), + CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"], + }, + ) + # submit writes to checkpointer + for task_id, writes in by_task.items(): + if self.checkpointer_put_writes_accepts_task_path and hasattr( + self, "tasks" + ): + task = self.tasks.get(task_id) + self.submit( + self.checkpointer_put_writes, + config, + writes, + task_id, + task_path_str(task.path) if task else "", + ) + else: + self.submit( + self.checkpointer_put_writes, + config, + writes, + task_id, + ) + def accept_push( self, task: PregelExecutableTask, write_idx: int, call: Optional[Call] = None ) -> Optional[PregelExecutableTask]: @@ -794,7 +834,9 @@ class PregelLoop(LoopProtocol): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: + # persist current checkpoint and writes if not self.checkpoint_during: + self._put_pending_writes() self._put_checkpoint(self.checkpoint_metadata, force=True) # suppress interrupt suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested From 7ebd6f5e1f9d463f78f9b77d072e991a70540e6c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 4 Apr 2025 10:01:35 -0700 Subject: [PATCH 04/13] Better test --- libs/langgraph/tests/test_interruption.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/libs/langgraph/tests/test_interruption.py b/libs/langgraph/tests/test_interruption.py index 9d262bffc..f305d94d6 100644 --- a/libs/langgraph/tests/test_interruption.py +++ b/libs/langgraph/tests/test_interruption.py @@ -3,8 +3,8 @@ from typing_extensions import TypedDict from langgraph.graph import END, START, StateGraph from tests.conftest import ( - ALL_CHECKPOINTERS_ASYNC, - ALL_CHECKPOINTERS_SYNC, + REGULAR_CHECKPOINTERS_ASYNC, + REGULAR_CHECKPOINTERS_SYNC, awith_checkpointer, ) @@ -12,7 +12,7 @@ pytestmark = pytest.mark.anyio @pytest.mark.parametrize("checkpoint_during", [True, False]) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_interruption_without_state_updates( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: @@ -42,16 +42,22 @@ def test_interruption_without_state_updates( graph.invoke(initial_input, thread, checkpoint_during=checkpoint_during) assert graph.get_state(thread).next == ("step_2",) + n_checkpoints = len([c for c in graph.get_state_history(thread)]) + assert n_checkpoints == (3 if checkpoint_during else 1) graph.invoke(None, thread, checkpoint_during=checkpoint_during) assert graph.get_state(thread).next == ("step_3",) + n_checkpoints = len([c for c in graph.get_state_history(thread)]) + assert n_checkpoints == (4 if checkpoint_during else 2) graph.invoke(None, thread, checkpoint_during=checkpoint_during) assert graph.get_state(thread).next == () + n_checkpoints = len([c for c in graph.get_state_history(thread)]) + assert n_checkpoints == (5 if checkpoint_during else 3) @pytest.mark.parametrize("checkpoint_during", [True, False]) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_interruption_without_state_updates_async( checkpointer_name: str, checkpoint_during: bool ) -> None: @@ -81,9 +87,15 @@ async def test_interruption_without_state_updates_async( await graph.ainvoke(initial_input, thread, checkpoint_during=checkpoint_during) assert (await graph.aget_state(thread)).next == ("step_2",) + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (3 if checkpoint_during else 1) await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during) assert (await graph.aget_state(thread)).next == ("step_3",) + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (4 if checkpoint_during else 2) await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during) assert (await graph.aget_state(thread)).next == () + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (5 if checkpoint_during else 3) From 4f353dac3107f03e4ce2bab298ee695337aecf0d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 4 Apr 2025 14:41:10 -0700 Subject: [PATCH 05/13] Fix assignment of pending writes --- libs/langgraph/langgraph/pregel/loop.py | 62 +++++++++++-------------- 1 file changed, 28 insertions(+), 34 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 4a42f95a5..6333c932f 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -358,16 +358,6 @@ class PregelLoop(LoopProtocol): by_task = defaultdict(list) for task_id, channel, value in self.checkpoint_pending_writes: by_task[task_id].append((channel, value)) - # patch config with checkpoint id - config = patch_configurable( - self.checkpoint_config, - { - CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( - CONFIG_KEY_CHECKPOINT_NS, "" - ), - CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"], - }, - ) # submit writes to checkpointer for task_id, writes in by_task.items(): if self.checkpointer_put_writes_accepts_task_path and hasattr( @@ -376,7 +366,7 @@ class PregelLoop(LoopProtocol): task = self.tasks.get(task_id) self.submit( self.checkpointer_put_writes, - config, + self.checkpoint_config, writes, task_id, task_path_str(task.path) if task else "", @@ -384,7 +374,7 @@ class PregelLoop(LoopProtocol): else: self.submit( self.checkpointer_put_writes, - config, + self.checkpoint_config, writes, task_id, ) @@ -749,31 +739,34 @@ class PregelLoop(LoopProtocol): ) return updated_channels - def _put_checkpoint( - self, metadata: CheckpointMetadata, force: bool = False - ) -> None: + def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: # assign step and parents - metadata["step"] = self.step - metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) - self.checkpoint_metadata = metadata - # debug flag - if self.debug: - print_step_checkpoint( - metadata, - self.channels, - ( - [self.stream_keys] - if isinstance(self.stream_keys, str) - else self.stream_keys - ), - ) + exiting = metadata is self.checkpoint_metadata + if not exiting: + metadata["step"] = self.step + metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) + self.checkpoint_metadata = metadata + # debug flag + if self.debug: + print_step_checkpoint( + metadata, + self.channels, + ( + [self.stream_keys] + if isinstance(self.stream_keys, str) + else self.stream_keys + ), + ) # do checkpoint? do_checkpoint = self._checkpointer_put_after_previous is not None and ( - force or self.checkpoint_during + exiting or self.checkpoint_during ) # create new checkpoint self.checkpoint = create_checkpoint( - self.checkpoint, self.channels if do_checkpoint else None, self.step + self.checkpoint, + self.channels if do_checkpoint else None, + self.step, + id=self.checkpoint["id"] if exiting else None, ) # bail if no checkpointer if do_checkpoint and self._checkpointer_put_after_previous is not None: @@ -822,8 +815,9 @@ class PregelLoop(LoopProtocol): CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"], }, } - # increment step - self.step += 1 + if not exiting: + # increment step + self.step += 1 def _update_mv(self, key: str, values: Sequence[Any]) -> None: raise NotImplementedError @@ -836,8 +830,8 @@ class PregelLoop(LoopProtocol): ) -> Optional[bool]: # persist current checkpoint and writes if not self.checkpoint_during: + self._put_checkpoint(self.checkpoint_metadata) self._put_pending_writes() - self._put_checkpoint(self.checkpoint_metadata, force=True) # suppress interrupt suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested if suppress: From a5495e84c8d2a1824ec23efa032d09b6d9ce7267 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 4 Apr 2025 15:42:11 -0700 Subject: [PATCH 06/13] Add another test --- libs/langgraph/langgraph/pregel/loop.py | 29 +++++++----------- libs/langgraph/tests/test_pregel.py | 35 ++++++++++++++++----- libs/langgraph/tests/test_pregel_async.py | 37 ++++++++++++++++++----- 3 files changed, 67 insertions(+), 34 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 6333c932f..86b9ba8dc 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -180,6 +180,7 @@ class PregelLoop(LoopProtocol): channels: Mapping[str, BaseChannel] managed: ManagedValueMapping checkpoint: Checkpoint + checkpoint_id_saved: str checkpoint_ns: tuple[str, ...] checkpoint_config: RunnableConfig checkpoint_metadata: CheckpointMetadata @@ -297,25 +298,12 @@ class PregelLoop(LoopProtocol): # deduplicate writes to special channels, last write wins if all(w[0] in WRITES_IDX_MAP for w in writes): writes = list({w[0]: w for w in writes}.values()) + # remove existing writes for this task + self.checkpoint_pending_writes = [ + w for w in self.checkpoint_pending_writes if w[0] != task_id + ] # save writes - for c, v in writes: - if ( - c in WRITES_IDX_MAP - and ( - idx := next( - ( - i - for i, w in enumerate(self.checkpoint_pending_writes) - if w[0] == task_id and w[1] == c - ), - None, - ) - ) - is not None - ): - self.checkpoint_pending_writes[idx] = (task_id, c, v) - else: - self.checkpoint_pending_writes.append((task_id, c, v)) + self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes) if self.checkpoint_during and self.checkpointer_put_writes is not None: config = patch_configurable( self.checkpoint_config, @@ -742,6 +730,9 @@ class PregelLoop(LoopProtocol): def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: # assign step and parents exiting = metadata is self.checkpoint_metadata + if exiting and self.checkpoint["id"] == self.checkpoint_id_saved: + # checkpoint already saved + return if not exiting: metadata["step"] = self.step metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) @@ -1049,6 +1040,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): }, } self.prev_checkpoint_config = saved.parent_config + self.checkpoint_id_saved = saved.checkpoint["id"] self.checkpoint = saved.checkpoint self.checkpoint_metadata = saved.metadata self.checkpoint_pending_writes = ( @@ -1198,6 +1190,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): }, } self.prev_checkpoint_config = saved.parent_config + self.checkpoint_id_saved = saved.checkpoint["id"] self.checkpoint = saved.checkpoint self.checkpoint_metadata = saved.metadata self.checkpoint_pending_writes = ( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 6d47f36e7..e03ceeb72 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1115,10 +1115,14 @@ def test_invoke_checkpoint_two( assert checkpoint["channel_values"].get("total") == 5 +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_pending_writes_resume( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + checkpointer: BaseCheckpointSaver = request.getfixturevalue( f"checkpointer_{checkpointer_name}" ) @@ -1144,17 +1148,19 @@ def test_pending_writes_resume( self.calls = 0 one = AwhileMaker(0.1, {"value": 2}) - two = AwhileMaker(0.3, ConnectionError("I'm not good")) + two = AwhileMaker(0.2, ConnectionError("I'm not good")) builder = StateGraph(State) builder.add_node("one", one) - builder.add_node("two", two, retry=RetryPolicy(max_attempts=2)) + builder.add_node( + "two", two, retry=RetryPolicy(max_attempts=2, initial_interval=0, jitter=False) + ) builder.add_edge(START, "one") builder.add_edge(START, "two") graph = builder.compile(checkpointer=checkpointer) thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} with pytest.raises(ConnectionError, match="I'm not good"): - graph.invoke({"value": 1}, thread1) + graph.invoke({"value": 1}, thread1, checkpoint_during=checkpoint_during) # both nodes should have been called once assert one.calls == 1 @@ -1200,7 +1206,7 @@ def test_pending_writes_resume( # resume execution with pytest.raises(ConnectionError, match="I'm not good"): - graph.invoke(None, thread1) + graph.invoke(None, thread1, checkpoint_during=checkpoint_during) # node "one" succeeded previously, so shouldn't be called again assert one.calls == 1 @@ -1214,7 +1220,9 @@ def test_pending_writes_resume( # resume execution, without exception two.rtn = {"value": 3} # both the pending write and the new write were applied, 1 + 2 + 3 = 6 - assert graph.invoke(None, thread1) == {"value": 6} + assert graph.invoke(None, thread1, checkpoint_during=checkpoint_during) == { + "value": 6 + } if "shallow" in checkpointer_name: assert len(list(checkpointer.list(thread1))) == 1 @@ -1223,7 +1231,7 @@ def test_pending_writes_resume( # check all final checkpoints checkpoints = [c for c in checkpointer.list(thread1)] # we should have 3 - assert len(checkpoints) == 3 + assert len(checkpoints) == (3 if checkpoint_during else 2) # the last one not too interesting for this test assert checkpoints[0] == CheckpointTuple( config={ @@ -1327,13 +1335,24 @@ def test_pending_writes_resume( "checkpoint_ns": "", "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"], } - }, + } + if checkpoint_during + else None, pending_writes=UnsortedSequence( (AnyStr(), "value", 2), (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), (AnyStr(), "value", 3), + ) + if checkpoint_during + else UnsortedSequence( + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), + # the write against the previous checkpoint is not saved, as it is + # produced in a run where only the next checkpoint (the last) is saved ), ) + if not checkpoint_during: + return assert checkpoints[2] == CheckpointTuple( config={ "configurable": { diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 327236d7f..9fbf042fc 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1947,10 +1947,14 @@ async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str) assert checkpoint["channel_values"].get("total") == 5 +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_pending_writes_resume( - request: pytest.FixtureRequest, checkpointer_name: str + checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + class State(TypedDict): value: Annotated[int, operator.add] @@ -1972,10 +1976,12 @@ async def test_pending_writes_resume( self.calls = 0 one = AwhileMaker(0.1, {"value": 2}) - two = AwhileMaker(0.3, ConnectionError("I'm not good")) + two = AwhileMaker(0.2, ConnectionError("I'm not good")) builder = StateGraph(State) builder.add_node("one", one) - builder.add_node("two", two, retry=RetryPolicy(max_attempts=2)) + builder.add_node( + "two", two, retry=RetryPolicy(max_attempts=2, initial_interval=0, jitter=False) + ) builder.add_edge(START, "one") builder.add_edge(START, "two") async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -1983,7 +1989,9 @@ async def test_pending_writes_resume( thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} with pytest.raises(ConnectionError, match="I'm not good"): - await graph.ainvoke({"value": 1}, thread1) + await graph.ainvoke( + {"value": 1}, thread1, checkpoint_during=checkpoint_during + ) # both nodes should have been called once assert one.calls == 1 @@ -2034,7 +2042,7 @@ async def test_pending_writes_resume( # resume execution with pytest.raises(ConnectionError, match="I'm not good"): - await graph.ainvoke(None, thread1) + await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during) # node "one" succeeded previously, so shouldn't be called again assert one.calls == 1 @@ -2048,7 +2056,9 @@ async def test_pending_writes_resume( # resume execution, without exception two.rtn = {"value": 3} # both the pending write and the new write were applied, 1 + 2 + 3 = 6 - assert await graph.ainvoke(None, thread1) == {"value": 6} + assert await graph.ainvoke( + None, thread1, checkpoint_during=checkpoint_during + ) == {"value": 6} if "shallow" in checkpointer_name: assert len([c async for c in checkpointer.alist(thread1)]) == 1 @@ -2057,7 +2067,7 @@ async def test_pending_writes_resume( # check all final checkpoints checkpoints = [c async for c in checkpointer.alist(thread1)] # we should have 3 - assert len(checkpoints) == 3 + assert len(checkpoints) == (3 if checkpoint_during else 2) # the last one not too interesting for this test assert checkpoints[0] == CheckpointTuple( config={ @@ -2165,13 +2175,24 @@ async def test_pending_writes_resume( "checkpoint_id" ], } - }, + } + if checkpoint_during + else None, pending_writes=UnsortedSequence( (AnyStr(), "value", 2), (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), (AnyStr(), "value", 3), + ) + if checkpoint_during + else UnsortedSequence( + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), + # the write against the previous checkpoint is not saved, as it is + # produced in a run where only the next checkpoint (the last) is saved ), ) + if not checkpoint_during: + return assert checkpoints[2] == CheckpointTuple( config={ "configurable": { From 4abfc7702d140a2e6cb2d2679f223c63a4c2b5f9 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 4 Apr 2025 16:00:22 -0700 Subject: [PATCH 07/13] Subgraphs inherit checkpoint mode --- libs/langgraph/langgraph/constants.py | 2 ++ libs/langgraph/langgraph/pregel/__init__.py | 23 +++++++++++++++------ 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 4fde5e7dd..8bcbb2b94 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -83,6 +83,8 @@ CONFIG_KEY_PREVIOUS = sys.intern("__pregel_previous") # holds the previous return value from a stateful Pregel graph. CONFIG_KEY_RUNNER_SUBMIT = sys.intern("__pregel_runner_submit") # holds a function that receives tasks from runner, executes them and returns results +CONFIG_KEY_CHECKPOINT_DURING = sys.intern("__pregel_checkpoint_during") +# holds a boolean indicating whether to checkpoint during the run (or only at the end) # --- Other constants --- PUSH = sys.intern("__pregel_push") diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index cbfbd12b9..3a182ea3c 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -54,6 +54,7 @@ from langgraph.checkpoint.base import ( ) from langgraph.constants import ( CONF, + CONFIG_KEY_CHECKPOINT_DURING, CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_CHECKPOINTER, @@ -2094,7 +2095,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, - checkpoint_during: bool = True, + checkpoint_during: Optional[bool] = None, debug: Optional[bool] = None, subgraphs: bool = False, ) -> Iterator[Union[dict[str, Any], Any]]: @@ -2278,6 +2279,9 @@ class Pregel(PregelProtocol): config[CONF][CONFIG_KEY_STREAM_WRITER] = lambda c: stream.put( ((), "custom", c) ) + # set checkpointing mode for subgraphs + if checkpoint_during is not None: + config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during with SyncPregelLoop( input, input_model=self.input_model, @@ -2293,7 +2297,9 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, - checkpoint_during=checkpoint_during, + checkpoint_during=checkpoint_during + if checkpoint_during is not None + else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True), trigger_to_nodes=self.trigger_to_nodes, migrate_checkpoint=self._migrate_checkpoint, ) as loop: @@ -2376,7 +2382,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, - checkpoint_during: bool = True, + checkpoint_during: Optional[bool] = None, debug: Optional[bool] = None, subgraphs: bool = False, ) -> AsyncIterator[Union[dict[str, Any], Any]]: @@ -2576,6 +2582,9 @@ class Pregel(PregelProtocol): stream.put_nowait, ((), "custom", c) ) ) + # set checkpointing mode for subgraphs + if checkpoint_during is not None: + config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during async with AsyncPregelLoop( input, input_model=self.input_model, @@ -2591,7 +2600,9 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, - checkpoint_during=checkpoint_during, + checkpoint_during=checkpoint_during + if checkpoint_during is not None + else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True), trigger_to_nodes=self.trigger_to_nodes, migrate_checkpoint=self._migrate_checkpoint, ) as loop: @@ -2667,7 +2678,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, - checkpoint_during: bool = True, + checkpoint_during: Optional[bool] = None, debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: @@ -2721,7 +2732,7 @@ class Pregel(PregelProtocol): output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, - checkpoint_during: bool = True, + checkpoint_during: Optional[bool] = None, debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: From 5a0228cb13b88d91a7f94d7273e28d48a45a8dcf Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 4 Apr 2025 16:00:28 -0700 Subject: [PATCH 08/13] Add test --- libs/langgraph/tests/test_pregel.py | 14 +++++++++++--- libs/langgraph/tests/test_pregel_async.py | 21 ++++++++++++++++----- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index e03ceeb72..97405e41c 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1510,8 +1510,14 @@ def test_send_sequences() -> None: ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> None: +def test_imp_task( + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool +) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") mapper_calls = 0 @@ -1577,7 +1583,7 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non } thread1 = {"configurable": {"thread_id": "1"}} - assert [*graph.stream([0, 1], thread1)] == [ + assert [*graph.stream([0, 1], thread1, checkpoint_during=checkpoint_during)] == [ {"mapper": "00"}, {"mapper": "11"}, { @@ -1593,7 +1599,9 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non ] assert mapper_calls == 2 - assert graph.invoke(Command(resume="answer"), thread1) == [ + assert graph.invoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ "00answer", "11answer", ] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 9fbf042fc..467e97f53 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2230,7 +2230,7 @@ async def test_pending_writes_resume( @pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_run_from_checkpoint_id_retains_previous_writes( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + checkpointer_name: str, ) -> None: class MyState(TypedDict): myval: Annotated[int, operator.add] @@ -2275,8 +2275,8 @@ async def test_run_from_checkpoint_id_retains_previous_writes( history = [c async for c in graph.aget_state_history(thread1)] assert len(history) == 4 - assert history[-1].values == {"myval": 0} assert history[0].values == {"myval": 4, "otherval": False} + assert history[-1].values == {"myval": 0} second_run_config = { **thread1, @@ -2453,8 +2453,12 @@ async def test_send_sequences(checkpointer_name: str) -> None: @NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_task(checkpointer_name: str) -> None: +async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + async with awith_checkpointer(checkpointer_name) as checkpointer: mapper_calls = 0 @@ -2474,7 +2478,12 @@ async def test_imp_task(checkpointer_name: str) -> None: tracer = FakeTracer() thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]} - assert [c async for c in graph.astream([0, 1], thread1)] == [ + assert [ + c + async for c in graph.astream( + [0, 1], thread1, checkpoint_during=checkpoint_during + ) + ] == [ {"mapper": "00"}, {"mapper": "11"}, { @@ -2498,7 +2507,9 @@ async def test_imp_task(checkpointer_name: str) -> None: assert any(r.inputs == {"input": 0} for r in mapper_runs) assert any(r.inputs == {"input": 1} for r in mapper_runs) - assert await graph.ainvoke(Command(resume="answer"), thread1) == [ + assert await graph.ainvoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ "00answer", "11answer", ] From d541ed90d571efeb783ba3d723efe1c294c06527 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 7 Apr 2025 16:45:26 -0700 Subject: [PATCH 09/13] Save Sends unconditionally --- libs/langgraph/langgraph/pregel/loop.py | 20 ++++- libs/langgraph/tests/test_large_cases.py | 32 ++++--- libs/langgraph/tests/test_pregel.py | 31 +++++-- libs/langgraph/tests/test_pregel_async.py | 101 +++++++++++++++++----- 4 files changed, 131 insertions(+), 53 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 86b9ba8dc..772564a50 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -63,6 +63,7 @@ from langgraph.constants import ( RESUME, SCHEDULED, TAG_HIDDEN, + TASKS, ) from langgraph.errors import ( CheckpointNotLatest, @@ -295,6 +296,7 @@ class PregelLoop(LoopProtocol): """Put writes for a task, to be read by the next tick.""" if not writes: return + checkpoint_during = self.checkpoint_during or any(w[0] == TASKS for w in writes) # deduplicate writes to special channels, last write wins if all(w[0] in WRITES_IDX_MAP for w in writes): writes = list({w[0]: w for w in writes}.values()) @@ -304,7 +306,7 @@ class PregelLoop(LoopProtocol): ] # save writes self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes) - if self.checkpoint_during and self.checkpointer_put_writes is not None: + if checkpoint_during and self.checkpointer_put_writes is not None: config = patch_configurable( self.checkpoint_config, { @@ -342,6 +344,16 @@ class PregelLoop(LoopProtocol): return if not self.checkpoint_pending_writes: return + # patch config + config = patch_configurable( + self.checkpoint_config, + { + CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( + CONFIG_KEY_CHECKPOINT_NS, "" + ), + CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"], + }, + ) # group by task id by_task = defaultdict(list) for task_id, channel, value in self.checkpoint_pending_writes: @@ -354,7 +366,7 @@ class PregelLoop(LoopProtocol): task = self.tasks.get(task_id) self.submit( self.checkpointer_put_writes, - self.checkpoint_config, + config, writes, task_id, task_path_str(task.path) if task else "", @@ -362,7 +374,7 @@ class PregelLoop(LoopProtocol): else: self.submit( self.checkpointer_put_writes, - self.checkpoint_config, + config, writes, task_id, ) @@ -748,6 +760,7 @@ class PregelLoop(LoopProtocol): else self.stream_keys ), ) + self.checkpoint_id_prev = self.checkpoint["id"] if self.step > -1 else None # do checkpoint? do_checkpoint = self._checkpointer_put_after_previous is not None and ( exiting or self.checkpoint_during @@ -776,6 +789,7 @@ class PregelLoop(LoopProtocol): **self.checkpoint_config, CONF: { **self.checkpoint_config[CONF], + CONFIG_KEY_CHECKPOINT_ID: self.checkpoint_id_prev, CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( CONFIG_KEY_CHECKPOINT_NS, "" ), diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index f15cd5c23..f5be91cee 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -7258,9 +7258,10 @@ def test_branch_then( ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +@pytest.mark.parametrize("checkpoint_during", [True, False]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_send_dedupe_on_resume( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @@ -7316,7 +7317,7 @@ def test_send_dedupe_on_resume( graph = builder.compile(checkpointer=checkpointer) thread1 = {"configurable": {"thread_id": "1"}} - assert graph.invoke(["0"], thread1, debug=1) == [ + assert graph.invoke(["0"], thread1, checkpoint_during=checkpoint_during) == [ "0", "1", "3.1", @@ -7333,12 +7334,11 @@ def test_send_dedupe_on_resume( pytest.xfail("TODO: shallow checkpointer reports wrong next set") assert state.next == ("flaky",) # check history - if "shallow" not in checkpointer_name: - history = [c for c in graph.get_state_history(thread1)] - assert len(history) == 4 + history = [c for c in graph.get_state_history(thread1)] + assert len(history) == (4 if checkpoint_during else 1) # resume execution - assert graph.invoke(None, thread1, debug=1) == [ + assert graph.invoke(None, thread1, checkpoint_during=checkpoint_during) == [ "0", "1", "3.1", @@ -7358,6 +7358,7 @@ def test_send_dedupe_on_resume( assert state.next == () # check history history = [c for c in graph.get_state_history(thread1)] + assert len(history) == (6 if checkpoint_during else 2) expected_history = [ StateSnapshot( values=[ @@ -7494,13 +7495,9 @@ def test_send_dedupe_on_resume( name="flaky", path=("__pregel_push", 1), error=None, - interrupts=( - Interrupt( - value="Bahh", resumable=False, ns=None, when="during" - ), - ), + interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),), state=None, - result=["flaky|4"], + result=["flaky|4"] if checkpoint_during else None, ), PregelTask( id=AnyStr(), @@ -7637,10 +7634,11 @@ def test_send_dedupe_on_resume( ), ), ] - if "shallow" in checkpointer_name: - expected_history = expected_history[:1] - - assert history == expected_history + if checkpoint_during: + assert history == expected_history + else: + assert history[0] == expected_history[0] + assert history[1] == expected_history[2] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 97405e41c..f1d45c192 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1333,11 +1333,11 @@ def test_pending_writes_resume( "configurable": { "thread_id": "1", "checkpoint_ns": "", - "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"], + "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"] + if checkpoint_during + else AnyStr(), } - } - if checkpoint_during - else None, + }, pending_writes=UnsortedSequence( (AnyStr(), "value", 2), (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), @@ -1608,10 +1608,14 @@ def test_imp_task( assert mapper_calls == 2 +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_nested( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") def mynode(input: list[str]) -> list[str]: @@ -1653,7 +1657,7 @@ def test_imp_nested( } thread1 = {"configurable": {"thread_id": "1"}} - assert [*graph.stream([0, 1], thread1)] == [ + assert [*graph.stream([0, 1], thread1, checkpoint_during=checkpoint_during)] == [ {"submapper": "0"}, {"mapper": "00"}, {"submapper": "1"}, @@ -1670,16 +1674,22 @@ def test_imp_nested( }, ] - assert graph.invoke(Command(resume="answer"), thread1) == [ + assert graph.invoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ "00answera", "11answera", ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_stream_order( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @task() @@ -1702,7 +1712,10 @@ def test_imp_stream_order( return fut_baz.result() thread1 = {"configurable": {"thread_id": "1"}} - assert [c for c in graph.stream({"a": "0"}, thread1)] == [ + assert [ + c + for c in graph.stream({"a": "0"}, thread1, checkpoint_during=checkpoint_during) + ] == [ { "foo": ( "0foo", diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 467e97f53..4cc115f6f 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2173,11 +2173,11 @@ async def test_pending_writes_resume( "checkpoint_ns": "", "checkpoint_id": checkpoints[2].config["configurable"][ "checkpoint_id" - ], + ] + if checkpoint_during + else AnyStr(), } - } - if checkpoint_during - else None, + }, pending_writes=UnsortedSequence( (AnyStr(), "value", 2), (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), @@ -2517,8 +2517,12 @@ async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None @NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_nested(checkpointer_name: str) -> None: +async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + async def mynode(input: list[str]) -> list[str]: return [it + "a" for it in input] @@ -2558,7 +2562,12 @@ async def test_imp_nested(checkpointer_name: str) -> None: } thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream([0, 1], thread1)] == [ + assert [ + c + async for c in graph.astream( + [0, 1], thread1, checkpoint_during=checkpoint_during + ) + ] == [ {"submapper": "0"}, {"mapper": "00"}, {"submapper": "1"}, @@ -2575,15 +2584,21 @@ async def test_imp_nested(checkpointer_name: str) -> None: }, ] - assert await graph.ainvoke(Command(resume="answer"), thread1) == [ + assert await graph.ainvoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ "00answera", "11answera", ] @NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_task_cancel(checkpointer_name: str) -> None: +async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + async with awith_checkpointer(checkpointer_name) as checkpointer: mapper_calls = 0 mapper_cancels = 0 @@ -2609,7 +2624,12 @@ async def test_imp_task_cancel(checkpointer_name: str) -> None: return [m + answer for m in mapped] thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream([0, 1], thread1)] == [ + assert [ + c + async for c in graph.astream( + [0, 1], thread1, checkpoint_during=checkpoint_during + ) + ] == [ {"mapper": "00"}, { "__interrupt__": ( @@ -2625,7 +2645,9 @@ async def test_imp_task_cancel(checkpointer_name: str) -> None: assert mapper_calls == 2 assert mapper_cancels == 1 - assert await graph.ainvoke(Command(resume="answer"), thread1) == [ + assert await graph.ainvoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ "00answer", ] assert mapper_calls == 3 @@ -2633,8 +2655,14 @@ async def test_imp_task_cancel(checkpointer_name: str) -> None: @NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_sync_from_async(checkpointer_name: str) -> None: +async def test_imp_sync_from_async( + checkpointer_name: str, checkpoint_during: bool +) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + async with awith_checkpointer(checkpointer_name) as checkpointer: @task() @@ -2657,7 +2685,12 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None: return fut_baz.result() thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream({"a": "0"}, thread1)] == [ + assert [ + c + async for c in graph.astream( + {"a": "0"}, thread1, checkpoint_during=checkpoint_during + ) + ] == [ {"foo": {"a": "0foo", "b": "bar"}}, {"bar": {"a": "0foobar", "c": "bark"}}, {"baz": {"a": "0foobarbaz", "c": "something else"}}, @@ -2666,8 +2699,14 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None: @NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_stream_order(checkpointer_name: str) -> None: +async def test_imp_stream_order( + checkpointer_name: str, checkpoint_during: bool +) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Checkpointing during execution not supported") + async with awith_checkpointer(checkpointer_name) as checkpointer: @task() @@ -2691,7 +2730,12 @@ async def test_imp_stream_order(checkpointer_name: str) -> None: return await fut_baz thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream({"a": "0"}, thread1)] == [ + assert [ + c + async for c in graph.astream( + {"a": "0"}, thread1, checkpoint_during=checkpoint_during + ) + ] == [ {"foo": {"a": "0foo", "b": "bar"}}, {"bar": {"a": "0foobar", "c": "bark"}}, {"baz": {"a": "0foobarbaz", "c": "something else"}}, @@ -2699,8 +2743,11 @@ async def test_imp_stream_order(checkpointer_name: str) -> None: ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) -async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: +async def test_send_dedupe_on_resume( + checkpointer_name: str, checkpoint_during: bool +) -> None: class InterruptOnce: ticks: int = 0 @@ -2751,7 +2798,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: graph = builder.compile(checkpointer=checkpointer) thread1 = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke(["0"], thread1, debug=1) == [ + assert await graph.ainvoke( + ["0"], thread1, checkpoint_during=checkpoint_during + ) == [ "0", "1", "3.1", @@ -2763,7 +2812,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: assert builder.nodes["2"].runnable.func.ticks == 3 assert builder.nodes["flaky"].runnable.func.ticks == 1 # resume execution - assert await graph.ainvoke(None, thread1, debug=1) == [ + assert await graph.ainvoke( + None, thread1, checkpoint_during=checkpoint_during + ) == [ "0", "1", "3.1", @@ -2780,7 +2831,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: assert builder.nodes["flaky"].runnable.func.ticks == 2 # check history history = [c async for c in graph.aget_state_history(thread1)] - assert history == [ + assert len(history) == (6 if checkpoint_during else 2) + expected_history = [ StateSnapshot( values=[ "0", @@ -2916,13 +2968,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: name="flaky", path=("__pregel_push", 1), error=None, - interrupts=( - Interrupt( - value="Bahh", resumable=False, ns=None, when="during" - ), - ), + interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),), state=None, - result=["flaky|4"], + result=["flaky|4"] if checkpoint_during else None, ), PregelTask( id=AnyStr(), @@ -3059,6 +3107,11 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: ), ), ] + if checkpoint_during: + assert history == expected_history + else: + assert history[0] == expected_history[0] + assert history[1] == expected_history[2] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) From e9aec77893932072cd0ca7f6d526f0416c5085de Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 8 Apr 2025 14:07:02 -0700 Subject: [PATCH 10/13] Add more tests --- libs/langgraph/tests/test_pregel.py | 138 +++++++++--- libs/langgraph/tests/test_pregel_async.py | 247 ++++++++++++++++++++-- 2 files changed, 337 insertions(+), 48 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f1d45c192..486b4cdf9 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -3683,10 +3683,14 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) class InnerState(TypedDict): @@ -3718,7 +3722,12 @@ def test_subgraph_checkpoint_true( app = graph.compile(checkpointer=checkpointer) config = {"configurable": {"thread_id": "2"}} - assert [c for c in app.stream({"my_key": ""}, config, subgraphs=True)] == [ + assert [ + c + for c in app.stream( + {"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during + ) + ] == [ (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), (("inner",), {"inner_2": {"my_key": " and there"}}), ((), {"inner": {"my_key": " got here and there"}}), @@ -3743,10 +3752,14 @@ def test_subgraph_checkpoint_true( ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) # Define subgraph @@ -3785,15 +3798,18 @@ def test_subgraph_checkpoint_true_interrupt( builder.add_edge(START, "node_1") builder.add_edge("node_1", "node_2") - checkpointer = MemorySaver() graph = builder.compile(checkpointer=checkpointer) config = {"configurable": {"thread_id": "1"}} - assert graph.invoke({"foo": "foo"}, config) == {"foo": "hi! foo"} + assert graph.invoke( + {"foo": "foo"}, config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foo"} assert graph.get_state(config, subgraphs=True).tasks[0].state.values == { "bar": "hi! foo" } - assert graph.invoke(Command(resume="baz"), config) == {"foo": "hi! foobaz"} + assert graph.invoke( + Command(resume="baz"), config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foobaz"} @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) @@ -3909,10 +3925,14 @@ def test_stream_buffering_single_node( ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_nested_graph_interrupts_parallel( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) class InnerState(TypedDict): @@ -3959,11 +3979,11 @@ def test_nested_graph_interrupts_parallel( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": ""}, config, debug=True) == { + assert app.invoke({"my_key": ""}, config, checkpoint_during=checkpoint_during) == { "my_key": " and parallel", } - assert app.invoke(None, config, debug=True) == { + assert app.invoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "got here and there and parallel and back again", } @@ -3972,13 +3992,17 @@ def test_nested_graph_interrupts_parallel( # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) # test stream updates w/ nested interrupt config = {"configurable": {"thread_id": "2"}} - assert [*app.stream({"my_key": ""}, config, subgraphs=True)] == [ + assert [ + *app.stream( + {"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during + ) + ] == [ # we got to parallel node first ((), {"outer_1": {"my_key": " and parallel"}}), ((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}), ((), {"__interrupt__": ()}), ] - assert [*app.stream(None, config)] == [ + assert [*app.stream(None, config, checkpoint_during=checkpoint_during)] == [ {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, {"inner": {"my_key": "got here and there"}}, {"outer_2": {"my_key": " and back again"}}, @@ -3986,11 +4010,22 @@ def test_nested_graph_interrupts_parallel( # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}} - assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -3999,15 +4034,28 @@ def test_nested_graph_interrupts_parallel( # test interrupts BEFORE the parallel node app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) config = {"configurable": {"thread_id": "4"}} - assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ - {"my_key": ""} - ] + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [{"my_key": ""}] # while we're waiting for the node w/ interrupt inside to finish - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -4016,24 +4064,43 @@ def test_nested_graph_interrupts_parallel( # test interrupts AFTER the parallel node app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) config = {"configurable": {"thread_id": "5"}} - assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_doubly_nested_graph_interrupts( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) class State(TypedDict): @@ -4087,11 +4154,13 @@ def test_doubly_nested_graph_interrupts( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": "my value"}, config, debug=True) == { + assert app.invoke( + {"my_key": "my value"}, config, checkpoint_during=checkpoint_during + ) == { "my_key": "hi my value", } - assert app.invoke(None, config, debug=True) == { + assert app.invoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "hi my value here and there and back again", } @@ -4100,12 +4169,14 @@ def test_doubly_nested_graph_interrupts( config = { "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} } - assert [*app.stream({"my_key": "my value"}, config)] == [ + assert [ + *app.stream({"my_key": "my value"}, config, checkpoint_during=checkpoint_during) + ] == [ {"parent_1": {"my_key": "hi my value"}}, {"__interrupt__": ()}, ] assert nodes == ["parent_1", "grandchild_1"] - assert [*app.stream(None, config)] == [ + assert [*app.stream(None, config, checkpoint_during=checkpoint_during)] == [ {"child": {"my_key": "hi my value here and there"}}, {"parent_2": {"my_key": "hi my value here and there and back again"}}, ] @@ -4120,11 +4191,22 @@ def test_doubly_nested_graph_interrupts( # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}} - assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ + assert [ + *app.stream( + {"my_key": "my value"}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ {"my_key": "my value"}, {"my_key": "hi my value"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "hi my value"}, {"my_key": "hi my value here and there"}, {"my_key": "hi my value here and there and back again"}, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 4cc115f6f..20e2b96b5 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5433,6 +5433,131 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 +@pytest.mark.parametrize("checkpoint_during", [True, False]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_subgraph_checkpoint_true( + checkpointer_name: str, checkpoint_during: bool +) -> None: + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + return {"my_key": " got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + return {"my_key": " and there"} + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + graph = StateGraph(State) + graph.add_node("inner", inner.compile(checkpointer=True)) + graph.add_edge(START, "inner") + graph.add_conditional_edges( + "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END + ) + + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) + + config = {"configurable": {"thread_id": "2"}} + assert [ + c + async for c in app.astream( + {"my_key": ""}, + config, + subgraphs=True, + checkpoint_during=checkpoint_during, + ) + ] == [ + (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ((), {"inner": {"my_key": " got here and there"}}), + ( + ("inner",), + { + "inner_1": { + "my_key": " got here", + "my_other_key": " got here and there got here and there", + } + }, + ), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ( + (), + { + "inner": { + "my_key": " got here and there got here and there got here and there" + } + }, + ), + ] + + +@pytest.mark.parametrize("checkpoint_during", [True, False]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_subgraph_checkpoint_true_interrupt( + checkpointer_name: str, checkpoint_during: bool +) -> None: + # Define subgraph + class SubgraphState(TypedDict): + # note that none of these keys are shared with the parent graph state + bar: str + baz: str + + def subgraph_node_1(state: SubgraphState): + baz_value = interrupt("Provide baz value") + return {"baz": baz_value} + + def subgraph_node_2(state: SubgraphState): + return {"bar": state["bar"] + state["baz"]} + + subgraph_builder = StateGraph(SubgraphState) + subgraph_builder.add_node(subgraph_node_1) + subgraph_builder.add_node(subgraph_node_2) + subgraph_builder.add_edge(START, "subgraph_node_1") + subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2") + subgraph = subgraph_builder.compile(checkpointer=True) + + class ParentState(TypedDict): + foo: str + + def node_1(state: ParentState): + return {"foo": "hi! " + state["foo"]} + + async def node_2(state: ParentState): + response = await subgraph.ainvoke({"bar": state["foo"]}) + return {"foo": response["bar"]} + + builder = StateGraph(ParentState) + builder.add_node("node_1", node_1) + builder.add_node("node_2", node_2) + builder.add_edge(START, "node_1") + builder.add_edge("node_1", "node_2") + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + assert await graph.ainvoke( + {"foo": "foo"}, config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foo"} + assert (await graph.aget_state(config, subgraphs=True)).tasks[ + 0 + ].state.values == {"bar": "hi! foo"} + assert await graph.ainvoke( + Command(resume="baz"), config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foobaz"} + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None: class InnerState(TypedDict): @@ -5541,8 +5666,11 @@ async def test_stream_buffering_single_node(checkpointer_name: str) -> None: ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: +async def test_nested_graph_interrupts_parallel( + checkpointer_name: str, checkpoint_during: bool +) -> None: class InnerState(TypedDict): my_key: Annotated[str, operator.add] my_other_key: str @@ -5591,11 +5719,13 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": ""}, config, debug=True) == { + assert await app.ainvoke( + {"my_key": ""}, config, checkpoint_during=checkpoint_during + ) == { "my_key": " and parallel", } - assert await app.ainvoke(None, config, debug=True) == { + assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "got here and there and parallel and back again", } @@ -5605,7 +5735,13 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test stream updates w/ nested interrupt config = {"configurable": {"thread_id": "2"}} assert [ - c async for c in app.astream({"my_key": ""}, config, subgraphs=True) + c + async for c in app.astream( + {"my_key": ""}, + config, + subgraphs=True, + checkpoint_during=checkpoint_during, + ) ] == [ # we got to parallel node first ((), {"outer_1": {"my_key": " and parallel"}}), @@ -5615,7 +5751,12 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: ), ((), {"__interrupt__": ()}), ] - assert [c async for c in app.astream(None, config)] == [ + assert [ + c + async for c in app.astream( + None, config, checkpoint_during=checkpoint_during + ) + ] == [ {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, {"inner": {"my_key": "got here and there"}}, {"outer_2": {"my_key": " and back again"}}, @@ -5624,12 +5765,23 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}} assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -5639,16 +5791,32 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) config = {"configurable": {"thread_id": "4"}} assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) ] == [ {"my_key": ""}, ] # while we're waiting for the node w/ interrupt inside to finish - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -5658,23 +5826,42 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) config = {"configurable": {"thread_id": "5"}} assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: +async def test_doubly_nested_graph_interrupts( + checkpointer_name: str, checkpoint_during: bool +) -> None: class State(TypedDict): my_key: str @@ -5727,11 +5914,13 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { + assert await app.ainvoke( + {"my_key": "my value"}, config, checkpoint_during=checkpoint_during + ) == { "my_key": "hi my value", } - assert await app.ainvoke(None, config, debug=True) == { + assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "hi my value here and there and back again", } @@ -5740,12 +5929,22 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: config = { "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} } - assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, checkpoint_during=checkpoint_during + ) + ] == [ {"parent_1": {"my_key": "hi my value"}}, {"__interrupt__": ()}, ] assert nodes == ["parent_1", "grandchild_1"] - assert [c async for c in app.astream(None, config)] == [ + assert [ + c + async for c in app.astream( + None, config, checkpoint_during=checkpoint_during + ) + ] == [ {"child": {"my_key": "hi my value here and there"}}, {"parent_2": {"my_key": "hi my value here and there and back again"}}, ] @@ -5763,13 +5962,21 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: assert [ c async for c in app.astream( - {"my_key": "my value"}, config, stream_mode="values" + {"my_key": "my value"}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, ) ] == [ {"my_key": "my value"}, {"my_key": "hi my value"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "hi my value"}, {"my_key": "hi my value here and there"}, {"my_key": "hi my value here and there and back again"}, From cbbfaba1fd85bd708942f763079b380714de051e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 8 Apr 2025 14:07:08 -0700 Subject: [PATCH 11/13] Add comments --- libs/langgraph/langgraph/pregel/loop.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 772564a50..e1de05602 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -296,6 +296,8 @@ class PregelLoop(LoopProtocol): """Put writes for a task, to be read by the next tick.""" if not writes: return + # always checkpoint writes containing Send, as they are fetched from the + # parent checkpoint, not the current one checkpoint_during = self.checkpoint_during or any(w[0] == TASKS for w in writes) # deduplicate writes to special channels, last write wins if all(w[0] in WRITES_IDX_MAP for w in writes): @@ -789,6 +791,7 @@ class PregelLoop(LoopProtocol): **self.checkpoint_config, CONF: { **self.checkpoint_config[CONF], + # this is guaranteed to be set by code above CONFIG_KEY_CHECKPOINT_ID: self.checkpoint_id_prev, CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( CONFIG_KEY_CHECKPOINT_NS, "" From b76dc8ae0a209fcf906216d6fcd4ef1b796f2bbe Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 8 Apr 2025 14:15:13 -0700 Subject: [PATCH 12/13] Fix --- libs/langgraph/tests/test_pregel_async.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 20e2b96b5..81f3c2289 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5533,8 +5533,8 @@ async def test_subgraph_checkpoint_true_interrupt( def node_1(state: ParentState): return {"foo": "hi! " + state["foo"]} - async def node_2(state: ParentState): - response = await subgraph.ainvoke({"bar": state["foo"]}) + async def node_2(state: ParentState, config: RunnableConfig): + response = await subgraph.ainvoke({"bar": state["foo"]}, config) return {"foo": response["bar"]} builder = StateGraph(ParentState) From 947a233fc5f2694780c2bd0855ff01be06d33121 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 8 Apr 2025 14:20:29 -0700 Subject: [PATCH 13/13] Fix --- libs/langgraph/tests/test_pregel_async.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 81f3c2289..dd146dbb7 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5502,6 +5502,7 @@ async def test_subgraph_checkpoint_true( ] +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_subgraph_checkpoint_true_interrupt( @@ -5534,7 +5535,7 @@ async def test_subgraph_checkpoint_true_interrupt( return {"foo": "hi! " + state["foo"]} async def node_2(state: ParentState, config: RunnableConfig): - response = await subgraph.ainvoke({"bar": state["foo"]}, config) + response = await subgraph.ainvoke({"bar": state["foo"]}) return {"foo": response["bar"]} builder = StateGraph(ParentState)