mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
Avoid saving checkpoints for subgraphs when checkpoint_during=False (#5051)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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",),
|
||||
|
||||
@@ -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",),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user