Avoid saving checkpoints for subgraphs when checkpoint_during=False (#5051)

This commit is contained in:
Nuno Campos
2025-06-11 11:11:02 -07:00
committed by GitHub
5 changed files with 117 additions and 193 deletions
+9 -1
View File
@@ -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
+1 -93
View File
@@ -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",),
+1 -95
View File
@@ -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",),
+56 -2
View File
@@ -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
@@ -4575,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
+50 -2
View File
@@ -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
@@ -6007,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