WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-28 10:49:56 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix concurrency issue in PregelScratchpad.consume_null_resume (#3888)

Browse Source 
- Need to use a single operation to check if present and remove item
from list
- This doesn't fix the separate issue that parallel tasks claiming a
single interrupt value have somewhat undefined behavior (in the sense
that they will race to be the first to take it). That will be fixed in a
future PR

Closes #3875
This commit is contained in:
Nuno Campos
2025-03-17 20:22:20 -07:00
committed by GitHub
parent 69dc29aaf9 eae1faa656
commit e24e141253
6 changed files with 34 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/algo.py
+13 -5
View File
@@ -1,4 +1,3 @@
import functools
import itertools
import sys
from collections import defaultdict, deque
@@ -768,6 +767,18 @@ def _scratchpad(
null_resume_write = next(
(w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None
)
def get_null_resume(consume: bool = False) -> Any:
if null_resume_write is None:
return None
if consume:
try:
pending_writes.remove(null_resume_write)
return null_resume_write[2]
except ValueError:
return None
return null_resume_write[2]
# using itertools.count as an atomic counter (+= 1 is not thread-safe)
return PregelScratchpad(
# call
@@ -777,10 +788,7 @@ def _scratchpad(
resume=next(
(w[2] for w in pending_writes if w[0] == task_id and w[1] == RESUME), []
),
null_resume=null_resume_write[2] if null_resume_write is not None else None,
_consume_null_resume=functools.partial(pending_writes.remove, null_resume_write)
if null_resume_write is not None
else lambda: None,
get_null_resume=get_null_resume,
# subgraph
subgraph_counter=itertools.count(0).__next__,
)
libs/langgraph/langgraph/pregel/loop.py
+4 -5
View File
@@ -591,11 +591,10 @@ class PregelLoop(LoopProtocol):
if scratchpad := cast(
Optional[PregelScratchpad], configurable.get(CONFIG_KEY_SCRATCHPAD)
):
if (
isinstance(scratchpad, PregelScratchpad)
and scratchpad.null_resume is not None
):
self.put_writes(NULL_TASK_ID, [(RESUME, scratchpad.null_resume)])
if isinstance(scratchpad, PregelScratchpad):
null_resume = scratchpad.get_null_resume(False)
if null_resume is not None:
self.put_writes(NULL_TASK_ID, [(RESUME, null_resume)])
# map command to writes
if isinstance(self.input, Command):
if self.input.resume is not None and not self.checkpointer:
libs/langgraph/langgraph/types.py
+3 -12
View File
@@ -351,20 +351,11 @@ class PregelScratchpad:
call_counter: Callable[[], int]
# interrupt
interrupt_counter: Callable[[], int]
get_null_resume: Callable[[bool], Any]
resume: list[Any]
null_resume: Optional[Any]
_consume_null_resume: Callable[[], None]
# subgraph
subgraph_counter: Callable[[], int]
def consume_null_resume(self) -> Any:
if self.null_resume is not None:
value = self.null_resume
self._consume_null_resume()
self.null_resume = None
return value
raise ValueError("No null resume to consume")
def interrupt(value: Any) -> Any:
"""Interrupt the graph with a resumable exception from within a node.
@@ -480,9 +471,9 @@ def interrupt(value: Any) -> Any:
if idx < len(scratchpad.resume):
return scratchpad.resume[idx]
# find current resume value
if scratchpad.null_resume is not None:
v = scratchpad.get_null_resume(True)
if v is not None:
assert len(scratchpad.resume) == idx, (scratchpad.resume, idx)
v = scratchpad.consume_null_resume()
scratchpad.resume.append(v)
conf[CONFIG_KEY_SEND]([(RESUME, scratchpad.resume)])
return v
libs/langgraph/tests/test_pregel_async.py
+2 -2
View File
@@ -938,7 +938,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
async for c in tool_two.astream(
{"my_key": "value ⛰️", "market": "DE"}, thread2
)
] == [
] == UnsortedSequence(
{
"__interrupt__": (
Interrupt(
@@ -951,7 +951,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
{
"tool_one": {"my_key": " one"},
},
]
)
# resume with answer
assert [
c async for c in tool_two.astream(Command(resume=" my answer"), thread2)
libs/scheduler-kafka/tests/test_subgraph.py
+6 -6
View File
@@ -202,7 +202,7 @@ async def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": None,
@@ -275,7 +275,7 @@ async def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -378,7 +378,7 @@ async def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -491,7 +491,7 @@ async def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": None,
@@ -559,7 +559,7 @@ async def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -683,7 +683,7 @@ async def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
libs/scheduler-kafka/tests/test_subgraph_sync.py
+6 -6
View File
@@ -201,7 +201,7 @@ def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": None,
@@ -274,7 +274,7 @@ def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -377,7 +377,7 @@ def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -489,7 +489,7 @@ def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": None,
@@ -557,7 +557,7 @@ def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -681,7 +681,7 @@ def test_subgraph_w_interrupt(
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"get_null_resume": None,
"resume": [],
},
"checkpoint_id": c.config["configurable"]["checkpoint_id"],