From c43a9a4bd0fcbdd33b0c97cf2aee5ddddbb95e8a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 16:07:59 -0800 Subject: [PATCH 1/5] Make scratchpad counters thread-safe - Same solution as used in python stdlib to name threads and asyncio tasks --- libs/langgraph/langgraph/pregel/algo.py | 9 +++++---- libs/langgraph/langgraph/pregel/loop.py | 7 ++++--- libs/langgraph/langgraph/pregel/runner.py | 12 ++++++------ libs/langgraph/langgraph/types.py | 9 ++++----- 4 files changed, 19 insertions(+), 18 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index eff62e0fb..eac0f1fdd 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -1,4 +1,5 @@ import functools +import itertools import sys from collections import defaultdict, deque from functools import partial @@ -767,12 +768,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), [] ), @@ -781,7 +782,7 @@ def _scratchpad( if null_resume_write is not None else lambda: None, # subgraph - subgraph_counter=0, + subgraph_counter=itertools.count(0).__next__, ) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 0559b02ea..b360a7440 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -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, diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 47ab5dfe6..6336bc5a1 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -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( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 806335f7e..bc6d6b645 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -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): From 23c73ae7195cde40977c431e39da60f9ad59b2ff Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 16:30:43 -0800 Subject: [PATCH 2/5] Update tests --- libs/langgraph/tests/test_pregel_async.py | 6 ++- libs/scheduler-kafka/tests/test_subgraph.py | 38 +++++++++---------- .../tests/test_subgraph_sync.py | 38 +++++++++---------- 3 files changed, 43 insertions(+), 39 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 166f8720a..fa1156859 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1612,7 +1612,11 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) + async def add_one_impl(x: int) -> int: + await asyncio.sleep(0.01 * x) + return x + 1 + + add_one = mocker.Mock(side_effect=add_one_impl) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = ( Channel.subscribe_to("inbox") diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index 2a6c9992a..934e3a614 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -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": [], }, diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index c2c9a8fc1..50b33a39d 100644 --- a/libs/scheduler-kafka/tests/test_subgraph_sync.py +++ b/libs/scheduler-kafka/tests/test_subgraph_sync.py @@ -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": [], }, From 44bf97ac0e58f633f064d378a4d0a391790a9fcc Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 16:32:59 -0800 Subject: [PATCH 3/5] Fix --- libs/langgraph/tests/test_pregel_async.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index fa1156859..77319fee7 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1616,7 +1616,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: await asyncio.sleep(0.01 * x) return x + 1 - add_one = mocker.Mock(side_effect=add_one_impl) + add_one = mocker.AsyncMock(side_effect=add_one_impl) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = ( Channel.subscribe_to("inbox") From 211fd4337da644c820f89949b9d5df9a49a9b4e7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 16:44:13 -0800 Subject: [PATCH 4/5] Undo --- libs/langgraph/tests/test_pregel_async.py | 22 +--------------------- 1 file changed, 1 insertion(+), 21 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 77319fee7..35d6143dd 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1612,11 +1612,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: - async def add_one_impl(x: int) -> int: - await asyncio.sleep(0.01 * x) - return x + 1 - - add_one = mocker.AsyncMock(side_effect=add_one_impl) + add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = ( Channel.subscribe_to("inbox") @@ -6412,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: From 8a4c4523177407d3e9c89c8987037b236908beac Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 23 Jan 2025 07:27:34 -0800 Subject: [PATCH 5/5] Make test less flaky --- libs/langgraph/tests/test_pregel.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index b48ce9d8d..fb69905c2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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]: