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

Make scratchpad counters thread-safe (#3158)

Browse Source 
- Same solution as used in python stdlib to name threads and asyncio
tasks
This commit is contained in:
Nuno Campos
2025-01-23 07:37:04 -08:00
committed by GitHub
parent 5dc5853161 8a4c452317
commit 6335963674
8 changed files with 60 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/algo.py
+5 -4
View File
@@ -1,4 +1,5 @@
import functools
import itertools
import sys
from collections import defaultdict, deque
from functools import partial
@@ -765,12 +766,12 @@ def _scratchpad(
null_resume_write = next(
(w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None
)
# using itertools.count as an atomic counter (+= 1 is not thread-safe)
return PregelScratchpad(
# call
call_counter=0,
call_counter=itertools.count(0).__next__,
# interrupt
interrupt_counter=-1,
interrupt_counter=itertools.count(0).__next__,
resume=next(
(w[2] for w in pending_writes if w[0] == task_id and w[1] == RESUME), []
),
@@ -779,7 +780,7 @@ def _scratchpad(
if null_resume_write is not None
else lambda: None,
# subgraph
subgraph_counter=0,
subgraph_counter=itertools.count(0).__next__,
)
libs/langgraph/langgraph/pregel/loop.py
+4 -3
View File
@@ -231,19 +231,20 @@ class PregelLoop(LoopProtocol):
if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and isinstance(
scratchpad, PregelScratchpad
):
if scratchpad.subgraph_counter:
# if count is > 0, append to checkpoint_ns
# if count is 0, leave as is
if cnt := scratchpad.subgraph_counter():
self.config = patch_configurable(
self.config,
{
CONFIG_KEY_CHECKPOINT_NS: NS_SEP.join(
(
config[CONF][CONFIG_KEY_CHECKPOINT_NS],
str(scratchpad.subgraph_counter),
str(cnt),
)
)
},
)
scratchpad.subgraph_counter += 1
if not self.is_nested and config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
self.config = patch_configurable(
self.config,
libs/langgraph/langgraph/pregel/runner.py
+6 -6
View File
@@ -143,9 +143,9 @@ class PregelRunner:
continue
# schedule the next task, if the callback returns one
wcall = calls[idx] if calls else None
cnt = scratchpad.call_counter
scratchpad.call_counter += 1
if next_task := self.schedule_task(task, cnt, wcall):
if next_task := self.schedule_task(
task, scratchpad.call_counter(), wcall
):
if fut := next(
(
f
@@ -331,9 +331,9 @@ class PregelRunner:
continue
# schedule the next task, if the callback returns one
wcall = calls[idx] if calls is not None else None
cnt = scratchpad.call_counter
scratchpad.call_counter += 1
if next_task := self.schedule_task(task, cnt, wcall):
if next_task := self.schedule_task(
task, scratchpad.call_counter(), wcall
):
# if the parent task was retried,
# the next task might already be running
if fut := next(
libs/langgraph/langgraph/types.py
+4 -5
View File
@@ -342,14 +342,14 @@ class LoopProtocol:
@dataclasses.dataclass(**{**_DC_KWARGS, "frozen": False})
class PregelScratchpad:
# call
call_counter: int
call_counter: Callable[[], int]
# interrupt
interrupt_counter: int
interrupt_counter: Callable[[], int]
resume: list[Any]
null_resume: Optional[Any]
_consume_null_resume: Callable[[], None]
# subgraph
subgraph_counter: int
subgraph_counter: Callable[[], int]
def consume_null_resume(self) -> Any:
if self.null_resume is not None:
@@ -468,8 +468,7 @@ def interrupt(value: Any) -> Any:
conf = get_config()["configurable"]
# track interrupt index
scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD]
scratchpad.interrupt_counter += 1
idx = scratchpad.interrupt_counter
idx = scratchpad.interrupt_counter()
# find previous resume values
if scratchpad.resume:
if idx < len(scratchpad.resume):
libs/langgraph/tests/test_pregel.py
+3 -1
View File
@@ -1539,12 +1539,14 @@ def test_imp_nested(
@task
def submapper(input: int) -> str:
time.sleep(input / 100)
return str(input)
@task()
def mapper(input: int) -> str:
sub = submapper(input)
time.sleep(input / 100)
return submapper(input).result() * 2
return sub.result() * 2
@entrypoint(checkpointer=checkpointer)
def graph(input: list[int]) -> list[str]:
libs/langgraph/tests/test_pregel_async.py
-16
View File
@@ -6408,22 +6408,6 @@ async def test_interrupt_task_functional(checkpointer_name: str) -> None:
res = await graph.ainvoke(Command(resume="bar"), config)
assert res == {"a": "foobar"}
# Test that we can interrupt the same task multiple times
config = {"configurable": {"thread_id": "2"}}
@entrypoint(checkpointer=checkpointer)
async def graph(inputs: dict) -> dict:
foo_result = await foo(inputs)
bar_result = await bar(foo_result)
baz_result = await bar(bar_result)
return baz_result
# First run, interrupted at bar
assert not await graph.ainvoke({"a": ""}, config)
# Provide resumes
assert not await graph.ainvoke(Command(resume="bar"), config)
assert await graph.ainvoke(Command(resume="baz"), config) == {"a": "foobarbaz"}
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_command_with_static_breakpoints(checkpointer_name: str) -> None:
libs/scheduler-kafka/tests/test_subgraph.py
+19 -19
View File
@@ -15,7 +15,7 @@ from langgraph.graph.state import StateGraph
from langgraph.pregel import Pregel
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from tests.any import AnyDict, AnyInt
from tests.any import AnyDict
from tests.drain import drain_topics_async
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -199,9 +199,9 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
@@ -272,9 +272,9 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
@@ -375,9 +375,9 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
@@ -488,9 +488,9 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
@@ -556,9 +556,9 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
@@ -680,9 +680,9 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
libs/scheduler-kafka/tests/test_subgraph_sync.py
+19 -19
View File
@@ -15,7 +15,7 @@ from langgraph.pregel import Pregel
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.default_sync import DefaultProducer
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from tests.any import AnyDict, AnyInt
from tests.any import AnyDict
from tests.drain import drain_topics
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -198,9 +198,9 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
@@ -271,9 +271,9 @@ def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
@@ -374,9 +374,9 @@ def test_subgraph_w_interrupt(
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_previous": None,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
@@ -486,9 +486,9 @@ def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
@@ -554,9 +554,9 @@ def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},
@@ -678,9 +678,9 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"subgraph_counter": AnyInt(),
"call_counter": 0,
"interrupt_counter": -1,
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
"null_resume": None,
"resume": [],
},