From a4a8934bd3563c8b0a36ba579b253f6422497ebf Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 10 Jun 2025 17:25:05 -0700 Subject: [PATCH 1/3] Avoid saving checkpoints for subgraphs when checkpoint_during=False - We can avoid saving checkpoints for successful subgraphs which do not request multi-turn memory --- libs/langgraph/langgraph/pregel/loop.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 3031ec525..7d43ca445 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -60,6 +60,7 @@ from langgraph.constants import ( INPUT, INTERRUPT, MISSING, + NS_END, NS_SEP, NULL_TASK_ID, PUSH, @@ -868,7 +869,14 @@ class PregelLoop: traceback: TracebackType | None, ) -> bool | None: # persist current checkpoint and writes - if not self.checkpoint_during: + if not self.checkpoint_during and ( + # if it's a top graph + not self.is_nested + # or a nested graph with error or interrupt + or exc_value is not None + # or a nested graph with checkpointer=True + or all(NS_END not in part for part in self.checkpoint_ns) + ): self._put_checkpoint(self.checkpoint_metadata) self._put_pending_writes() # suppress interrupt From 3b98044f2fa5f4ac719c660c4ec020f0569da336 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Tue, 10 Jun 2025 17:29:03 -0700 Subject: [PATCH 2/3] Add tests --- libs/langgraph/tests/test_pregel.py | 51 +++++++++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 45 ++++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f8e99fa0e..0a1aad429 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -3276,6 +3276,57 @@ def test_subgraph_checkpoint_true( ), ] + checkpoints = list(app.get_state_history(config)) + if checkpoint_during: + assert len(checkpoints) == 4 + else: + assert len(checkpoints) == 1 + + +def test_subgraph_checkpoint_during_false_inherited() -> None: + sync_checkpointer = InMemorySaver() + + 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 + + inner_app = inner.compile(checkpointer=sync_checkpointer) + graph = StateGraph(State) + graph.add_node("inner", inner_app) + graph.add_edge(START, "inner") + graph.add_conditional_edges( + "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END + ) + app = graph.compile(checkpointer=sync_checkpointer) + for checkpoint_during in [True, False]: + thread_id = str(uuid.uuid4()) + config = {"configurable": {"thread_id": thread_id}} + app.invoke( + {"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during + ) + if checkpoint_during: + checkpoints = list(sync_checkpointer.list(config)) + assert len(checkpoints) == 12 + else: + checkpoints = list(sync_checkpointer.list(config)) + assert len(checkpoints) == 1 + def test_subgraph_checkpoint_true_interrupt( sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index f3fb807be..7f3572a2b 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5029,6 +5029,51 @@ async def test_subgraph_checkpoint_true( ] +async def test_subgraph_checkpoint_during_false_inherited() -> None: + async_checkpointer = InMemorySaver() + + 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 + + inner_app = inner.compile(checkpointer=async_checkpointer) + graph = StateGraph(State) + graph.add_node("inner", inner_app) + graph.add_edge(START, "inner") + graph.add_conditional_edges( + "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END + ) + app = graph.compile(checkpointer=async_checkpointer) + for checkpoint_during in [True, False]: + thread_id = str(uuid.uuid4()) + config = {"configurable": {"thread_id": thread_id}} + await app.ainvoke( + {"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during + ) + if checkpoint_during: + checkpoints = list(async_checkpointer.list(config)) + assert len(checkpoints) == 12 + else: + checkpoints = list(async_checkpointer.list(config)) + assert len(checkpoints) == 1 + + @NEEDS_CONTEXTVARS async def test_subgraph_checkpoint_true_interrupt( async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool From 6a9ca8d67e084b9a7f145dd54d02b6059014c77c Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Tue, 10 Jun 2025 17:59:41 -0700 Subject: [PATCH 3/3] Update existing --- libs/langgraph/tests/test_large_cases.py | 94 +----------------- .../langgraph/tests/test_large_cases_async.py | 96 +------------------ libs/langgraph/tests/test_pregel.py | 7 +- libs/langgraph/tests/test_pregel_async.py | 7 +- 4 files changed, 12 insertions(+), 192 deletions(-) diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index c19202bf8..04ff8e5c6 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -4213,44 +4213,6 @@ def test_doubly_nested_graph_state( # get child graph history child_history = list(app.get_state_history(outer_history[1].tasks[0].state)) assert child_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child"], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=(), - interrupts=(), - ), StateSnapshot( values={"my_key": "hi my value"}, next=("child_1",), @@ -4295,62 +4257,8 @@ def test_doubly_nested_graph_state( ), ] # get grandchild graph history - grandchild_history = list(app.get_state_history(child_history[1].tasks[0].state)) + grandchild_history = list(app.get_state_history(child_history[0].tasks[0].state)) assert grandchild_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "loop", - "step": 2, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": ["branch:to:child_1"], - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - tasks=(), - interrupts=(), - ), StateSnapshot( values={"my_key": "hi my value here"}, next=("grandchild_2",), diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 2b3a99c13..b144e89e7 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -3028,44 +3028,6 @@ async def test_doubly_nested_graph_state( c async for c in app.aget_state_history(outer_history[1].tasks[0].state) ] assert child_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child"], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=(), - interrupts=(), - ), StateSnapshot( values={"my_key": "hi my value"}, next=("child_1",), @@ -3111,65 +3073,9 @@ async def test_doubly_nested_graph_state( ] # get grandchild graph history grandchild_history = [ - c async for c in app.aget_state_history(child_history[1].tasks[0].state) + c async for c in app.aget_state_history(child_history[0].tasks[0].state) ] assert grandchild_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "loop", - "step": 2, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [ - "branch:to:child_1", - ], - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - tasks=(), - interrupts=(), - ), StateSnapshot( values={"my_key": "hi my value here"}, next=("grandchild_2",), diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0a1aad429..8434f4158 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -4626,11 +4626,14 @@ def test_debug_nested_subgraphs( return clean_config - for checkpoint_events, checkpoint_history in zip( - stream_ns.values(), history_ns.values() + for checkpoint_events, checkpoint_history, ns in zip( + stream_ns.values(), history_ns.values(), stream_ns.keys() ): if not checkpoint_during: checkpoint_events = checkpoint_events[-1:] + if ns: # Save no checkpoints for subgraphs when checkpoint_during=False + assert not checkpoint_history + continue assert len(checkpoint_events) == len(checkpoint_history) for stream, history in zip(checkpoint_events, checkpoint_history): assert stream["values"] == history.values diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 7f3572a2b..db7c619c8 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6052,11 +6052,14 @@ async def test_debug_nested_subgraphs( return clean_config - for checkpoint_events, checkpoint_history in zip( - stream_ns.values(), history_ns.values() + for checkpoint_events, checkpoint_history, ns in zip( + stream_ns.values(), history_ns.values(), stream_ns.keys() ): if not checkpoint_during: checkpoint_events = checkpoint_events[-1:] + if ns: # Save no checkpoints for subgraphs when checkpoint_during=False + assert not checkpoint_history + continue assert len(checkpoint_events) == len(checkpoint_history) for stream, history in zip(checkpoint_events, checkpoint_history): assert stream["values"] == history.values