Make scratchpad counters thread-safe

- Same solution as used in python stdlib to name threads and asyncio tasks
langchain-ai · Jan 23, 2025 · c43a9a4 · c43a9a4
1 parent c961392
commit c43a9a4
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 18 deletions.
diff --git 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
@@ -231,19 +231,20 @@ def __init__(
         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
@@ -143,9 +143,9 @@ def writer(
                     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 @@ def writer(
                     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
@@ -342,14 +342,14 @@ def __init__(
 @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 node(state: State):
     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):