From 020d10138dcda2829c2401fa9b84ce460184d086 Mon Sep 17 00:00:00 2001 From: Chester Curme Date: Wed, 22 Jan 2025 12:39:18 -0500 Subject: [PATCH 1/7] add test case --- libs/langgraph/tests/test_pregel.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 319ceed1c..d71c18306 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5087,6 +5087,23 @@ def test_interrupt_task_functional( res = graph.invoke(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) + def graph(inputs: dict) -> dict: + foo_result = foo(inputs).result() + bar_result = bar(foo_result).result() + baz_result = bar(bar_result).result() + return baz_result + + # First run, interrupted at bar + graph.invoke({"a": ""}, config) + # Provide resumes + graph.invoke(Command(resume="bar"), config) + res_2 = graph.invoke(Command(resume="baz"), config) + assert res_2 == {"a": "foobarbaz"} # Errors, produces "foobarbar" instead + def test_root_mixed_return() -> None: def my_node(state: list[str]): From 3a997be08879baacc0d46f70449a7fdefbc69bf7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 15:40:32 -0800 Subject: [PATCH 2/7] Fix --- libs/langgraph/langgraph/pregel/algo.py | 22 ++++---- libs/langgraph/langgraph/pregel/io.py | 2 +- libs/langgraph/langgraph/pregel/loop.py | 15 +++--- libs/langgraph/langgraph/pregel/runner.py | 16 +++--- libs/langgraph/langgraph/types.py | 39 +++++++++------ libs/langgraph/tests/test_pregel.py | 9 ++-- libs/langgraph/tests/test_pregel_async.py | 16 ++++++ libs/scheduler-kafka/tests/test_subgraph.py | 50 +++---------------- .../tests/test_subgraph_sync.py | 50 +++---------------- 9 files changed, 87 insertions(+), 132 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 38db8ff67..492341581 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -1,3 +1,4 @@ +import functools import sys from collections import defaultdict, deque from functools import partial @@ -46,7 +47,6 @@ from langgraph.constants import ( EMPTY_SEQ, ERROR, INTERRUPT, - MISSING, NO_WRITES, NS_END, NS_SEP, @@ -341,7 +341,7 @@ def prepare_next_tasks( @overload def prepare_next_tasks( checkpoint: Checkpoint, - pending_writes: Sequence[PendingWrite], + pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -357,7 +357,7 @@ def prepare_next_tasks( def prepare_next_tasks( checkpoint: Checkpoint, - pending_writes: Sequence[PendingWrite], + pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -418,7 +418,7 @@ def prepare_single_task( task_id_checksum: Optional[str], *, checkpoint: Checkpoint, - pending_writes: Sequence[PendingWrite], + pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -761,9 +761,13 @@ def prepare_single_task( def _scratchpad( - pending_writes: Sequence[PendingWrite], + pending_writes: list[PendingWrite], task_id: str, ) -> PregelScratchpad: + null_resume_write = next( + (w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None + ) + return PregelScratchpad( # call call_counter=0, @@ -772,10 +776,10 @@ def _scratchpad( resume=next( (w[2] for w in pending_writes if w[0] == task_id and w[1] == RESUME), [] ), - null_resume=next( - (w[2] for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), - MISSING, - ), + 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, # subgraph subgraph_counter=0, ) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 363ff375b..0f064f518 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -89,7 +89,7 @@ def map_command( raise TypeError( f"In Command.goto, expected Send/str, got {type(send).__name__}" ) - if cmd.resume: + if cmd.resume is not None: if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume): for tid, resume in cmd.resume.items(): existing: list[Any] = next( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index b329aa3d5..9496847ca 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -54,7 +54,6 @@ from langgraph.constants import ( ERROR, INPUT, INTERRUPT, - MISSING, NS_SEP, NULL_TASK_ID, PUSH, @@ -230,19 +229,19 @@ class PregelLoop(LoopProtocol): self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD) if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and scratchpad is not None: - if scratchpad["subgraph_counter"]: + if 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(scratchpad.subgraph_counter), ) ) }, ) - scratchpad["subgraph_counter"] += 1 + scratchpad.subgraph_counter += 1 if not self.is_nested and config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): self.config = patch_configurable( self.config, @@ -563,9 +562,11 @@ class PregelLoop(LoopProtocol): ) # take resume value from parent - if scratchpad := configurable.get(CONFIG_KEY_SCRATCHPAD): - if scratchpad["null_resume"] is not MISSING: - self.put_writes(NULL_TASK_ID, [(RESUME, scratchpad["null_resume"])]) + if scratchpad := cast( + Optional[PregelScratchpad], configurable.get(CONFIG_KEY_SCRATCHPAD) + ): + if scratchpad.null_resume is not None: + self.put_writes(NULL_TASK_ID, [(RESUME, scratchpad.null_resume)]) # map command to writes if isinstance(self.input, Command): if self.input.resume is not None and not self.checkpointer: diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 3c899c985..47ab5dfe6 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -39,7 +39,7 @@ from langgraph.errors import GraphBubbleUp, GraphInterrupt from langgraph.pregel.algo import Call from langgraph.pregel.executor import Submit from langgraph.pregel.retry import arun_with_retry, run_with_retry -from langgraph.types import PregelExecutableTask, RetryPolicy +from langgraph.types import PregelExecutableTask, PregelScratchpad, RetryPolicy from langgraph.utils.future import chain_future F = TypeVar("F", concurrent.futures.Future, asyncio.Future) @@ -135,8 +135,7 @@ class PregelRunner: return task.config[CONF][CONFIG_KEY_SEND](writes) # schedule PUSH tasks, collect futures - scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] - scratchpad.setdefault("call_counter", 0) + scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] rtn: dict[int, Optional[concurrent.futures.Future]] = {} for idx, w in enumerate(writes): # bail if not a PUSH write @@ -144,8 +143,8 @@ 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 + cnt = scratchpad.call_counter + scratchpad.call_counter += 1 if next_task := self.schedule_task(task, cnt, wcall): if fut := next( ( @@ -324,8 +323,7 @@ class PregelRunner: return task.config[CONF][CONFIG_KEY_SEND](writes) # schedule PUSH tasks, collect futures - scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] - scratchpad.setdefault("call_counter", 0) + scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] rtn: dict[int, Optional[asyncio.Future]] = {} for idx, w in enumerate(writes): # bail if not a PUSH write @@ -333,8 +331,8 @@ 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 + cnt = scratchpad.call_counter + scratchpad.call_counter += 1 if next_task := self.schedule_task(task, cnt, wcall): # if the parent task was retried, # the next task might already be running diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 09f777d01..806335f7e 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -19,7 +19,7 @@ from typing import ( ) from langchain_core.runnables import Runnable, RunnableConfig -from typing_extensions import Self, TypedDict +from typing_extensions import Self from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata @@ -339,16 +339,26 @@ class LoopProtocol: self.stop = stop -class PregelScratchpad(TypedDict): +@dataclasses.dataclass(**{**_DC_KWARGS, "frozen": False}) +class PregelScratchpad: # call call_counter: int # interrupt interrupt_counter: int resume: list[Any] - null_resume: Any + null_resume: Optional[Any] + _consume_null_resume: Callable[[], None] # subgraph subgraph_counter: 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. @@ -449,7 +459,6 @@ def interrupt(value: Any) -> Any: CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_SEND, - MISSING, NS_SEP, RESUME, ) @@ -459,19 +468,19 @@ 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"] + scratchpad.interrupt_counter += 1 + idx = scratchpad.interrupt_counter # find previous resume values - if scratchpad["resume"]: - if idx < len(scratchpad["resume"]): - return scratchpad["resume"][idx] + if scratchpad.resume: + if idx < len(scratchpad.resume): + return scratchpad.resume[idx] # find current resume value - if scratchpad["null_resume"] is not MISSING: - assert len(scratchpad["resume"]) == idx, (scratchpad["resume"], idx) - v = scratchpad["null_resume"] - scratchpad["null_resume"] = MISSING - scratchpad["resume"].append(v) - conf[CONFIG_KEY_SEND]([(RESUME, scratchpad["resume"])]) + if scratchpad.null_resume is not None: + assert len(scratchpad.resume) == idx, (scratchpad.resume, idx) + print("consume null resume", scratchpad.null_resume) + v = scratchpad.consume_null_resume() + scratchpad.resume.append(v) + conf[CONFIG_KEY_SEND]([(RESUME, scratchpad.resume)]) return v # no resume value found raise GraphInterrupt( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d71c18306..b48ce9d8d 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5082,7 +5082,7 @@ def test_interrupt_task_functional( config = {"configurable": {"thread_id": "1"}} # First run, interrupted at bar - graph.invoke({"a": ""}, config) + assert not graph.invoke({"a": ""}, config) # Resume with an answer res = graph.invoke(Command(resume="bar"), config) assert res == {"a": "foobar"} @@ -5098,11 +5098,10 @@ def test_interrupt_task_functional( return baz_result # First run, interrupted at bar - graph.invoke({"a": ""}, config) + assert not graph.invoke({"a": ""}, config) # Provide resumes - graph.invoke(Command(resume="bar"), config) - res_2 = graph.invoke(Command(resume="baz"), config) - assert res_2 == {"a": "foobarbaz"} # Errors, produces "foobarbar" instead + assert not graph.invoke(Command(resume="bar"), config) + assert graph.invoke(Command(resume="baz"), config) == {"a": "foobarbaz"} def test_root_mixed_return() -> None: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 35d6143dd..166f8720a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6408,6 +6408,22 @@ 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: diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index 2a6c9992a..89a092b83 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 @@ -198,13 +198,7 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": None, "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -271,13 +265,7 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -374,13 +362,7 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -487,13 +469,7 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": None, "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -555,13 +531,7 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -679,13 +649,7 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index c2c9a8fc1..e7e7bdfb0 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 @@ -197,13 +197,7 @@ def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": None, "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -270,13 +264,7 @@ def test_subgraph_w_interrupt( "__pregel_resuming": False, "__pregel_previous": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -373,13 +361,7 @@ def test_subgraph_w_interrupt( "__pregel_resuming": False, "__pregel_task_id": history[0].tasks[0].id, "__pregel_previous": None, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -485,13 +467,7 @@ def test_subgraph_w_interrupt( "__pregel_resuming": True, "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": None, "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -553,13 +529,7 @@ def test_subgraph_w_interrupt( "__pregel_resuming": True, "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -677,13 +647,7 @@ def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, - "null_resume": None, - "resume": [], - }, + "__pregel_scratchpad": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] From 17dacb83a21d5651fea33e710c2e02cc2f602c2f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 15:46:52 -0800 Subject: [PATCH 3/7] Fix --- .../langgraph/scheduler/kafka/orchestrator.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 9ed72cd02..3e4499266 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -13,10 +13,8 @@ from typing_extensions import Self import langgraph.scheduler.kafka.serde as serde from langgraph.constants import ( - CONF, CONFIG_KEY_DEDUPE_TASKS, CONFIG_KEY_ENSURE_LATEST, - CONFIG_KEY_SCRATCHPAD, INTERRUPT, SCHEDULED, ) @@ -178,8 +176,6 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): CONFIG_KEY_ENSURE_LATEST: True, }, ) - if CONFIG_KEY_SCRATCHPAD in config[CONF]: - config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0 # send messages to executor futures = await asyncio.gather( *( @@ -366,8 +362,6 @@ class KafkaOrchestrator(AbstractContextManager): CONFIG_KEY_ENSURE_LATEST: True, }, ) - if CONFIG_KEY_SCRATCHPAD in config[CONF]: - config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0 # send messages to executor futures = [ self.producer.send( From cbad17fa7d4300eca86fffbdc4c4a5a3cbc61950 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 15:51:11 -0800 Subject: [PATCH 4/7] Fix --- libs/langgraph/langgraph/pregel/loop.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 9496847ca..943cfdfc2 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -228,7 +228,9 @@ class PregelLoop(LoopProtocol): if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD) - if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and scratchpad is not None: + if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and isinstance( + scratchpad, PregelScratchpad + ): if scratchpad.subgraph_counter: self.config = patch_configurable( self.config, From 3f2557c9c9993eecc7e7c4c288015e8618cf1ffa Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 15:52:13 -0800 Subject: [PATCH 5/7] Fix --- libs/langgraph/langgraph/pregel/loop.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 943cfdfc2..0559b02ea 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -567,7 +567,10 @@ class PregelLoop(LoopProtocol): if scratchpad := cast( Optional[PregelScratchpad], configurable.get(CONFIG_KEY_SCRATCHPAD) ): - if scratchpad.null_resume is not None: + if ( + isinstance(scratchpad, PregelScratchpad) + and scratchpad.null_resume is not None + ): self.put_writes(NULL_TASK_ID, [(RESUME, scratchpad.null_resume)]) # map command to writes if isinstance(self.input, Command): From c697c2aa04cd5fc0dbd5a6251f5ff5ed0e28480b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 15:57:00 -0800 Subject: [PATCH 6/7] Undo --- libs/scheduler-kafka/tests/test_subgraph.py | 50 ++++++++++++++++--- .../tests/test_subgraph_sync.py | 50 ++++++++++++++++--- 2 files changed, 86 insertions(+), 14 deletions(-) diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index 89a092b83..2a6c9992a 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 +from tests.any import AnyDict, AnyInt from tests.drain import drain_topics_async from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -198,7 +198,13 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": None, "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -265,7 +271,13 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -362,7 +374,13 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -469,7 +487,13 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": None, "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -531,7 +555,13 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -649,7 +679,13 @@ async def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index e7e7bdfb0..c2c9a8fc1 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 +from tests.any import AnyDict, AnyInt from tests.drain import drain_topics from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -197,7 +197,13 @@ def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": None, "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -264,7 +270,13 @@ def test_subgraph_w_interrupt( "__pregel_resuming": False, "__pregel_previous": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -361,7 +373,13 @@ def test_subgraph_w_interrupt( "__pregel_resuming": False, "__pregel_task_id": history[0].tasks[0].id, "__pregel_previous": None, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -467,7 +485,13 @@ def test_subgraph_w_interrupt( "__pregel_resuming": True, "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": None, "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -529,7 +553,13 @@ def test_subgraph_w_interrupt( "__pregel_resuming": True, "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -647,7 +677,13 @@ def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] From c9613927dc84b5a363dac45ece56466a1641b408 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 16:05:34 -0800 Subject: [PATCH 7/7] Lint --- libs/langgraph/langgraph/pregel/algo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 492341581..eff62e0fb 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -324,7 +324,7 @@ def apply_writes( @overload def prepare_next_tasks( checkpoint: Checkpoint, - pending_writes: Sequence[PendingWrite], + pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping,