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):