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": {