WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-29 03:09:45 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix Send order after interrupt/resume (#3037)

Browse Source 
- order was incorrectly based on task id, instead of the correct task
path
- this requires storing task paths on checkpointers
- addition of task_path to put_writes is made backwards compatible by
checking signature on call, and treating it as an optional arg
This commit is contained in:
Nuno Campos
2025-01-15 11:42:08 -08:00
committed by GitHub
parent be1d035aba 0adbd89d9a
commit aab6fdf3f3
7 changed files with 201 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/algo.py
+5 -5
View File
@@ -236,7 +236,7 @@ def apply_writes(
# sort tasks on path, to ensure deterministic order for update application
# any path parts after the 3rd are ignored for sorting
# (we use them for eg. task ids which aren't good for sorting)
tasks = sorted(tasks, key=lambda t: _tuple_str(t.path[:3]))
tasks = sorted(tasks, key=lambda t: task_path_str(t.path[:3]))
# if no task has triggers this is applying writes from the null task only
# so we don't do anything other than update the channels written to
bump_step = any(t.triggers for t in tasks)
@@ -450,7 +450,7 @@ def prepare_single_task(
str(step),
name,
PUSH,
_tuple_str(task_path[1]),
task_path_str(task_path[1]),
str(task_path[2]),
)
task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
@@ -813,10 +813,10 @@ def _uuid5_str(namespace: bytes, *parts: str) -> str:
return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}"
def _tuple_str(tup: Union[str, int, tuple]) -> str:
"""Generate a string representation of a tuple."""
def task_path_str(tup: Union[str, int, tuple]) -> str:
"""Generate a string representation of the task path."""
return (
f"~{', '.join(_tuple_str(x) for x in tup)}"
f"~{', '.join(task_path_str(x) for x in tup)}"
if isinstance(tup, (tuple, list))
else f"{tup:010d}"
if isinstance(tup, int)
libs/langgraph/langgraph/pregel/loop.py
+40 -13
View File
@@ -2,6 +2,7 @@ import asyncio
import concurrent.futures
from collections import defaultdict, deque
from contextlib import AsyncExitStack, ExitStack
from inspect import signature
from types import TracebackType
from typing import (
Any,
@@ -81,6 +82,7 @@ from langgraph.pregel.algo import (
prepare_next_tasks,
prepare_single_task,
should_interrupt,
task_path_str,
)
from langgraph.pregel.debug import (
map_debug_checkpoint,
@@ -151,6 +153,7 @@ class PregelLoop(LoopProtocol):
checkpointer_put_writes: Optional[
Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any]
]
checkpointer_put_writes_accepts_task_path: bool
_checkpointer_put_after_previous: Optional[
Callable[
[
@@ -288,20 +291,34 @@ class PregelLoop(LoopProtocol):
else:
self.checkpoint_pending_writes.append((task_id, c, v))
if self.checkpointer_put_writes is not None:
self.submit(
self.checkpointer_put_writes,
patch_configurable(
self.checkpoint_config,
{
CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
CONFIG_KEY_CHECKPOINT_NS, ""
),
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
},
),
writes,
task_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"],
},
)
if self.checkpointer_put_writes_accepts_task_path:
if hasattr(self, "tasks"):
task = self.tasks.get(task_id)
else:
task = None
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,
)
# output writes
if hasattr(self, "tasks"):
self._output_writes(task_id, writes)
@@ -813,10 +830,15 @@ class SyncPregelLoop(PregelLoop, ContextManager):
if checkpointer:
self.checkpointer_get_next_version = checkpointer.get_next_version
self.checkpointer_put_writes = checkpointer.put_writes
self.checkpointer_put_writes_accepts_task_path = (
signature(checkpointer.put_writes).parameters.get("task_path")
is not None
)
else:
self.checkpointer_get_next_version = increment
self._checkpointer_put_after_previous = None # type: ignore[assignment]
self.checkpointer_put_writes = None
self.checkpointer_put_writes_accepts_task_path = False
def _checkpointer_put_after_previous(
self,
@@ -945,10 +967,15 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
if checkpointer:
self.checkpointer_get_next_version = checkpointer.get_next_version
self.checkpointer_put_writes = checkpointer.aput_writes
self.checkpointer_put_writes_accepts_task_path = (
signature(checkpointer.aput_writes).parameters.get("task_path")
is not None
)
else:
self.checkpointer_get_next_version = increment
self._checkpointer_put_after_previous = None # type: ignore[assignment]
self.checkpointer_put_writes = None
self.checkpointer_put_writes_accepts_task_path = False
async def _checkpointer_put_after_previous(
self,
libs/langgraph/poetry.lock
Generated
+6 -6
View File
@@ -1348,7 +1348,7 @@ typing-extensions = ">=4.7"
[[package]]
name = "langgraph-checkpoint"
version = "2.0.9"
version = "2.0.10"
description = "Library with base interfaces for LangGraph checkpoint savers."
optional = false
python-versions = "^3.9.0,<4.0"
@@ -1366,7 +1366,7 @@ url = "../checkpoint"
[[package]]
name = "langgraph-checkpoint-postgres"
version = "2.0.11"
version = "2.0.12"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
optional = false
python-versions = "^3.9.0,<4.0"
@@ -1375,7 +1375,7 @@ files = []
develop = true
[package.dependencies]
langgraph-checkpoint = "^2.0.7"
langgraph-checkpoint = "^2.0.10"
orjson = ">=3.10.1"
psycopg = "^3.2.0"
psycopg-pool = "^3.2.0"
@@ -1386,7 +1386,7 @@ url = "../checkpoint-postgres"
[[package]]
name = "langgraph-checkpoint-sqlite"
version = "2.0.2"
version = "2.0.3"
description = "Library with a SQLite implementation of LangGraph checkpoint saver."
optional = false
python-versions = "^3.9.0"
@@ -1396,7 +1396,7 @@ develop = true
[package.dependencies]
aiosqlite = "^0.20.0"
langgraph-checkpoint = "^2.0.2"
langgraph-checkpoint = "^2.0.10"
[package.source]
type = "directory"
@@ -3491,4 +3491,4 @@ type = ["pytest-mypy"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9.0,<4.0"
content-hash = "356f7e84cf1375119bd3a118ecf4e9476d95913736f8c426a76d55717eb39161"
content-hash = "caf943b02b6913c05d15c37fda6d216669f789e2a059b7e8e2490b2bdcd23e0e"
libs/langgraph/pyproject.toml
+1 -1
View File
@@ -10,7 +10,7 @@ repository = "https://www.github.com/langchain-ai/langgraph"
[tool.poetry.dependencies]
python = ">=3.9.0,<4.0"
langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.20,!=0.3.21,!=0.3.22"
langgraph-checkpoint = "^2.0.4"
langgraph-checkpoint = "^2.0.10"
langgraph-sdk = "^0.1.42"
[tool.poetry.group.dev.dependencies]
libs/langgraph/tests/test_algo.py
+6 -6
View File
@@ -1,6 +1,6 @@
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.constants import PULL, PUSH
from langgraph.pregel.algo import _tuple_str, prepare_next_tasks
from langgraph.pregel.algo import prepare_next_tasks, task_path_str
from langgraph.pregel.manager import ChannelsManager
@@ -49,16 +49,16 @@ def test_tuple_str() -> None:
push_path_b = (PUSH, push_path_a, 1)
push_path_c = (PUSH, push_path_b, 3)
assert _tuple_str(push_path_a) == f"~{PUSH}, 0000000002"
assert _tuple_str(push_path_b) == f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001"
assert task_path_str(push_path_a) == f"~{PUSH}, 0000000002"
assert task_path_str(push_path_b) == f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001"
assert (
_tuple_str(push_path_c)
task_path_str(push_path_c)
== f"~{PUSH}, ~{PUSH}, ~{PUSH}, 0000000002, 0000000001, 0000000003"
)
assert _tuple_str(pull_path_a) == f"~{PULL}, abc"
assert task_path_str(pull_path_a) == f"~{PULL}, abc"
path_list = [push_path_b, push_path_a, pull_path_a, push_path_c]
assert sorted(map(_tuple_str, path_list)) == [
assert sorted(map(task_path_str, path_list)) == [
f"~{PULL}, abc",
f"~{PUSH}, 0000000002",
f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001",
libs/langgraph/tests/test_large_cases.py
+140 -88
View File
@@ -7255,7 +7255,6 @@ def test_branch_then(
)
@pytest.mark.skip("TODO: re-enable in next PR")
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_send_dedupe_on_resume(
request: pytest.FixtureRequest, checkpointer_name: str
@@ -7314,30 +7313,35 @@ def test_send_dedupe_on_resume(
assert graph.invoke(["0"], thread1, debug=1) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"2|3",
]
assert builder.nodes["2"].runnable.func.ticks == 3
assert builder.nodes["flaky"].runnable.func.ticks == 1
# check state
state = graph.get_state(thread1)
if "shallow" in checkpointer_name:
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) == 2
assert len(history) == 4
# resume execution
assert graph.invoke(None, thread1, debug=1) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"2|3",
"flaky|4",
"3",
"3.1",
]
# node "2" doesn't get called again, as we recover writes saved before
assert builder.nodes["2"].runnable.func.ticks == 3
@@ -7353,12 +7357,13 @@ def test_send_dedupe_on_resume(
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"2|3",
"flaky|4",
"3",
"3.1",
],
next=(),
config={
@@ -7370,35 +7375,33 @@ def test_send_dedupe_on_resume(
},
metadata={
"source": "loop",
"writes": {"3": ["3"], "3.1": ["3.1"]},
"writes": {"3": ["3"]},
"thread_id": "1",
"step": 2,
"step": 4,
"parents": {},
},
created_at=AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
),
},
tasks=(),
),
StateSnapshot(
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"2|3",
"flaky|4",
],
next=("3", "3.1"),
next=("3",),
config={
"configurable": {
"thread_id": "1",
@@ -7408,17 +7411,9 @@ def test_send_dedupe_on_resume(
},
metadata={
"source": "loop",
"writes": {
"1": ["1"],
"2": [
["2|Command(goto=Send(node='2', arg=3))"],
["2|Command(goto=Send(node='flaky', arg=4))"],
["2|3"],
],
"flaky": ["flaky|4"],
},
"writes": {"2": ["2|3"], "3": ["3"], "flaky": ["flaky|4"]},
"thread_id": "1",
"step": 1,
"step": 3,
"parents": {},
},
created_at=AnyStr(),
@@ -7439,6 +7434,123 @@ def test_send_dedupe_on_resume(
state=None,
result=["3"],
),
),
),
StateSnapshot(
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
],
next=("2", "flaky", "3"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"2": [
["2|Command(goto=Send(node='2', arg=3))"],
["2|Command(goto=Send(node='flaky', arg=4))"],
],
"3.1": ["3.1"],
},
"thread_id": "1",
"step": 2,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
result=["2|3"],
),
PregelTask(
id=AnyStr(),
name="flaky",
path=("__pregel_push", 1),
error=None,
interrupts=(
Interrupt(
value="Bahh", resumable=False, ns=None, when="during"
),
),
state=None,
result=["flaky|4"],
),
PregelTask(
id=AnyStr(),
name="3",
path=("__pregel_pull", "3"),
error=None,
interrupts=(),
state=None,
result=["3"],
),
),
),
StateSnapshot(
values=["0", "1"],
next=("2", "2", "3.1"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"1": ["1"]},
"thread_id": "1",
"step": 1,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='2', arg=3))"],
),
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 1),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='flaky', arg=4))"],
),
PregelTask(
id=AnyStr(),
name="3.1",
@@ -7452,7 +7564,7 @@ def test_send_dedupe_on_resume(
),
StateSnapshot(
values=["0"],
next=("1", "2", "2", "2", "flaky"),
next=("1",),
config={
"configurable": {
"thread_id": "1",
@@ -7485,66 +7597,6 @@ def test_send_dedupe_on_resume(
state=None,
result=["1"],
),
PregelTask(
id=AnyStr(),
name="2",
path=(
"__pregel_push",
("__pregel_pull", "1"),
2,
),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='2', arg=3))"],
),
PregelTask(
id=AnyStr(),
name="2",
path=(
"__pregel_push",
("__pregel_pull", "1"),
3,
),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='flaky', arg=4))"],
),
PregelTask(
id=AnyStr(),
name="2",
path=(
"__pregel_push",
(
"__pregel_push",
("__pregel_pull", "1"),
2,
),
2,
),
error=None,
interrupts=(),
state=None,
result=["2|3"],
),
PregelTask(
id=AnyStr(),
name="flaky",
path=(
"__pregel_push",
(
"__pregel_push",
("__pregel_pull", "1"),
3,
),
2,
),
error=None,
interrupts=(Interrupt(value="Bahh", when="during"),),
state=None,
result=["flaky|4"],
),
),
),
StateSnapshot(
libs/langgraph/tests/test_pregel_async.py
+3 -6
View File
@@ -2373,7 +2373,6 @@ async def test_concurrent_emit_sends() -> None:
)
@pytest.mark.skip("TODO: re-enable in next PR")
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_send_sequences(checkpointer_name: str) -> None:
class Node:
@@ -2621,7 +2620,6 @@ async def test_imp_stream_order(checkpointer_name: str) -> None:
]
@pytest.mark.skip("TODO: re-enable in next PR")
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
class InterruptOnce:
@@ -2685,7 +2683,6 @@ 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
print((await graph.aget_state(thread1)).tasks)
# resume execution
assert await graph.ainvoke(None, thread1, debug=1) == [
"0",
@@ -2694,8 +2691,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"flaky|4",
"2|3",
"flaky|4",
"3",
]
# node "2" doesn't get called again, as we recover writes saved before
@@ -2713,8 +2710,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"flaky|4",
"2|3",
"flaky|4",
"3",
],
next=(),
@@ -2750,8 +2747,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"flaky|4",
"2|3",
"flaky|4",
],
next=("3",),
config={