From 7d8205633d32f05d544f517e63c0c7754ccc75cd Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 8 Nov 2024 13:05:45 -0800 Subject: [PATCH 01/26] Add `call` function to call a node and get a future - Whereas Send is for fire-and-forget type of calls, new `call` and `acall` functions are for flows where you want to wait for the node to finish before doing something else - Because we return regular python future objects (concurrent.futures.Future or asyncio.Future) all the python primitives for working with futures work, eg. wait, gather, etc --- libs/langgraph/langgraph/constants.py | 2 + libs/langgraph/langgraph/pregel/__init__.py | 5 + libs/langgraph/langgraph/pregel/algo.py | 40 ++++++-- libs/langgraph/langgraph/pregel/loop.py | 8 +- libs/langgraph/langgraph/pregel/retry.py | 20 ++-- libs/langgraph/langgraph/pregel/runner.py | 108 ++++++++++++++------ libs/langgraph/langgraph/types.py | 38 +++++++ libs/langgraph/tests/test_pregel.py | 48 ++++----- libs/langgraph/tests/test_pregel_async.py | 74 ++++---------- 9 files changed, 197 insertions(+), 146 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index e2d9f069a..b9d6f37f7 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -46,6 +46,8 @@ CONFIG_KEY_SEND = sys.intern("__pregel_send") # holds the `write` function that accepts writes to state/edges/reserved keys CONFIG_KEY_READ = sys.intern("__pregel_read") # holds the `read` function that returns a copy of the current state +CONFIG_KEY_CALL = sys.intern("__pregel_call") +# holds the `call` function that accepts a node/func, args and returns a future CONFIG_KEY_CHECKPOINTER = sys.intern("__pregel_checkpointer") # holds a `BaseCheckpointSaver` passed from parent graph to child graphs CONFIG_KEY_STREAM = sys.intern("__pregel_stream") diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index e714afe21..52cd2bf09 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -541,6 +541,11 @@ class Pregel(PregelProtocol): next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: apply_writes(saved.checkpoint, channels, tasks, None) + print( + "next_tasks", + [(t.id, t.name, t.input, bool(t.writes)) for t in next_tasks.values()], + ) + print("pending_writes", saved.pending_writes) # assemble the state snapshot return StateSnapshot( read_channels(channels, self.stream_channels_asis), diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 0885f12aa..a8d39674c 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -97,6 +97,11 @@ class PregelTaskWrites(NamedTuple): triggers: Sequence[str] +class Call: + func: str | Callable + input: Any + + def should_interrupt( checkpoint: Checkpoint, interrupt_nodes: Union[All, Sequence[str]], @@ -490,10 +495,16 @@ def prepare_single_task( PUSH, str(idx), ) - elif len(task_path) == 4: + elif len(task_path) >= 4: # new PUSH tasks, executed in superstep n # (PUSH, parent task path, idx of PUSH write, id of parent task) - task_path_t = cast(tuple[str, tuple, int, str], task_path) + task_path_t = cast( + Union[ + tuple[str, tuple, int, str], + tuple[str, tuple, int, str, Optional[Call]], + ], + task_path, + ) writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]] if task_path_t[2] >= len(writes_for_path): logger.warning( @@ -501,6 +512,13 @@ def prepare_single_task( ) return packet = writes_for_path[task_path_t[2]][2] + if packet is None: + if len(task_path_t) == 5: + packet = task_path_t[4] + else: + # no packet to replay, this is a "call" task + return + # TODO handle Call packets if not isinstance(packet, Send): logger.warning( f"Ignoring invalid packet type {type(packet)} in pending writes" @@ -533,7 +551,7 @@ def prepare_single_task( "langgraph_step": step, "langgraph_node": packet.node, "langgraph_triggers": triggers, - "langgraph_path": task_path, + "langgraph_path": task_path[:3], "langgraph_checkpoint_ns": task_checkpoint_ns, } if task_id_checksum is not None: @@ -572,7 +590,7 @@ def prepare_single_task( channels, managed, PregelTaskWrites( - task_path, packet.node, writes, triggers + task_path[:3], packet.node, writes, triggers ), config, ), @@ -602,12 +620,12 @@ def prepare_single_task( proc.retry_policy, None, task_id, - task_path, + task_path[:3], writers=proc.flat_writers, ) else: - return PregelTask(task_id, packet.node, task_path) + return PregelTask(task_id, packet.node, task_path[:3]) elif task_path[0] == PULL: # (PULL, node name) name = cast(str, task_path[1]) @@ -657,7 +675,7 @@ def prepare_single_task( "langgraph_step": step, "langgraph_node": name, "langgraph_triggers": triggers, - "langgraph_path": task_path, + "langgraph_path": task_path[:3], "langgraph_checkpoint_ns": task_checkpoint_ns, } if task_id_checksum is not None: @@ -696,7 +714,9 @@ def prepare_single_task( checkpoint, channels, managed, - PregelTaskWrites(task_path, name, writes, triggers), + PregelTaskWrites( + task_path[:3], name, writes, triggers + ), config, ), CONFIG_KEY_STORE: ( @@ -725,11 +745,11 @@ def prepare_single_task( proc.retry_policy, None, task_id, - task_path, + task_path[:3], writers=proc.flat_writers, ) else: - return PregelTask(task_id, name, task_path) + return PregelTask(task_id, name, task_path[:3]) def _proc_input( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index d9af9279e..31397d22b 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -73,6 +73,7 @@ from langgraph.managed.base import ( WritableManagedValue, ) from langgraph.pregel.algo import ( + Call, GetNextVersion, PregelTaskWrites, apply_writes, @@ -307,12 +308,9 @@ class PregelLoop(LoopProtocol): self._output_writes(task_id, writes) def accept_push( - self, task: PregelExecutableTask, write_idx: int + self, task: PregelExecutableTask, write_idx: int, call: Optional[Call] = None ) -> Optional[PregelExecutableTask]: """Accept a PUSH from a task, potentially returning a new task to start.""" - # don't start if an earlier PUSH has already triggered an interrupt - if self.to_interrupt: - return # don't start if we should interrupt *after* the original task if should_interrupt(self.checkpoint, self.interrupt_after, [task]): self.to_interrupt.append(task) @@ -320,7 +318,7 @@ class PregelLoop(LoopProtocol): if pushed := cast( Optional[PregelExecutableTask], prepare_single_task( - (PUSH, task.path, write_idx, task.id), + (PUSH, task.path, write_idx, task.id, call), None, checkpoint=self.checkpoint, pending_writes=[(task.id, *w) for w in task.writes], diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 2d0f2b6da..093859579 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -4,14 +4,12 @@ import random import sys import time from dataclasses import replace -from functools import partial -from typing import Any, Callable, Optional, Sequence +from typing import Any, Optional, Sequence from langgraph.constants import ( CONF, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_RESUMING, - CONFIG_KEY_SEND, NS_SEP, ) from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand @@ -25,17 +23,15 @@ SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) def run_with_retry( task: PregelExecutableTask, retry_policy: Optional[RetryPolicy], - writer: Optional[ - Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None] - ] = None, + configurable: Optional[dict[str, Any]] = None, ) -> None: """Run a task with retries.""" retry_policy = task.retry_policy or retry_policy interval = retry_policy.initial_interval if retry_policy else 0 attempts = 0 config = task.config - if writer is not None: - config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)}) + if configurable is not None: + config = patch_configurable(config, configurable) while True: try: # clear any writes from previous attempts @@ -115,17 +111,15 @@ async def arun_with_retry( task: PregelExecutableTask, retry_policy: Optional[RetryPolicy], stream: bool = False, - writer: Optional[ - Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None] - ] = None, + configurable: Optional[dict[str, Any]] = None, ) -> None: """Run a task asynchronously with retries.""" retry_policy = task.retry_policy or retry_policy interval = retry_policy.initial_interval if retry_policy else 0 attempts = 0 config = task.config - if writer is not None: - config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)}) + if configurable is not None: + config = patch_configurable(config, configurable) while True: try: # clear any writes from previous attempts diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index f46210459..067394bdb 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -1,9 +1,11 @@ import asyncio import concurrent.futures import time +from functools import partial from typing import ( Any, AsyncIterator, + Awaitable, Callable, Iterable, Iterator, @@ -16,6 +18,7 @@ from typing import ( from langgraph.constants import ( CONF, + CONFIG_KEY_CALL, CONFIG_KEY_SEND, ERROR, INTERRUPT, @@ -25,6 +28,7 @@ from langgraph.constants import ( TAG_HIDDEN, ) 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 @@ -41,7 +45,7 @@ class PregelRunner: submit: Submit, put_writes: Callable[[str, Sequence[tuple[str, Any]]], None], schedule_task: Callable[ - [PregelExecutableTask, int], Optional[PregelExecutableTask] + [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] ], use_astream: bool = False, node_finished: Optional[Callable[[str], None]] = None, @@ -67,11 +71,10 @@ class PregelRunner: prev_length = len(task.writes) # delegate to the underlying writer task.config[CONF][CONFIG_KEY_SEND](writes) - for idx, w in enumerate(task.writes): - # find the index for the newly inserted writes - if idx < prev_length: - continue - assert writes[idx - prev_length] is w + # confirm no other concurrent writes were added + # TODO could use a lock here instead, if writes can come from many threads + assert len(task.writes) == prev_length + len(writes) + for idx, w in enumerate(writes, start=prev_length): # bail if not a PUSH write if w[0] != PUSH: continue @@ -89,7 +92,10 @@ class PregelRunner: run_with_retry, next_task, retry_policy, - writer=writer, + configurable={ + CONFIG_KEY_SEND: partial(writer, next_task), + # CONFIG_KEY_CALL: partial(call, next_task), + }, __reraise_on_exit__=reraise, ) ] = next_task @@ -102,7 +108,14 @@ class PregelRunner: if len(tasks) == 1 and timeout is None and get_waiter is None: t = tasks[0] try: - run_with_retry(t, retry_policy, writer=writer) + run_with_retry( + t, + retry_policy, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + # CONFIG_KEY_CALL: partial(call, t), + }, + ) self.commit(t, None) except Exception as exc: self.commit(t, exc) @@ -123,7 +136,10 @@ class PregelRunner: run_with_retry, t, retry_policy, - writer=writer, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + # CONFIG_KEY_CALL: partial(call, t), + }, __reraise_on_exit__=reraise, ) ] = t @@ -172,21 +188,28 @@ class PregelRunner: get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, ) -> AsyncIterator[None]: def writer( - task: PregelExecutableTask, writes: Sequence[tuple[str, Any]] - ) -> None: + task: PregelExecutableTask, + writes: Sequence[tuple[str, Any]], + *, + calls: Optional[Sequence[Call]] = None, + ) -> Sequence[Optional[asyncio.Future]]: prev_length = len(task.writes) # delegate to the underlying writer task.config[CONF][CONFIG_KEY_SEND](writes) - for idx, w in enumerate(task.writes): - # find the index for the newly inserted writes - if idx < prev_length: - continue - assert writes[idx - prev_length] is w + # confirm no other concurrent writes were added + # TODO could use a lock here instead, if writes can come from many threads + assert len(task.writes) == prev_length + len(writes) + rtn: dict[int, Optional[asyncio.Future]] = {} + for idx, w in enumerate(writes, start=prev_length): # bail if not a PUSH write if w[0] != PUSH: continue # schedule the next task, if the callback returns one - if next_task := self.schedule_task(task, idx): + if next_task := self.schedule_task( + task, + idx, + calls[idx] if calls is not None else None, + ): # if the parent task was retried, # the next task might already be running if any( @@ -194,21 +217,29 @@ class PregelRunner: ): continue # schedule the next task - futures[ - cast( - asyncio.Future, - self.submit( - arun_with_retry, - next_task, - retry_policy, - stream=self.use_astream, - writer=writer, - __name__=t.name, - __cancel_on_exit__=True, - __reraise_on_exit__=reraise, - ), - ) - ] = next_task + fut = self.submit( + arun_with_retry, + next_task, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_SEND: partial(writer, next_task), + CONFIG_KEY_CALL: partial(call, next_task), + }, + __name__=t.name, + __cancel_on_exit__=True, + __reraise_on_exit__=reraise, + ) + futures[cast(asyncio.Future, fut)] = next_task + rtn[idx] = fut + return [rtn.get(i) for i in range(len(writes))] + + def call( + task, func: str | Callable[[Any], Union[Awaitable[Any], Any]], input: Any + ) -> asyncio.Future[Any]: + (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)]) + assert fut is not None, "writer did not return a future for call" + return fut loop = asyncio.get_event_loop() tasks = tuple(tasks) @@ -220,7 +251,13 @@ class PregelRunner: t = tasks[0] try: await arun_with_retry( - t, retry_policy, stream=self.use_astream, writer=writer + t, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + CONFIG_KEY_CALL: partial(call, t), + }, ) self.commit(t, None) except Exception as exc: @@ -245,7 +282,10 @@ class PregelRunner: t, retry_policy, stream=self.use_astream, - writer=writer, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + CONFIG_KEY_CALL: partial(call, t), + }, __name__=t.name, __cancel_on_exit__=True, __reraise_on_exit__=reraise, diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 67c7e53f8..b49b1c45e 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,9 +1,13 @@ +import asyncio +import concurrent +import concurrent.futures import dataclasses import sys from collections import deque from typing import ( TYPE_CHECKING, Any, + Awaitable, Callable, ClassVar, Generic, @@ -11,6 +15,7 @@ from typing import ( Literal, NamedTuple, Optional, + ParamSpec, Sequence, Type, TypedDict, @@ -363,3 +368,36 @@ def interrupt(value: Any) -> Any: ), ) ) + + +P = ParamSpec("P") +T = TypeVar("T") + + +def call( + func: str | Callable[P, T], *args: P.args, **kwargs: P.kwargs +) -> concurrent.futures.Future[T]: + from langgraph.constants import CONFIG_KEY_CALL + from langgraph.utils.config import get_configurable + + conf = get_configurable() + impl = conf[CONFIG_KEY_CALL] + fut = impl(func, *args, **kwargs) + if not isinstance(fut, concurrent.futures.Future): + raise RuntimeError("In an async context, use acall() instead of call()") + return fut + + +def acall( + func: str | Callable[P, Union[T, Awaitable[T]]], *args: P.args, **kwargs: P.kwargs +) -> asyncio.Future[T]: + from langgraph.constants import CONFIG_KEY_CALL + from langgraph.utils.config import get_configurable + + conf = get_configurable() + impl = conf[CONFIG_KEY_CALL] + fut = impl(func, *args, **kwargs) + if isinstance(fut, concurrent.futures.Future): + fut = asyncio.wrap_future(fut) + fut.cancel + return fut diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f69d36ed3..120b3deac 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -2496,7 +2496,7 @@ def test_send_react_interrupt( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -2653,7 +2653,7 @@ def test_send_react_interrupt( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -2740,7 +2740,7 @@ def test_send_react_interrupt( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", (), 0, AnyStr()), + path=("__pregel_push", (), 0), error=None, interrupts=(), state=None, @@ -2965,7 +2965,7 @@ def test_send_react_interrupt_control( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -6043,9 +6043,7 @@ def test_state_graph_packets( ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), ), next=("tools",), config=(app_w_interrupt.checkpointer.get_tuple(config)).config, @@ -6085,7 +6083,7 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], @@ -6214,12 +6212,8 @@ def test_state_graph_packets( ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), ), next=("tools", "tools"), config=app_w_interrupt.checkpointer.get_tuple(config).config, @@ -6366,9 +6360,7 @@ def test_state_graph_packets( ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), ), next=("tools",), config=(app_w_interrupt.checkpointer.get_tuple(config)).config, @@ -6408,7 +6400,7 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], @@ -6537,12 +6529,8 @@ def test_state_graph_packets( ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), ), next=("tools", "tools"), config=app_w_interrupt.checkpointer.get_tuple(config).config, @@ -12786,7 +12774,7 @@ def test_send_to_nested_graphs( PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 1), state={ "configurable": { "thread_id": "1", @@ -12797,7 +12785,7 @@ def test_send_to_nested_graphs( PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 2), state={ "configurable": { "thread_id": "1", @@ -12850,7 +12838,7 @@ def test_send_to_nested_graphs( "checkpoint_ns": AnyStr("generate_joke:"), "langgraph_checkpoint_ns": AnyStr("generate_joke:"), "langgraph_node": "generate_joke", - "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1, AnyStr()], + "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1], "langgraph_step": 0, "langgraph_triggers": [PUSH], }, @@ -12895,7 +12883,7 @@ def test_send_to_nested_graphs( "checkpoint_ns": AnyStr("generate_joke:"), "langgraph_checkpoint_ns": AnyStr("generate_joke:"), "langgraph_node": "generate_joke", - "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2, AnyStr()], + "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2], "langgraph_step": 0, "langgraph_triggers": [PUSH], }, @@ -13021,7 +13009,7 @@ def test_send_to_nested_graphs( PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 1), state={ "configurable": { "thread_id": "1", @@ -13033,7 +13021,7 @@ def test_send_to_nested_graphs( PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 2), state={ "configurable": { "thread_id": "1", diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 514703781..d8aeedd9a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2864,12 +2864,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="2", - path=( - "__pregel_push", - ("__pregel_pull", "1"), - 2, - AnyStr(), - ), + path=("__pregel_push", ("__pregel_pull", "1"), 2), error=None, interrupts=(), state=None, @@ -2878,12 +2873,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="2", - path=( - "__pregel_push", - ("__pregel_pull", "1"), - 3, - AnyStr(), - ), + path=("__pregel_push", ("__pregel_pull", "1"), 3), error=None, interrupts=(), state=None, @@ -2894,14 +2884,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: name="2", path=( "__pregel_push", - ( - "__pregel_push", - ("__pregel_pull", "1"), - 2, - AnyStr(), - ), + ("__pregel_push", ("__pregel_pull", "1"), 2), 2, - AnyStr(), ), error=None, interrupts=(), @@ -2913,14 +2897,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: name="flaky", path=( "__pregel_push", - ( - "__pregel_push", - ("__pregel_pull", "1"), - 3, - AnyStr(), - ), + ("__pregel_push", ("__pregel_pull", "1"), 3), 2, - AnyStr(), ), error=None, interrupts=(Interrupt(value="Bahh", when="during"),), @@ -3157,7 +3135,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -3314,7 +3292,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -3401,7 +3379,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", (), 0, AnyStr()), + path=("__pregel_push", (), 0), error=None, interrupts=(), state=None, @@ -3625,7 +3603,7 @@ async def test_send_react_interrupt_control( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -6420,9 +6398,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, @@ -6465,7 +6441,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), next=("tools",), config=tup.config, created_at=tup.checkpoint["ts"], @@ -6596,12 +6572,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), ), next=("tools", "tools"), config=tup.config, @@ -6751,9 +6723,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, @@ -6796,7 +6766,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), next=("tools",), config=tup.config, created_at=tup.checkpoint["ts"], @@ -6929,12 +6899,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), ), next=("tools", "tools"), config=tup.config, @@ -11612,7 +11578,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 1), state={ "configurable": { "thread_id": "1", @@ -11623,7 +11589,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 2), state={ "configurable": { "thread_id": "1", @@ -11764,7 +11730,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 1), state={ "configurable": { "thread_id": "1", @@ -11776,7 +11742,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 2), state={ "configurable": { "thread_id": "1", From 0461d45d76d4af9f1bd214a8b77792d88e59fa2b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 8 Nov 2024 16:33:07 -0800 Subject: [PATCH 02/26] Finish impl --- libs/langgraph/langgraph/pregel/algo.py | 107 +++++++++++++++++--- libs/langgraph/langgraph/pregel/call.py | 113 ++++++++++++++++++++++ libs/langgraph/langgraph/pregel/retry.py | 10 +- libs/langgraph/langgraph/pregel/runner.py | 52 ++++++---- 4 files changed, 241 insertions(+), 41 deletions(-) create mode 100644 libs/langgraph/langgraph/pregel/call.py diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index a8d39674c..c95c699c8 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -58,6 +58,7 @@ from langgraph.constants import ( ) from langgraph.errors import EmptyChannelError, InvalidUpdateError from langgraph.managed.base import ManagedValueMapping +from langgraph.pregel.call import get_runnable_for_func from langgraph.pregel.io import read_channel, read_channels from langgraph.pregel.log import logger from langgraph.pregel.manager import ChannelsManager @@ -98,9 +99,15 @@ class PregelTaskWrites(NamedTuple): class Call: + __slots__ = ("func", "input") + func: str | Callable input: Any + def __init__(self, func: str | Callable, input: Any) -> None: + self.func = func + self.input = input + def should_interrupt( checkpoint: Checkpoint, @@ -184,7 +191,7 @@ def local_write( """Function injected under CONFIG_KEY_SEND in task config, to write to channels. Validates writes and forwards them to `commit` function.""" for chan, value in writes: - if chan in (PUSH, TASKS): + if chan in (PUSH, TASKS) and value is not None: if not isinstance(value, Send): raise InvalidUpdateError(f"Expected Send, got {value}") if value.node not in process_keys: @@ -464,7 +471,87 @@ def prepare_single_task( configurable = config.get(CONF, {}) parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "") - if task_path[0] == PUSH: + if task_path[0] == PUSH and isinstance(task_path[-1], Call): + # (PUSH, parent task path, idx of PUSH write, id of parent task, Call) + task_path_t = cast(tuple[str, tuple, int, str, Optional[Call]], task_path) + call = task_path_t[-1] + proc = get_runnable_for_func(call.func) + name = proc.name + if name is None: + raise ValueError("`call` functions must have a `__name__` attribute") + # create task id + triggers = [PUSH] + checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name + task_id = _uuid5_str( + checkpoint_id, + checkpoint_ns, + str(step), + name, + PUSH, + _tuple_str(task_path[1]), + str(task_path[2]), + ) + task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" + metadata = { + "langgraph_step": step, + "langgraph_node": name, + "langgraph_triggers": triggers, + "langgraph_path": task_path[:3], + "langgraph_checkpoint_ns": task_checkpoint_ns, + } + if task_id_checksum is not None: + assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}" + if for_execution: + writes: deque[tuple[str, Any]] = deque() + return PregelExecutableTask( + name, + call.input, + proc, + writes, + patch_config( + merge_configs(config, {"metadata": metadata, "tags": proc.tags}), + run_name=name, + callbacks=( + manager.get_child(f"graph:step:{step}") if manager else None + ), + configurable={ + CONFIG_KEY_TASK_ID: task_id, + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + local_write, + writes.extend, + processes.keys(), + ), + CONFIG_KEY_READ: partial( + local_read, + step, + checkpoint, + channels, + managed, + PregelTaskWrites(task_path[:3], name, writes, triggers), + config, + ), + CONFIG_KEY_STORE: (store or configurable.get(CONFIG_KEY_STORE)), + CONFIG_KEY_CHECKPOINTER: ( + checkpointer or configurable.get(CONFIG_KEY_CHECKPOINTER) + ), + CONFIG_KEY_CHECKPOINT_MAP: { + **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}), + parent_ns: checkpoint["id"], + }, + CONFIG_KEY_CHECKPOINT_ID: None, + CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, + }, + ), + triggers, + None, + None, + task_id, + task_path[:3], + ) + else: + return PregelTask(task_id, name, task_path[:3]) + elif task_path[0] == PUSH: if len(task_path) == 2: # TODO: remove branch in 1.0 # legacy SEND tasks, executed in superstep n+1 # (PUSH, idx of pending send) @@ -498,13 +585,7 @@ def prepare_single_task( elif len(task_path) >= 4: # new PUSH tasks, executed in superstep n # (PUSH, parent task path, idx of PUSH write, id of parent task) - task_path_t = cast( - Union[ - tuple[str, tuple, int, str], - tuple[str, tuple, int, str, Optional[Call]], - ], - task_path, - ) + task_path_t = cast(tuple[str, tuple, int, str], task_path) writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]] if task_path_t[2] >= len(writes_for_path): logger.warning( @@ -513,12 +594,7 @@ def prepare_single_task( return packet = writes_for_path[task_path_t[2]][2] if packet is None: - if len(task_path_t) == 5: - packet = task_path_t[4] - else: - # no packet to replay, this is a "call" task - return - # TODO handle Call packets + return if not isinstance(packet, Send): logger.warning( f"Ignoring invalid packet type {type(packet)} in pending writes" @@ -623,7 +699,6 @@ def prepare_single_task( task_path[:3], writers=proc.flat_writers, ) - else: return PregelTask(task_id, packet.node, task_path[:3]) elif task_path[0] == PULL: diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py new file mode 100644 index 000000000..8932697b1 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/call.py @@ -0,0 +1,113 @@ +import sys +import types + +from langgraph.utils.runnable import RunnableCallable + +""" +Utilities borrowed from cloudpickle. +https://github.com/cloudpipe/cloudpickle/blob/6220b0ce83ffee5e47e06770a1ee38ca9e47c850/cloudpickle/cloudpickle.py#L265 +""" + + +def _getattribute(obj, name): + for subpath in name.split("."): + if subpath == "": + raise AttributeError( + "Can't get local attribute {!r} on {!r}".format(name, obj) + ) + try: + parent = obj + obj = getattr(obj, subpath) + except AttributeError: + raise AttributeError( + "Can't get attribute {!r} on {!r}".format(name, obj) + ) from None + return obj, parent + + +def _whichmodule(obj, name): + """Find the module an object belongs to. + + This function differs from ``pickle.whichmodule`` in two ways: + - it does not mangle the cases where obj's module is __main__ and obj was + not found in any module. + - Errors arising during module introspection are ignored, as those errors + are considered unwanted side effects. + """ + module_name = getattr(obj, "__module__", None) + + if module_name is not None: + return module_name + # Protect the iteration by using a copy of sys.modules against dynamic + # modules that trigger imports of other modules upon calls to getattr or + # other threads importing at the same time. + for module_name, module in sys.modules.copy().items(): + # Some modules such as coverage can inject non-module objects inside + # sys.modules + if ( + module_name == "__main__" + or module_name == "__mp_main__" + or module is None + or not isinstance(module, types.ModuleType) + ): + continue + try: + if _getattribute(module, name)[0] is obj: + return module_name + except Exception: + pass + return None + + +def _lookup_module_and_qualname(obj, name=None): + if name is None: + name = getattr(obj, "__qualname__", None) + if name is None: # pragma: no cover + # This used to be needed for Python 2.7 support but is probably not + # needed anymore. However we keep the __name__ introspection in case + # users of cloudpickle rely on this old behavior for unknown reasons. + name = getattr(obj, "__name__", None) + + module_name = _whichmodule(obj, name) + + if module_name is None: + # In this case, obj.__module__ is None AND obj was not found in any + # imported module. obj is thus treated as dynamic. + return None + + if module_name == "__main__": + return None + + # Note: if module_name is in sys.modules, the corresponding module is + # assumed importable at unpickling time. See #357 + module = sys.modules.get(module_name, None) + if module is None: + # The main reason why obj's module would not be imported is that this + # module has been dynamically created, using for example + # types.ModuleType. The other possibility is that module was removed + # from sys.modules after obj was created/imported. But this case is not + # supported, as the standard pickle does not support it either. + return None + + try: + obj2, parent = _getattribute(module, name) + except AttributeError: + # obj was not found inside the module it points to + return None + if obj2 is not obj: + return None + return module, name + + +def get_runnable_for_func( + func: types.FunctionType, +) -> RunnableCallable: + if func in CACHE: + return CACHE[func] + elif not _lookup_module_and_qualname(func): + return RunnableCallable(func) + else: + return CACHE.setdefault(func, RunnableCallable(func)) + + +CACHE: dict[types.FunctionType, RunnableCallable] = {} diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 093859579..29faaab21 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -37,9 +37,7 @@ def run_with_retry( # clear any writes from previous attempts task.writes.clear() # run the task - task.proc.invoke(task.input, config) - # if successful, end - break + return task.proc.invoke(task.input, config) except ParentCommand as exc: ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS] cmd = exc.args[0] @@ -128,10 +126,10 @@ async def arun_with_retry( if stream: async for _ in task.proc.astream(task.input, config): pass + # if successful, end + break else: - await task.proc.ainvoke(task.input, config) - # if successful, end - break + return await task.proc.ainvoke(task.input, config) except ParentCommand as exc: ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS] cmd = exc.args[0] diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 067394bdb..6155ce74e 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -66,20 +66,26 @@ class PregelRunner: get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: def writer( - task: PregelExecutableTask, writes: Sequence[tuple[str, Any]] - ) -> None: + task: PregelExecutableTask, + writes: Sequence[tuple[str, Any]], + *, + calls: Optional[Sequence[Call]] = None, + ) -> Sequence[Optional[concurrent.futures.Future]]: prev_length = len(task.writes) # delegate to the underlying writer task.config[CONF][CONFIG_KEY_SEND](writes) # confirm no other concurrent writes were added # TODO could use a lock here instead, if writes can come from many threads assert len(task.writes) == prev_length + len(writes) + rtn: dict[int, Optional[concurrent.futures.Future]] = {} for idx, w in enumerate(writes, start=prev_length): # bail if not a PUSH write if w[0] != PUSH: continue # schedule the next task, if the callback returns one - if next_task := self.schedule_task(task, idx): + if next_task := self.schedule_task( + task, idx, calls[idx - prev_length] if calls else None + ): # if the parent task was retried, # the next task might already be running if any( @@ -87,18 +93,26 @@ class PregelRunner: ): continue # schedule the next task - futures[ - self.submit( - run_with_retry, - next_task, - retry_policy, - configurable={ - CONFIG_KEY_SEND: partial(writer, next_task), - # CONFIG_KEY_CALL: partial(call, next_task), - }, - __reraise_on_exit__=reraise, - ) - ] = next_task + fut = self.submit( + run_with_retry, + next_task, + retry_policy, + configurable={ + CONFIG_KEY_SEND: partial(writer, next_task), + CONFIG_KEY_CALL: partial(call, next_task), + }, + __reraise_on_exit__=reraise, + ) + futures[fut] = next_task + rtn[idx - prev_length] = fut + return [rtn.get(i) for i in range(len(writes))] + + def call( + task, func: str | Callable[[Any], Union[Awaitable[Any], Any]], input: Any + ) -> concurrent.futures.Future[Any]: + (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)]) + assert fut is not None, "writer did not return a future for call" + return fut tasks = tuple(tasks) futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {} @@ -113,7 +127,7 @@ class PregelRunner: retry_policy, configurable={ CONFIG_KEY_SEND: partial(writer, t), - # CONFIG_KEY_CALL: partial(call, t), + CONFIG_KEY_CALL: partial(call, t), }, ) self.commit(t, None) @@ -138,7 +152,7 @@ class PregelRunner: retry_policy, configurable={ CONFIG_KEY_SEND: partial(writer, t), - # CONFIG_KEY_CALL: partial(call, t), + CONFIG_KEY_CALL: partial(call, t), }, __reraise_on_exit__=reraise, ) @@ -208,7 +222,7 @@ class PregelRunner: if next_task := self.schedule_task( task, idx, - calls[idx] if calls is not None else None, + calls[idx - prev_length] if calls is not None else None, ): # if the parent task was retried, # the next task might already be running @@ -231,7 +245,7 @@ class PregelRunner: __reraise_on_exit__=reraise, ) futures[cast(asyncio.Future, fut)] = next_task - rtn[idx] = fut + rtn[idx - prev_length] = fut return [rtn.get(i) for i in range(len(writes))] def call( From 01a3c23a2901d06854d2cf9171452a174c0ce417 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 10:05:34 -0800 Subject: [PATCH 03/26] WIP --- libs/langgraph/langgraph/func/__init__.py | 85 +++++++++++++++++++++++ libs/langgraph/tests/test_pregel.py | 24 +++++++ 2 files changed, 109 insertions(+) create mode 100644 libs/langgraph/langgraph/func/__init__.py diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py new file mode 100644 index 000000000..66acba23c --- /dev/null +++ b/libs/langgraph/langgraph/func/__init__.py @@ -0,0 +1,85 @@ +import asyncio +import concurrent +import concurrent.futures +import types +from functools import partial, update_wrapper +from typing import ( + Any, + Awaitable, + Callable, + Coroutine, + Optional, + ParamSpec, + TypeVar, + Union, + overload, +) + +from langgraph.channels.ephemeral_value import EphemeralValue +from langgraph.channels.last_value import LastValue +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.constants import END, START, TAG_HIDDEN +from langgraph.pregel import Pregel +from langgraph.pregel.call import get_runnable_for_func +from langgraph.pregel.read import PregelNode +from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry +from langgraph.store.base import BaseStore +from langgraph.types import RetryPolicy, acall, call + +P = ParamSpec("P") +T = TypeVar("T") + + +@overload +def task( + *, retry: Optional[RetryPolicy] = None +) -> Callable[ + [Callable[P, Coroutine[None, None, T]]], Callable[P, asyncio.Future[T]] +]: ... + + +@overload +def task( + *, retry: Optional[RetryPolicy] = None +) -> Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]]: ... + + +def task( + *, retry: Optional[RetryPolicy] = None +) -> Callable[ + [Callable[P, Union[T, Awaitable[T]]]], + Callable[P, Union[concurrent.futures.Future[T], asyncio.Future[T]]], +]: + def _task(func: Callable[P, T]) -> Callable[P, concurrent.futures.Future[T]]: + if asyncio.iscoroutinefunction(func): + return update_wrapper(partial(acall, func), func) + else: + return update_wrapper(partial(call, func), func) + + return _task + + +def imp( + *, + checkpointer: Optional[BaseCheckpointSaver] = None, + store: Optional[BaseStore] = None, +) -> Callable[[types.FunctionType], Pregel]: + def _imp(func: types.FunctionType): + return Pregel( + nodes={ + func.__name__: PregelNode( + bound=get_runnable_for_func(func), + triggers=[START], + channels=[START], + writers=[ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN])], + ) + }, + channels={START: EphemeralValue(Any, START), END: LastValue(Any, END)}, + input_channels=START, + output_channels=END, + stream_mode="updates", + checkpointer=checkpointer, + store=store, + ) + + return _imp diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 120b3deac..356ecbadf 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -65,6 +65,7 @@ from langgraph.constants import ( START, ) from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt +from langgraph.func import imp, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.managed.shared_value import SharedValue @@ -1969,6 +1970,29 @@ def test_send_sequences() -> None: ) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + @task() + def mapper(input: str) -> str: + print(f"mapper {input}") + return input * 2 + + @imp(checkpointer=checkpointer) + def graph(input: list[str]) -> list[str]: + futures = [mapper(i) for i in input] + mapped = [f.result() for f in futures] + # answer = interrupt("question") + # TODO raises NodeInterrupt if no answer provided yet + # returns answer (saved in writes?) if provided + # what is the API for passing the answer? + return mapped + + thread1 = {"configurable": {"thread_id": "1"}} + assert graph.invoke(["0", "1"], thread1) == ["00", "11"] + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_dedupe_on_resume( request: pytest.FixtureRequest, checkpointer_name: str From 872f54adf1df688ab3ed90a6bdfbe15ea6d09061 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 14 Nov 2024 17:05:39 -0800 Subject: [PATCH 04/26] Get it working with interrupt (sync) --- libs/langgraph/langgraph/constants.py | 2 + libs/langgraph/langgraph/pregel/algo.py | 5 +- libs/langgraph/langgraph/pregel/call.py | 21 ++++++-- libs/langgraph/langgraph/pregel/io.py | 30 +++++------ libs/langgraph/langgraph/pregel/loop.py | 5 +- libs/langgraph/langgraph/pregel/runner.py | 62 ++++++++++++++++------- libs/langgraph/tests/test_pregel.py | 34 ++++++++++--- 7 files changed, 110 insertions(+), 49 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index b9d6f37f7..cd847f9be 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -40,6 +40,8 @@ SCHEDULED = sys.intern("__scheduled__") # marker to signal node was scheduled (in distributed mode) TASKS = sys.intern("__pregel_tasks") # for Send objects returned by nodes/edges, corresponds to PUSH below +RETURN = sys.intern("__return__") +# for writes of a task where we simply record the return value # --- Reserved config.configurable keys --- CONFIG_KEY_SEND = sys.intern("__pregel_send") diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c95c699c8..78242a9ac 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -52,6 +52,7 @@ from langgraph.constants import ( PUSH, RESERVED, RESUME, + RETURN, TAG_HIDDEN, TASKS, Send, @@ -259,7 +260,7 @@ def apply_writes( pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list) for task in tasks: for chan, val in task.writes: - if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT): + if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN): pass elif chan == TASKS: # TODO: remove branch in 1.0 checkpoint["pending_sends"].append(val) @@ -509,7 +510,7 @@ def prepare_single_task( proc, writes, patch_config( - merge_configs(config, {"metadata": metadata, "tags": proc.tags}), + merge_configs(config, {"metadata": metadata}), run_name=name, callbacks=( manager.get_child(f"graph:step:{step}") if manager else None diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index 8932697b1..6d218f702 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -1,7 +1,9 @@ import sys import types -from langgraph.utils.runnable import RunnableCallable +from langgraph.constants import RETURN +from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry +from langgraph.utils.runnable import RunnableCallable, RunnableSeq """ Utilities borrowed from cloudpickle. @@ -101,13 +103,24 @@ def _lookup_module_and_qualname(obj, name=None): def get_runnable_for_func( func: types.FunctionType, -) -> RunnableCallable: +) -> RunnableSeq: if func in CACHE: return CACHE[func] elif not _lookup_module_and_qualname(func): - return RunnableCallable(func) + return RunnableSeq( + RunnableCallable(func, trace=False), + ChannelWrite([ChannelWriteEntry(RETURN)]), + name=func.__name__, + ) else: - return CACHE.setdefault(func, RunnableCallable(func)) + return CACHE.setdefault( + func, + RunnableSeq( + RunnableCallable(func, trace=False), + ChannelWrite([ChannelWriteEntry(RETURN)]), + name=func.__name__, + ), + ) CACHE: dict[types.FunctionType, RunnableCallable] = {} diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index b2596d3ad..df841ffd5 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -13,6 +13,7 @@ from langgraph.constants import ( NULL_TASK_ID, PUSH, RESUME, + RETURN, TAG_HIDDEN, TASKS, ) @@ -171,22 +172,21 @@ def map_output_updates( ] if not output_tasks: return - if isinstance(output_channels, str): - updated = ( - (task.name, value) - for task, writes in output_tasks - for chan, value in writes - if chan == output_channels - ) - else: - updated = ( - ( - task.name, - {chan: value for chan, value in writes if chan in output_channels}, + updated: list[tuple[str, Any]] = [] + for task, writes in output_tasks: + if rtn := next((value for chan, value in writes if chan == RETURN), None): + updated.append((task.name, rtn)) + elif isinstance(output_channels, str): + updated.extend( + (task.name, value) for chan, value in writes if chan == output_channels + ) + elif any(chan in output_channels for chan, _ in writes): + updated.append( + ( + task.name, + {chan: value for chan, value in writes if chan in output_channels}, + ) ) - for task, writes in output_tasks - if any(chan in output_channels for chan, _ in writes) - ) grouped: dict[str, list[Any]] = {t.name: [] for t, _ in output_tasks} for node, value in updated: grouped[node].append(value) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 31397d22b..cb141c4a1 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -347,9 +347,8 @@ class PregelLoop(LoopProtocol): # match any pending writes to the new task if self.skip_done_tasks: self._match_writes({pushed.id: pushed}) - # return the new task, to be started, if not run before - if not pushed.writes: - return pushed + # return the new task, to be started if not run before + return pushed def tick( self, diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 6155ce74e..3dec1d84c 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -25,6 +25,7 @@ from langgraph.constants import ( NO_WRITES, PUSH, RESUME, + RETURN, TAG_HIDDEN, ) from langgraph.errors import GraphBubbleUp, GraphInterrupt @@ -88,23 +89,42 @@ class PregelRunner: ): # if the parent task was retried, # the next task might already be running - if any( - t == next_task.id for t in futures.values() if t is not None + if fut := next( + ( + f + for f, t in futures.items() + if t is not None and t == next_task.id + ), + None, ): - continue - # schedule the next task - fut = self.submit( - run_with_retry, - next_task, - retry_policy, - configurable={ - CONFIG_KEY_SEND: partial(writer, next_task), - CONFIG_KEY_CALL: partial(call, next_task), - }, - __reraise_on_exit__=reraise, - ) - futures[fut] = next_task - rtn[idx - prev_length] = fut + rtn[idx - prev_length] = fut + elif next_task.writes: + fut = concurrent.futures.Future() + if val := next(v for c, v in next_task.writes if c == RETURN): + fut.set_result(val) + elif exc := next(v for c, v in next_task.writes if c == ERROR): + fut.set_exception( + exc + if isinstance(exc, BaseException) + else Exception(exc) + ) + else: + fut.set_result(None) + rtn[idx - prev_length] = fut + else: + # schedule the next task + fut = self.submit( + run_with_retry, + next_task, + retry_policy, + configurable={ + CONFIG_KEY_SEND: partial(writer, next_task), + CONFIG_KEY_CALL: partial(call, next_task), + }, + __reraise_on_exit__=reraise, + ) + futures[fut] = next_task + rtn[idx - prev_length] = fut return [rtn.get(i) for i in range(len(writes))] def call( @@ -116,6 +136,7 @@ class PregelRunner: tasks = tuple(tasks) futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {} + done_futures: set[concurrent.futures.Future] = set() # give control back to the caller yield # fast path if single task with no timeout and no waiter @@ -133,7 +154,12 @@ class PregelRunner: self.commit(t, None) except Exception as exc: self.commit(t, exc) - if reraise: + if reraise and futures: + # will be re-raised after futures are done + fut = concurrent.futures.Future() + fut.set_exception(exc) + done_futures.add(fut) + elif reraise: raise if not futures: # maybe `t` schuduled another task return @@ -157,7 +183,6 @@ class PregelRunner: __reraise_on_exit__=reraise, ) ] = t - done_futures: set[concurrent.futures.Future] = set() end_time = timeout + time.monotonic() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): done, inflight = concurrent.futures.wait( @@ -183,6 +208,7 @@ class PregelRunner: del fut, task # maybe stop other tasks if _should_stop_others(done): + print("stopping others") break # give control back to the caller yield diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 356ecbadf..49ee6bec5 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1973,24 +1973,44 @@ def test_send_sequences() -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> None: checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + mapper_calls = 0 @task() def mapper(input: str) -> str: - print(f"mapper {input}") + nonlocal mapper_calls + mapper_calls += 1 return input * 2 @imp(checkpointer=checkpointer) def graph(input: list[str]) -> list[str]: futures = [mapper(i) for i in input] mapped = [f.result() for f in futures] - # answer = interrupt("question") - # TODO raises NodeInterrupt if no answer provided yet - # returns answer (saved in writes?) if provided - # what is the API for passing the answer? - return mapped + answer = interrupt("question") + return [m + answer for m in mapped] thread1 = {"configurable": {"thread_id": "1"}} - assert graph.invoke(["0", "1"], thread1) == ["00", "11"] + assert [*graph.stream(["0", "1"], thread1)] == [ + # TODO make test not depend on order of execution (which is not guaranteed) + {"mapper": "00"}, + {"mapper": "11"}, + { + "__interrupt__": ( + Interrupt( + value="question", + resumable=True, + ns=[AnyStr("graph:")], + when="during", + ), + ) + }, + ] + assert mapper_calls == 2 + + assert graph.invoke(Command(resume="answer"), thread1) == [ + "00answer", + "11answer", + ] + assert mapper_calls == 2 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) From 76a209835f81c3521f790130d2401bc7a0aad787 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 14 Nov 2024 17:52:17 -0800 Subject: [PATCH 05/26] Comments --- libs/langgraph/langgraph/pregel/runner.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 3dec1d84c..a21c3465e 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -87,8 +87,6 @@ class PregelRunner: if next_task := self.schedule_task( task, idx, calls[idx - prev_length] if calls else None ): - # if the parent task was retried, - # the next task might already be running if fut := next( ( f @@ -97,8 +95,12 @@ class PregelRunner: ), None, ): + # if the parent task was retried, + # the next task might already be running rtn[idx - prev_length] = fut elif next_task.writes: + # if it already ran, return the result + # TODO we could also set the result for non-RETURN writes fut = concurrent.futures.Future() if val := next(v for c, v in next_task.writes if c == RETURN): fut.set_result(val) From 2895a69678fe68288e740d66496029a344d5eb6c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 19 Nov 2024 17:04:22 -0800 Subject: [PATCH 06/26] Lint --- libs/langgraph/langgraph/channels/base.py | 4 +- libs/langgraph/langgraph/func/__init__.py | 17 ++++----- libs/langgraph/langgraph/pregel/algo.py | 24 ++++++------ libs/langgraph/langgraph/pregel/call.py | 17 +++++---- libs/langgraph/langgraph/pregel/runner.py | 39 ++++++++++++-------- libs/langgraph/tests/test_pregel_async.py | 45 +++++++++++++++++++++++ 6 files changed, 99 insertions(+), 47 deletions(-) diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py index 61f8908c0..4aaeb5681 100644 --- a/libs/langgraph/langgraph/channels/base.py +++ b/libs/langgraph/langgraph/channels/base.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Any, Generic, Optional, Sequence, Type, TypeVar +from typing import Any, Generic, Optional, Sequence, TypeVar from typing_extensions import Self @@ -13,7 +13,7 @@ C = TypeVar("C") class BaseChannel(Generic[Value, Update, C], ABC): __slots__ = ("key", "typ") - def __init__(self, typ: Type[Any], key: str = "") -> None: + def __init__(self, typ: Any, key: str = "") -> None: self.typ = typ self.key = key diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 66acba23c..6826a8740 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -7,7 +7,6 @@ from typing import ( Any, Awaitable, Callable, - Coroutine, Optional, ParamSpec, TypeVar, @@ -33,22 +32,20 @@ T = TypeVar("T") @overload def task( *, retry: Optional[RetryPolicy] = None -) -> Callable[ - [Callable[P, Coroutine[None, None, T]]], Callable[P, asyncio.Future[T]] -]: ... +) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]]: ... @overload -def task( +def task( # type: ignore[overload-cannot-match] *, retry: Optional[RetryPolicy] = None ) -> Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]]: ... def task( *, retry: Optional[RetryPolicy] = None -) -> Callable[ - [Callable[P, Union[T, Awaitable[T]]]], - Callable[P, Union[concurrent.futures.Future[T], asyncio.Future[T]]], +) -> Union[ + Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]], + Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]], ]: def _task(func: Callable[P, T]) -> Callable[P, concurrent.futures.Future[T]]: if asyncio.iscoroutinefunction(func): @@ -64,7 +61,7 @@ def imp( checkpointer: Optional[BaseCheckpointSaver] = None, store: Optional[BaseStore] = None, ) -> Callable[[types.FunctionType], Pregel]: - def _imp(func: types.FunctionType): + def _imp(func: types.FunctionType) -> Pregel: return Pregel( nodes={ func.__name__: PregelNode( @@ -74,7 +71,7 @@ def imp( writers=[ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN])], ) }, - channels={START: EphemeralValue(Any, START), END: LastValue(Any, END)}, + channels={START: EphemeralValue(Any), END: LastValue(Any, END)}, input_channels=START, output_channels=END, stream_mode="updates", diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 78242a9ac..330d46449 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -102,10 +102,10 @@ class PregelTaskWrites(NamedTuple): class Call: __slots__ = ("func", "input") - func: str | Callable + func: Callable input: Any - def __init__(self, func: str | Callable, input: Any) -> None: + def __init__(self, func: Callable, input: Any) -> None: self.func = func self.input = input @@ -451,7 +451,7 @@ def prepare_next_tasks( def prepare_single_task( - task_path: tuple[Union[str, int, tuple], ...], + task_path: tuple[Any, ...], task_id_checksum: Optional[str], *, checkpoint: Checkpoint, @@ -474,10 +474,10 @@ def prepare_single_task( if task_path[0] == PUSH and isinstance(task_path[-1], Call): # (PUSH, parent task path, idx of PUSH write, id of parent task, Call) - task_path_t = cast(tuple[str, tuple, int, str, Optional[Call]], task_path) + task_path_t = cast(tuple[str, tuple, int, str, Call], task_path) call = task_path_t[-1] - proc = get_runnable_for_func(call.func) - name = proc.name + proc_ = get_runnable_for_func(call.func) + name = proc_.name if name is None: raise ValueError("`call` functions must have a `__name__` attribute") # create task id @@ -507,7 +507,7 @@ def prepare_single_task( return PregelExecutableTask( name, call.input, - proc, + proc_, writes, patch_config( merge_configs(config, {"metadata": metadata}), @@ -586,14 +586,14 @@ def prepare_single_task( elif len(task_path) >= 4: # new PUSH tasks, executed in superstep n # (PUSH, parent task path, idx of PUSH write, id of parent task) - task_path_t = cast(tuple[str, tuple, int, str], task_path) - writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]] - if task_path_t[2] >= len(writes_for_path): + task_path_tt = cast(tuple[str, tuple, int, str], task_path) + writes_for_path = [w for w in pending_writes if w[0] == task_path_tt[3]] + if task_path_tt[2] >= len(writes_for_path): logger.warning( f"Ignoring invalid write index {task_path[2]} in pending writes" ) return - packet = writes_for_path[task_path_t[2]][2] + packet = writes_for_path[task_path_tt[2]][2] if packet is None: return if not isinstance(packet, Send): @@ -638,7 +638,7 @@ def prepare_single_task( if node := proc.node: if proc.metadata: metadata.update(proc.metadata) - writes: deque[tuple[str, Any]] = deque() + writes = deque() return PregelExecutableTask( packet.node, packet.arg, diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index 6d218f702..53ece995a 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -1,5 +1,6 @@ import sys import types +from typing import Any, Callable, Optional from langgraph.constants import RETURN from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry @@ -11,7 +12,7 @@ https://github.com/cloudpipe/cloudpickle/blob/6220b0ce83ffee5e47e06770a1ee38ca9e """ -def _getattribute(obj, name): +def _getattribute(obj: Any, name: str) -> Any: for subpath in name.split("."): if subpath == "": raise AttributeError( @@ -27,7 +28,7 @@ def _getattribute(obj, name): return obj, parent -def _whichmodule(obj, name): +def _whichmodule(obj: Any, name: str) -> Optional[str]: """Find the module an object belongs to. This function differs from ``pickle.whichmodule`` in two ways: @@ -61,7 +62,9 @@ def _whichmodule(obj, name): return None -def _lookup_module_and_qualname(obj, name=None): +def _lookup_module_and_qualname( + obj: Any, name: Optional[str] = None +) -> Optional[tuple[types.ModuleType, str]]: if name is None: name = getattr(obj, "__qualname__", None) if name is None: # pragma: no cover @@ -69,6 +72,8 @@ def _lookup_module_and_qualname(obj, name=None): # needed anymore. However we keep the __name__ introspection in case # users of cloudpickle rely on this old behavior for unknown reasons. name = getattr(obj, "__name__", None) + if name is None: + return None module_name = _whichmodule(obj, name) @@ -101,9 +106,7 @@ def _lookup_module_and_qualname(obj, name=None): return module, name -def get_runnable_for_func( - func: types.FunctionType, -) -> RunnableSeq: +def get_runnable_for_func(func: Callable[..., Any]) -> RunnableSeq: if func in CACHE: return CACHE[func] elif not _lookup_module_and_qualname(func): @@ -123,4 +126,4 @@ def get_runnable_for_func( ) -CACHE: dict[types.FunctionType, RunnableCallable] = {} +CACHE: dict[Callable[..., Any], RunnableSeq] = {} diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index a21c3465e..150e84f30 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -130,7 +130,9 @@ class PregelRunner: return [rtn.get(i) for i in range(len(writes))] def call( - task, func: str | Callable[[Any], Union[Awaitable[Any], Any]], input: Any + task: PregelExecutableTask, + func: Callable[[Any], Union[Awaitable[Any], Any]], + input: Any, ) -> concurrent.futures.Future[Any]: (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)]) assert fut is not None, "writer did not return a future for call" @@ -158,7 +160,7 @@ class PregelRunner: self.commit(t, exc) if reraise and futures: # will be re-raised after futures are done - fut = concurrent.futures.Future() + fut: concurrent.futures.Future = concurrent.futures.Future() fut.set_exception(exc) done_futures.add(fut) elif reraise: @@ -259,25 +261,30 @@ class PregelRunner: ): continue # schedule the next task - fut = self.submit( - arun_with_retry, - next_task, - retry_policy, - stream=self.use_astream, - configurable={ - CONFIG_KEY_SEND: partial(writer, next_task), - CONFIG_KEY_CALL: partial(call, next_task), - }, - __name__=t.name, - __cancel_on_exit__=True, - __reraise_on_exit__=reraise, + fut = cast( + asyncio.Future, + self.submit( + arun_with_retry, + next_task, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_SEND: partial(writer, next_task), + CONFIG_KEY_CALL: partial(call, next_task), + }, + __name__=t.name, + __cancel_on_exit__=True, + __reraise_on_exit__=reraise, + ), ) - futures[cast(asyncio.Future, fut)] = next_task + futures[fut] = next_task rtn[idx - prev_length] = fut return [rtn.get(i) for i in range(len(writes))] def call( - task, func: str | Callable[[Any], Union[Awaitable[Any], Any]], input: Any + task: PregelExecutableTask, + func: Callable[[Any], Union[Awaitable[Any], Any]], + input: Any, ) -> asyncio.Future[Any]: (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)]) assert fut is not None, "writer did not return a future for call" diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index d8aeedd9a..450606d80 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -62,6 +62,7 @@ from langgraph.constants import ( START, ) from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt +from langgraph.func import imp, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.managed.shared_value import SharedValue @@ -2647,6 +2648,50 @@ async def test_send_sequences(checkpointer_name: str) -> None: ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_imp_task(checkpointer_name: str) -> None: + mapper_calls = 0 + + @task() + async def mapper(input: str) -> str: + nonlocal mapper_calls + mapper_calls += 1 + return input * 2 + + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @imp(checkpointer=checkpointer) + async def graph(input: list[str]) -> list[str]: + futures = [mapper(i) for i in input] + mapped = await asyncio.gather(*futures) + answer = interrupt("question") + return [m + answer for m in mapped] + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c async for c in graph.astream(["0", "1"], thread1)] == [ + # TODO make test not depend on order of execution (which is not guaranteed) + {"mapper": "00"}, + {"mapper": "11"}, + { + "__interrupt__": ( + Interrupt( + value="question", + resumable=True, + ns=[AnyStr("graph:")], + when="during", + ), + ) + }, + ] + assert mapper_calls == 2 + + assert await graph.ainvoke(Command(resume="answer"), thread1) == [ + "00answer", + "11answer", + ] + assert mapper_calls == 2 + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: if not FF_SEND_V2: From 2e9aea6fc8b08d66ae33f177bacb790bddec0c72 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 19 Nov 2024 17:11:12 -0800 Subject: [PATCH 07/26] Lint --- libs/langgraph/langgraph/types.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index b49b1c45e..a0949613c 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -15,7 +15,6 @@ from typing import ( Literal, NamedTuple, Optional, - ParamSpec, Sequence, Type, TypedDict, @@ -25,7 +24,7 @@ from typing import ( ) from langchain_core.runnables import Runnable, RunnableConfig -from typing_extensions import Self +from typing_extensions import ParamSpec, Self from langgraph.checkpoint.base import ( BaseCheckpointSaver, From a443b3b25677b83950ec3128b1f8701f456b4e84 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 19 Nov 2024 17:16:24 -0800 Subject: [PATCH 08/26] Fix --- libs/langgraph/langgraph/pregel/call.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index 53ece995a..a837ff4a8 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -1,3 +1,4 @@ +import asyncio import sys import types from typing import Any, Callable, Optional @@ -111,7 +112,9 @@ def get_runnable_for_func(func: Callable[..., Any]) -> RunnableSeq: return CACHE[func] elif not _lookup_module_and_qualname(func): return RunnableSeq( - RunnableCallable(func, trace=False), + RunnableCallable(None, func, trace=False) + if asyncio.iscoroutinefunction(func) + else RunnableCallable(func, trace=False), ChannelWrite([ChannelWriteEntry(RETURN)]), name=func.__name__, ) @@ -119,7 +122,9 @@ def get_runnable_for_func(func: Callable[..., Any]) -> RunnableSeq: return CACHE.setdefault( func, RunnableSeq( - RunnableCallable(func, trace=False), + RunnableCallable(None, func, trace=False) + if asyncio.iscoroutinefunction(func) + else RunnableCallable(func, trace=False), ChannelWrite([ChannelWriteEntry(RETURN)]), name=func.__name__, ), From 90dd2b01b6bbad852fefc4d0faeecf4ccbf669e3 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 19 Nov 2024 17:19:30 -0800 Subject: [PATCH 09/26] Comment --- libs/langgraph/langgraph/pregel/runner.py | 1 + 1 file changed, 1 insertion(+) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 150e84f30..a7a17fc09 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -248,6 +248,7 @@ class PregelRunner: # bail if not a PUSH write if w[0] != PUSH: continue + # TODO apply changes from sync version # schedule the next task, if the callback returns one if next_task := self.schedule_task( task, From 287c29fbdc641607ff71e362a5d359a8599e94a2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 20 Nov 2024 15:38:43 -0800 Subject: [PATCH 10/26] Fix async --- libs/langgraph/langgraph/pregel/runner.py | 119 +++++++++++++++------- libs/langgraph/langgraph/types.py | 1 - 2 files changed, 80 insertions(+), 40 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index a7a17fc09..7c10a72ac 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -1,5 +1,6 @@ import asyncio import concurrent.futures +import threading import time from functools import partial from typing import ( @@ -66,18 +67,26 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: + locks: dict[str, threading.Lock] = {} + def writer( task: PregelExecutableTask, writes: Sequence[tuple[str, Any]], *, calls: Optional[Sequence[Call]] = None, ) -> Sequence[Optional[concurrent.futures.Future]]: - prev_length = len(task.writes) - # delegate to the underlying writer - task.config[CONF][CONFIG_KEY_SEND](writes) - # confirm no other concurrent writes were added - # TODO could use a lock here instead, if writes can come from many threads - assert len(task.writes) == prev_length + len(writes) + if all(w[0] != PUSH for w in writes): + return task.config[CONF][CONFIG_KEY_SEND](writes) + + if task.id not in locks: + locks[task.id] = threading.Lock() + with locks[task.id]: + prev_length = len(task.writes) + # delegate to the underlying writer + task.config[CONF][CONFIG_KEY_SEND](writes) + # confirm no other concurrent writes were added + assert len(task.writes) == prev_length + len(writes) + # schedule PUSH tasks, collect futures rtn: dict[int, Optional[concurrent.futures.Future]] = {} for idx, w in enumerate(writes, start=prev_length): # bail if not a PUSH write @@ -100,7 +109,6 @@ class PregelRunner: rtn[idx - prev_length] = fut elif next_task.writes: # if it already ran, return the result - # TODO we could also set the result for non-RETURN writes fut = concurrent.futures.Future() if val := next(v for c, v in next_task.writes if c == RETURN): fut.set_result(val) @@ -212,7 +220,6 @@ class PregelRunner: del fut, task # maybe stop other tasks if _should_stop_others(done): - print("stopping others") break # give control back to the caller yield @@ -231,24 +238,31 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, ) -> AsyncIterator[None]: + locks: dict[str, threading.Lock] = {} + def writer( task: PregelExecutableTask, writes: Sequence[tuple[str, Any]], *, calls: Optional[Sequence[Call]] = None, ) -> Sequence[Optional[asyncio.Future]]: - prev_length = len(task.writes) - # delegate to the underlying writer - task.config[CONF][CONFIG_KEY_SEND](writes) - # confirm no other concurrent writes were added - # TODO could use a lock here instead, if writes can come from many threads - assert len(task.writes) == prev_length + len(writes) + if all(w[0] != PUSH for w in writes): + return task.config[CONF][CONFIG_KEY_SEND](writes) + + if task.id not in locks: + locks[task.id] = threading.Lock() + with locks[task.id]: + prev_length = len(task.writes) + # delegate to the underlying writer + task.config[CONF][CONFIG_KEY_SEND](writes) + # confirm no other concurrent writes were added + assert len(task.writes) == prev_length + len(writes) + # schedule PUSH tasks, collect futures rtn: dict[int, Optional[asyncio.Future]] = {} for idx, w in enumerate(writes, start=prev_length): # bail if not a PUSH write if w[0] != PUSH: continue - # TODO apply changes from sync version # schedule the next task, if the callback returns one if next_task := self.schedule_task( task, @@ -257,29 +271,51 @@ class PregelRunner: ): # if the parent task was retried, # the next task might already be running - if any( - t == next_task.id for t in futures.values() if t is not None - ): - continue - # schedule the next task - fut = cast( - asyncio.Future, - self.submit( - arun_with_retry, - next_task, - retry_policy, - stream=self.use_astream, - configurable={ - CONFIG_KEY_SEND: partial(writer, next_task), - CONFIG_KEY_CALL: partial(call, next_task), - }, - __name__=t.name, - __cancel_on_exit__=True, - __reraise_on_exit__=reraise, + if fut := next( + ( + f + for f, t in futures.items() + if t is not None and t == next_task.id ), - ) - futures[fut] = next_task - rtn[idx - prev_length] = fut + None, + ): + # if the parent task was retried, + # the next task might already be running + rtn[idx - prev_length] = fut + elif next_task.writes: + # if it already ran, return the result + fut = asyncio.Future() + if val := next(v for c, v in next_task.writes if c == RETURN): + fut.set_result(val) + elif exc := next(v for c, v in next_task.writes if c == ERROR): + fut.set_exception( + exc + if isinstance(exc, BaseException) + else Exception(exc) + ) + else: + fut.set_result(None) + rtn[idx - prev_length] = fut + else: + # schedule the next task + fut = cast( + asyncio.Future, + self.submit( + arun_with_retry, + next_task, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_SEND: partial(writer, next_task), + CONFIG_KEY_CALL: partial(call, next_task), + }, + __name__=t.name, + __cancel_on_exit__=True, + __reraise_on_exit__=reraise, + ), + ) + futures[fut] = next_task + rtn[idx - prev_length] = fut return [rtn.get(i) for i in range(len(writes))] def call( @@ -294,6 +330,7 @@ class PregelRunner: loop = asyncio.get_event_loop() tasks = tuple(tasks) futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {} + done_futures: set[asyncio.Future] = set() # give control back to the caller yield # fast path if single task with no waiter and no timeout @@ -312,7 +349,12 @@ class PregelRunner: self.commit(t, None) except Exception as exc: self.commit(t, exc) - if reraise: + if reraise and futures: + # will be re-raised after futures are done + fut: asyncio.Future = loop.create_future() + fut.set_exception(exc) + done_futures.add(fut) + elif reraise: raise if not futures: # maybe `t` schuduled another task return @@ -342,7 +384,6 @@ class PregelRunner: ), ) ] = t - done_futures: set[asyncio.Future] = set() end_time = timeout + loop.time() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): done, inflight = await asyncio.wait( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index a0949613c..3bdbf2049 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -398,5 +398,4 @@ def acall( fut = impl(func, *args, **kwargs) if isinstance(fut, concurrent.futures.Future): fut = asyncio.wrap_future(fut) - fut.cancel return fut From d93be914c7187861bdb0bd6031abee53acb993d2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 20 Nov 2024 17:09:54 -0800 Subject: [PATCH 11/26] Fix stream order --- libs/langgraph/langgraph/func/__init__.py | 7 +- libs/langgraph/langgraph/pregel/call.py | 25 ++--- libs/langgraph/langgraph/pregel/runner.py | 94 +++++++++-------- libs/langgraph/langgraph/types.py | 18 ---- libs/langgraph/langgraph/utils/future.py | 121 ++++++++++++++++++++++ libs/langgraph/tests/test_pregel.py | 44 +++++++- libs/langgraph/tests/test_pregel_async.py | 84 +++++++++++++-- 7 files changed, 292 insertions(+), 101 deletions(-) create mode 100644 libs/langgraph/langgraph/utils/future.py diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 6826a8740..b48cf32e2 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -23,7 +23,7 @@ from langgraph.pregel.call import get_runnable_for_func from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import RetryPolicy, acall, call +from langgraph.types import RetryPolicy, call P = ParamSpec("P") T = TypeVar("T") @@ -48,10 +48,7 @@ def task( Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]], ]: def _task(func: Callable[P, T]) -> Callable[P, concurrent.futures.Future[T]]: - if asyncio.iscoroutinefunction(func): - return update_wrapper(partial(acall, func), func) - else: - return update_wrapper(partial(call, func), func) + return update_wrapper(partial(call, func), func) return _task diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index a837ff4a8..a9986102d 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -1,11 +1,10 @@ -import asyncio import sys import types from typing import Any, Callable, Optional from langgraph.constants import RETURN from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry -from langgraph.utils.runnable import RunnableCallable, RunnableSeq +from langgraph.utils.runnable import RunnableSeq, coerce_to_runnable """ Utilities borrowed from cloudpickle. @@ -110,25 +109,15 @@ def _lookup_module_and_qualname( def get_runnable_for_func(func: Callable[..., Any]) -> RunnableSeq: if func in CACHE: return CACHE[func] - elif not _lookup_module_and_qualname(func): - return RunnableSeq( - RunnableCallable(None, func, trace=False) - if asyncio.iscoroutinefunction(func) - else RunnableCallable(func, trace=False), + else: + seq = RunnableSeq( + coerce_to_runnable(func, name=None, trace=False), ChannelWrite([ChannelWriteEntry(RETURN)]), name=func.__name__, ) - else: - return CACHE.setdefault( - func, - RunnableSeq( - RunnableCallable(None, func, trace=False) - if asyncio.iscoroutinefunction(func) - else RunnableCallable(func, trace=False), - ChannelWrite([ChannelWriteEntry(RETURN)]), - name=func.__name__, - ), - ) + if not _lookup_module_and_qualname(func): + return seq + return CACHE.setdefault(func, seq) CACHE: dict[Callable[..., Any], RunnableSeq] = {} diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 7c10a72ac..8decbedec 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -34,6 +34,7 @@ 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.utils.future import chain_future class PregelRunner: @@ -133,6 +134,7 @@ class PregelRunner: }, __reraise_on_exit__=reraise, ) + fut.add_done_callback(partial(self.commit, next_task)) futures[fut] = next_task rtn[idx - prev_length] = fut return [rtn.get(i) for i in range(len(writes))] @@ -165,7 +167,7 @@ class PregelRunner: ) self.commit(t, None) except Exception as exc: - self.commit(t, exc) + self.commit(t, None, exc) if reraise and futures: # will be re-raised after futures are done fut: concurrent.futures.Future = concurrent.futures.Future() @@ -183,18 +185,18 @@ class PregelRunner: # yield updates/debug output as each task finishes for t in tasks: if not t.writes: - futures[ - self.submit( - run_with_retry, - t, - retry_policy, - configurable={ - CONFIG_KEY_SEND: partial(writer, t), - CONFIG_KEY_CALL: partial(call, t), - }, - __reraise_on_exit__=reraise, - ) - ] = t + fut = self.submit( + run_with_retry, + t, + retry_policy, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + CONFIG_KEY_CALL: partial(call, t), + }, + __reraise_on_exit__=reraise, + ) + fut.add_done_callback(partial(self.commit, t)) + futures[fut] = t end_time = timeout + time.monotonic() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): done, inflight = concurrent.futures.wait( @@ -213,8 +215,6 @@ class PregelRunner: else: # store for panic check done_futures.add(fut) - # task finished, commit writes - self.commit(task, _exception(fut)) else: # remove references to loop vars del fut, task @@ -264,11 +264,8 @@ class PregelRunner: if w[0] != PUSH: continue # schedule the next task, if the callback returns one - if next_task := self.schedule_task( - task, - idx, - calls[idx - prev_length] if calls is not None else None, - ): + wcall = calls[idx - prev_length] if calls is not None else None + if next_task := self.schedule_task(task, idx, wcall): # if the parent task was retried, # the next task might already be running if fut := next( @@ -314,6 +311,7 @@ class PregelRunner: __reraise_on_exit__=reraise, ), ) + fut.add_done_callback(partial(self.commit, next_task)) futures[fut] = next_task rtn[idx - prev_length] = fut return [rtn.get(i) for i in range(len(writes))] @@ -322,10 +320,15 @@ class PregelRunner: task: PregelExecutableTask, func: Callable[[Any], Union[Awaitable[Any], Any]], input: Any, - ) -> asyncio.Future[Any]: + ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]: (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)]) assert fut is not None, "writer did not return a future for call" - return fut + if asyncio.iscoroutinefunction(func): + return fut + # adapted from asyncio.run_coroutine_threadsafe + sfut = concurrent.futures.Future() + loop.call_soon_threadsafe(chain_future, fut, sfut) + return sfut loop = asyncio.get_event_loop() tasks = tuple(tasks) @@ -348,7 +351,7 @@ class PregelRunner: ) self.commit(t, None) except Exception as exc: - self.commit(t, exc) + self.commit(t, None, exc) if reraise and futures: # will be re-raised after futures are done fut: asyncio.Future = loop.create_future() @@ -366,24 +369,24 @@ class PregelRunner: # yield updates/debug output as each task finishes for t in tasks: if not t.writes: - futures[ - cast( - asyncio.Future, - self.submit( - arun_with_retry, - t, - retry_policy, - stream=self.use_astream, - configurable={ - CONFIG_KEY_SEND: partial(writer, t), - CONFIG_KEY_CALL: partial(call, t), - }, - __name__=t.name, - __cancel_on_exit__=True, - __reraise_on_exit__=reraise, - ), - ) - ] = t + fut = cast( + asyncio.Future, + self.submit( + arun_with_retry, + t, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + CONFIG_KEY_CALL: partial(call, t), + }, + __name__=t.name, + __cancel_on_exit__=True, + __reraise_on_exit__=reraise, + ), + ) + fut.add_done_callback(partial(self.commit, t)) + futures[fut] = t end_time = timeout + loop.time() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): done, inflight = await asyncio.wait( @@ -402,8 +405,6 @@ class PregelRunner: else: # store for panic check done_futures.add(fut) - # task finished, commit writes - self.commit(task, _exception(fut)) else: # remove references to loop vars del fut, task @@ -423,8 +424,13 @@ class PregelRunner: ) def commit( - self, task: PregelExecutableTask, exception: Optional[BaseException] + self, + task: PregelExecutableTask, + fut: Union[None, concurrent.futures.Future[Any], asyncio.Future[Any]], + exception: Optional[BaseException] = None, ) -> None: + if fut is not None: + exception = _exception(fut) if exception: if isinstance(exception, GraphInterrupt): # save interrupt to checkpointer diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3bdbf2049..5ce64b52b 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,4 +1,3 @@ -import asyncio import concurrent import concurrent.futures import dataclasses @@ -7,7 +6,6 @@ from collections import deque from typing import ( TYPE_CHECKING, Any, - Awaitable, Callable, ClassVar, Generic, @@ -382,20 +380,4 @@ def call( conf = get_configurable() impl = conf[CONFIG_KEY_CALL] fut = impl(func, *args, **kwargs) - if not isinstance(fut, concurrent.futures.Future): - raise RuntimeError("In an async context, use acall() instead of call()") - return fut - - -def acall( - func: str | Callable[P, Union[T, Awaitable[T]]], *args: P.args, **kwargs: P.kwargs -) -> asyncio.Future[T]: - from langgraph.constants import CONFIG_KEY_CALL - from langgraph.utils.config import get_configurable - - conf = get_configurable() - impl = conf[CONFIG_KEY_CALL] - fut = impl(func, *args, **kwargs) - if isinstance(fut, concurrent.futures.Future): - fut = asyncio.wrap_future(fut) return fut diff --git a/libs/langgraph/langgraph/utils/future.py b/libs/langgraph/langgraph/utils/future.py new file mode 100644 index 000000000..d7b8718f6 --- /dev/null +++ b/libs/langgraph/langgraph/utils/future.py @@ -0,0 +1,121 @@ +import asyncio +import concurrent.futures +from typing import Union + +AnyFuture = Union[asyncio.Future, concurrent.futures.Future] + + +def _get_loop(fut: asyncio.Future) -> asyncio.AbstractEventLoop: + # Tries to call Future.get_loop() if it's available. + # Otherwise fallbacks to using the old '_loop' property. + try: + get_loop = fut.get_loop + except AttributeError: + pass + else: + return get_loop() + return fut._loop + + +def _convert_future_exc(exc): + exc_class = type(exc) + if exc_class is concurrent.futures.CancelledError: + return asyncio.CancelledError(*exc.args) + elif exc_class is concurrent.futures.TimeoutError: + return asyncio.TimeoutError(*exc.args) + elif exc_class is concurrent.futures.InvalidStateError: + return asyncio.InvalidStateError(*exc.args) + else: + return exc + + +def _set_concurrent_future_state(concurrent, source): + """Copy state from a future to a concurrent.futures.Future.""" + assert source.done() + if source.cancelled(): + concurrent.cancel() + if not concurrent.set_running_or_notify_cancel(): + return + exception = source.exception() + if exception is not None: + concurrent.set_exception(_convert_future_exc(exception)) + else: + result = source.result() + concurrent.set_result(result) + + +def _copy_future_state(source, dest): + """Internal helper to copy state from another Future. + + The other Future may be a concurrent.futures.Future. + """ + assert source.done() + if dest.cancelled(): + return + assert not dest.done() + if source.cancelled(): + dest.cancel() + else: + exception = source.exception() + if exception is not None: + dest.set_exception(_convert_future_exc(exception)) + else: + result = source.result() + dest.set_result(result) + + +def _chain_future(source: AnyFuture, destination: AnyFuture) -> None: + """Chain two futures so that when one completes, so does the other. + + The result (or exception) of source will be copied to destination. + If destination is cancelled, source gets cancelled too. + Compatible with both asyncio.Future and concurrent.futures.Future. + """ + if not asyncio.isfuture(source) and not isinstance( + source, concurrent.futures.Future + ): + raise TypeError("A future is required for source argument") + if not asyncio.isfuture(destination) and not isinstance( + destination, concurrent.futures.Future + ): + raise TypeError("A future is required for destination argument") + source_loop = _get_loop(source) if asyncio.isfuture(source) else None + dest_loop = _get_loop(destination) if asyncio.isfuture(destination) else None + + def _set_state(future, other): + if asyncio.isfuture(future): + _copy_future_state(other, future) + else: + _set_concurrent_future_state(future, other) + + def _call_check_cancel(destination): + if destination.cancelled(): + if source_loop is None or source_loop is dest_loop: + source.cancel() + else: + source_loop.call_soon_threadsafe(source.cancel) + + def _call_set_state(source): + if destination.cancelled() and dest_loop is not None and dest_loop.is_closed(): + return + if dest_loop is None or dest_loop is source_loop: + _set_state(destination, source) + else: + if dest_loop.is_closed(): + return + dest_loop.call_soon_threadsafe(_set_state, destination, source) + + destination.add_done_callback(_call_check_cancel) + source.add_done_callback(_call_set_state) + + +def chain_future(source: AnyFuture, destination: AnyFuture) -> None: + # adapted from asyncio.run_coroutine_threadsafe + try: + _chain_future(source, destination) + except (SystemExit, KeyboardInterrupt): + raise + except BaseException as exc: + if destination.set_running_or_notify_cancel(): + destination.set_exception(exc) + raise diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 49ee6bec5..ab675ea79 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1976,21 +1976,21 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non mapper_calls = 0 @task() - def mapper(input: str) -> str: + def mapper(input: int) -> str: nonlocal mapper_calls mapper_calls += 1 - return input * 2 + time.sleep(input / 100) + return str(input) * 2 @imp(checkpointer=checkpointer) - def graph(input: list[str]) -> list[str]: + def graph(input: list[int]) -> list[str]: futures = [mapper(i) for i in input] mapped = [f.result() for f in futures] answer = interrupt("question") return [m + answer for m in mapped] thread1 = {"configurable": {"thread_id": "1"}} - assert [*graph.stream(["0", "1"], thread1)] == [ - # TODO make test not depend on order of execution (which is not guaranteed) + assert [*graph.stream([0, 1], thread1)] == [ {"mapper": "00"}, {"mapper": "11"}, { @@ -2013,6 +2013,40 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non assert mapper_calls == 2 +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_imp_stream_order( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + @task() + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo", "b": "bar"} + + @task() + def bar(state: dict) -> dict: + return {"a": state["a"] + state["b"], "c": "bark"} + + @task() + def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} + + @imp(checkpointer=checkpointer) + def graph(state: dict) -> dict: + fut_foo = foo(state) + fut_bar = bar(fut_foo.result()) + fut_baz = baz(fut_bar.result()) + return fut_baz.result() + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c for c in graph.stream({"a": "0"}, thread1)] == [ + {"foo": {"a": "0foo", "b": "bar"}}, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_dedupe_on_resume( request: pytest.FixtureRequest, checkpointer_name: str diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 450606d80..07e16a538 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2650,26 +2650,24 @@ async def test_send_sequences(checkpointer_name: str) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_task(checkpointer_name: str) -> None: - mapper_calls = 0 - - @task() - async def mapper(input: str) -> str: - nonlocal mapper_calls - mapper_calls += 1 - return input * 2 - async with awith_checkpointer(checkpointer_name) as checkpointer: + mapper_calls = 0 + + @task() + async def mapper(input: int) -> str: + nonlocal mapper_calls + mapper_calls += 1 + return str(input) * 2 @imp(checkpointer=checkpointer) - async def graph(input: list[str]) -> list[str]: + async def graph(input: list[int]) -> list[str]: futures = [mapper(i) for i in input] mapped = await asyncio.gather(*futures) answer = interrupt("question") return [m + answer for m in mapped] thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream(["0", "1"], thread1)] == [ - # TODO make test not depend on order of execution (which is not guaranteed) + assert [c async for c in graph.astream([0, 1], thread1)] == [ {"mapper": "00"}, {"mapper": "11"}, { @@ -2692,6 +2690,70 @@ async def test_imp_task(checkpointer_name: str) -> None: assert mapper_calls == 2 +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_imp_sync_from_async(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @task() + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo", "b": "bar"} + + @task() + def bar(state: dict) -> dict: + return {"a": state["a"] + state["b"], "c": "bark"} + + @task() + def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} + + @imp(checkpointer=checkpointer) + def graph(state: dict) -> dict: + fut_foo = foo(state) + fut_bar = bar(fut_foo.result()) + fut_baz = baz(fut_bar.result()) + return fut_baz.result() + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c async for c in graph.astream({"a": "0"}, thread1)] == [ + {"foo": {"a": "0foo", "b": "bar"}}, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_imp_stream_order(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @task() + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo", "b": "bar"} + + @task() + async def bar(state: dict) -> dict: + return {"a": state["a"] + state["b"], "c": "bark"} + + @task() + async def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} + + @imp(checkpointer=checkpointer) + async def graph(state: dict) -> dict: + fut_foo = foo(state) + fut_bar = bar(await fut_foo) + fut_baz = baz(await fut_bar) + return await fut_baz + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c async for c in graph.astream({"a": "0"}, thread1)] == [ + {"foo": {"a": "0foo", "b": "bar"}}, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: if not FF_SEND_V2: From a91dbf9b70f7cd397e6ea666f87b390e2ecb51ad Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 21 Nov 2024 08:10:39 -0800 Subject: [PATCH 12/26] Lint --- libs/langgraph/langgraph/pregel/__init__.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 52cd2bf09..e714afe21 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -541,11 +541,6 @@ class Pregel(PregelProtocol): next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: apply_writes(saved.checkpoint, channels, tasks, None) - print( - "next_tasks", - [(t.id, t.name, t.input, bool(t.writes)) for t in next_tasks.values()], - ) - print("pending_writes", saved.pending_writes) # assemble the state snapshot return StateSnapshot( read_channels(channels, self.stream_channels_asis), From 0663d46c478936923988d25a5fa36edfb41d0728 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 21 Nov 2024 08:11:23 -0800 Subject: [PATCH 13/26] Rename --- libs/langgraph/langgraph/func/__init__.py | 2 +- libs/langgraph/tests/test_pregel.py | 6 +++--- libs/langgraph/tests/test_pregel_async.py | 8 ++++---- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index b48cf32e2..69e9b6c5d 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -53,7 +53,7 @@ def task( return _task -def imp( +def entrypoint( *, checkpointer: Optional[BaseCheckpointSaver] = None, store: Optional[BaseStore] = None, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ab675ea79..8de87ef0e 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -65,7 +65,7 @@ from langgraph.constants import ( START, ) from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt -from langgraph.func import imp, task +from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.managed.shared_value import SharedValue @@ -1982,7 +1982,7 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non time.sleep(input / 100) return str(input) * 2 - @imp(checkpointer=checkpointer) + @entrypoint(checkpointer=checkpointer) def graph(input: list[int]) -> list[str]: futures = [mapper(i) for i in input] mapped = [f.result() for f in futures] @@ -2031,7 +2031,7 @@ def test_imp_stream_order( def baz(state: dict) -> dict: return {"a": state["a"] + "baz", "c": "something else"} - @imp(checkpointer=checkpointer) + @entrypoint(checkpointer=checkpointer) def graph(state: dict) -> dict: fut_foo = foo(state) fut_bar = bar(fut_foo.result()) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 07e16a538..c4e5174f6 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -62,7 +62,7 @@ from langgraph.constants import ( START, ) from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt -from langgraph.func import imp, task +from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.managed.shared_value import SharedValue @@ -2659,7 +2659,7 @@ async def test_imp_task(checkpointer_name: str) -> None: mapper_calls += 1 return str(input) * 2 - @imp(checkpointer=checkpointer) + @entrypoint(checkpointer=checkpointer) async def graph(input: list[int]) -> list[str]: futures = [mapper(i) for i in input] mapped = await asyncio.gather(*futures) @@ -2706,7 +2706,7 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None: def baz(state: dict) -> dict: return {"a": state["a"] + "baz", "c": "something else"} - @imp(checkpointer=checkpointer) + @entrypoint(checkpointer=checkpointer) def graph(state: dict) -> dict: fut_foo = foo(state) fut_bar = bar(fut_foo.result()) @@ -2738,7 +2738,7 @@ async def test_imp_stream_order(checkpointer_name: str) -> None: async def baz(state: dict) -> dict: return {"a": state["a"] + "baz", "c": "something else"} - @imp(checkpointer=checkpointer) + @entrypoint(checkpointer=checkpointer) async def graph(state: dict) -> dict: fut_foo = foo(state) fut_bar = bar(await fut_foo) From 09ca96471402d8594b3e9ca71e06d15a71d9868d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 21 Nov 2024 08:55:07 -0800 Subject: [PATCH 14/26] Wire up retry policy --- libs/langgraph/langgraph/func/__init__.py | 2 +- libs/langgraph/langgraph/pregel/algo.py | 18 ++++++++++++++---- libs/langgraph/langgraph/pregel/runner.py | 12 ++++++++++-- libs/langgraph/langgraph/types.py | 6 ++++-- 4 files changed, 29 insertions(+), 9 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 69e9b6c5d..4db234b81 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -48,7 +48,7 @@ def task( Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]], ]: def _task(func: Callable[P, T]) -> Callable[P, concurrent.futures.Future[T]]: - return update_wrapper(partial(call, func), func) + return update_wrapper(partial(call, func, retry=retry), func) return _task diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 330d46449..4e17dd39d 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -65,7 +65,13 @@ from langgraph.pregel.log import logger from langgraph.pregel.manager import ChannelsManager from langgraph.pregel.read import PregelNode from langgraph.store.base import BaseStore -from langgraph.types import All, LoopProtocol, PregelExecutableTask, PregelTask +from langgraph.types import ( + All, + LoopProtocol, + PregelExecutableTask, + PregelTask, + RetryPolicy, +) from langgraph.utils.config import merge_configs, patch_config GetNextVersion = Callable[[Optional[V], BaseChannel], V] @@ -100,14 +106,18 @@ class PregelTaskWrites(NamedTuple): class Call: - __slots__ = ("func", "input") + __slots__ = ("func", "input", "retry") func: Callable input: Any + retry: Optional[RetryPolicy] - def __init__(self, func: Callable, input: Any) -> None: + def __init__( + self, func: Callable, input: Any, *, retry: Optional[RetryPolicy] + ) -> None: self.func = func self.input = input + self.retry = retry def should_interrupt( @@ -545,7 +555,7 @@ def prepare_single_task( }, ), triggers, - None, + call.retry, None, task_id, task_path[:3], diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 8decbedec..8e3f52ace 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -143,8 +143,12 @@ class PregelRunner: task: PregelExecutableTask, func: Callable[[Any], Union[Awaitable[Any], Any]], input: Any, + *, + retry: Optional[RetryPolicy] = None, ) -> concurrent.futures.Future[Any]: - (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)]) + (fut,) = writer( + task, [(PUSH, None)], calls=[Call(func, input, retry=retry)] + ) assert fut is not None, "writer did not return a future for call" return fut @@ -320,8 +324,12 @@ class PregelRunner: task: PregelExecutableTask, func: Callable[[Any], Union[Awaitable[Any], Any]], input: Any, + *, + retry: Optional[RetryPolicy] = None, ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]: - (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)]) + (fut,) = writer( + task, [(PUSH, None)], calls=[Call(func, input, retry=retry)] + ) assert fut is not None, "writer did not return a future for call" if asyncio.iscoroutinefunction(func): return fut diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 5ce64b52b..c3bf175a2 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -372,12 +372,14 @@ T = TypeVar("T") def call( - func: str | Callable[P, T], *args: P.args, **kwargs: P.kwargs + func: str | Callable[P, T], + *args: P.args, + retry: Optional[RetryPolicy] = None, ) -> concurrent.futures.Future[T]: from langgraph.constants import CONFIG_KEY_CALL from langgraph.utils.config import get_configurable conf = get_configurable() impl = conf[CONFIG_KEY_CALL] - fut = impl(func, *args, **kwargs) + fut = impl(func, *args, retry=retry) return fut From 2fe38f394021c2e3f4d7169c791a1ee8d84054b9 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 21 Nov 2024 09:07:19 -0800 Subject: [PATCH 15/26] Fix get_state --- libs/langgraph/langgraph/func/__init__.py | 1 + libs/langgraph/tests/test_pregel.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 4db234b81..d990ae4a5 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -71,6 +71,7 @@ def entrypoint( channels={START: EphemeralValue(Any), END: LastValue(Any, END)}, input_channels=START, output_channels=END, + stream_channels=END, stream_mode="updates", checkpointer=checkpointer, store=store, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 8de87ef0e..cb7998430 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -2046,6 +2046,8 @@ def test_imp_stream_order( {"graph": {"a": "0foobarbaz", "c": "something else"}}, ] + assert graph.get_state(thread1).values == {"a": "0foobarbaz", "c": "something else"} + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_dedupe_on_resume( From 4c6323c585dc0d2e6999b5fe73a9699366004292 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Dec 2024 13:54:01 -0800 Subject: [PATCH 16/26] Lint --- libs/langgraph/langgraph/utils/future.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/utils/future.py b/libs/langgraph/langgraph/utils/future.py index d7b8718f6..eaad8e64d 100644 --- a/libs/langgraph/langgraph/utils/future.py +++ b/libs/langgraph/langgraph/utils/future.py @@ -17,7 +17,7 @@ def _get_loop(fut: asyncio.Future) -> asyncio.AbstractEventLoop: return fut._loop -def _convert_future_exc(exc): +def _convert_future_exc(exc: BaseException) -> BaseException: exc_class = type(exc) if exc_class is concurrent.futures.CancelledError: return asyncio.CancelledError(*exc.args) @@ -29,7 +29,10 @@ def _convert_future_exc(exc): return exc -def _set_concurrent_future_state(concurrent, source): +def _set_concurrent_future_state( + concurrent: concurrent.futures.Future, + source: AnyFuture, +) -> None: """Copy state from a future to a concurrent.futures.Future.""" assert source.done() if source.cancelled(): @@ -44,7 +47,7 @@ def _set_concurrent_future_state(concurrent, source): concurrent.set_result(result) -def _copy_future_state(source, dest): +def _copy_future_state(source: AnyFuture, dest: asyncio.Future) -> None: """Internal helper to copy state from another Future. The other Future may be a concurrent.futures.Future. @@ -82,20 +85,20 @@ def _chain_future(source: AnyFuture, destination: AnyFuture) -> None: source_loop = _get_loop(source) if asyncio.isfuture(source) else None dest_loop = _get_loop(destination) if asyncio.isfuture(destination) else None - def _set_state(future, other): + def _set_state(future: AnyFuture, other: AnyFuture) -> None: if asyncio.isfuture(future): _copy_future_state(other, future) else: _set_concurrent_future_state(future, other) - def _call_check_cancel(destination): + def _call_check_cancel(destination: AnyFuture) -> None: if destination.cancelled(): if source_loop is None or source_loop is dest_loop: source.cancel() else: source_loop.call_soon_threadsafe(source.cancel) - def _call_set_state(source): + def _call_set_state(source: AnyFuture) -> None: if destination.cancelled() and dest_loop is not None and dest_loop.is_closed(): return if dest_loop is None or dest_loop is source_loop: @@ -109,7 +112,7 @@ def _chain_future(source: AnyFuture, destination: AnyFuture) -> None: source.add_done_callback(_call_set_state) -def chain_future(source: AnyFuture, destination: AnyFuture) -> None: +def chain_future(source: AnyFuture, destination: concurrent.futures.Future) -> None: # adapted from asyncio.run_coroutine_threadsafe try: _chain_future(source, destination) From ec7bbe14b214f8b0e145899a866bb3254c4b95c7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Dec 2024 13:54:55 -0800 Subject: [PATCH 17/26] Lint --- libs/langgraph/langgraph/pregel/runner.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 8e3f52ace..6e68df66a 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -184,9 +184,7 @@ class PregelRunner: # add waiter task if requested if get_waiter is not None: futures[get_waiter()] = None - # execute tasks, and wait for one to fail or all to finish. - # each task is independent from all other concurrent tasks - # yield updates/debug output as each task finishes + # schedule tasks for t in tasks: if not t.writes: fut = self.submit( @@ -201,6 +199,9 @@ class PregelRunner: ) fut.add_done_callback(partial(self.commit, t)) futures[fut] = t + # execute tasks, and wait for one to fail or all to finish. + # each task is independent from all other concurrent tasks + # yield updates/debug output as each task finishes end_time = timeout + time.monotonic() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): done, inflight = concurrent.futures.wait( @@ -334,7 +335,7 @@ class PregelRunner: if asyncio.iscoroutinefunction(func): return fut # adapted from asyncio.run_coroutine_threadsafe - sfut = concurrent.futures.Future() + sfut: concurrent.futures.Future = concurrent.futures.Future() loop.call_soon_threadsafe(chain_future, fut, sfut) return sfut @@ -372,9 +373,7 @@ class PregelRunner: # add waiter task if requested if get_waiter is not None: futures[get_waiter()] = None - # execute tasks, and wait for one to fail or all to finish. - # each task is independent from all other concurrent tasks - # yield updates/debug output as each task finishes + # schedule tasks for t in tasks: if not t.writes: fut = cast( @@ -395,6 +394,9 @@ class PregelRunner: ) fut.add_done_callback(partial(self.commit, t)) futures[fut] = t + # execute tasks, and wait for one to fail or all to finish. + # each task is independent from all other concurrent tasks + # yield updates/debug output as each task finishes end_time = timeout + loop.time() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): done, inflight = await asyncio.wait( From 40d16593c7c1260dba8f7f97cd54de9a8720d7e5 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Dec 2024 13:59:21 -0800 Subject: [PATCH 18/26] Lint --- libs/langgraph/langgraph/func/__init__.py | 18 +++++++++++++++++- libs/langgraph/langgraph/types.py | 22 +--------------------- 2 files changed, 18 insertions(+), 22 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index d990ae4a5..d608f41f7 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -23,12 +23,28 @@ from langgraph.pregel.call import get_runnable_for_func from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import RetryPolicy, call +from langgraph.types import RetryPolicy P = ParamSpec("P") +P1 = TypeVar("P1") T = TypeVar("T") +def call( + func: Callable[[P1], T], + input: P1, + *, + retry: Optional[RetryPolicy] = None, +) -> concurrent.futures.Future[T]: + from langgraph.constants import CONFIG_KEY_CALL + from langgraph.utils.config import get_configurable + + conf = get_configurable() + impl = conf[CONFIG_KEY_CALL] + fut = impl(func, input, retry=retry) + return fut + + @overload def task( *, retry: Optional[RetryPolicy] = None diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index c3bf175a2..67c7e53f8 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,5 +1,3 @@ -import concurrent -import concurrent.futures import dataclasses import sys from collections import deque @@ -22,7 +20,7 @@ from typing import ( ) from langchain_core.runnables import Runnable, RunnableConfig -from typing_extensions import ParamSpec, Self +from typing_extensions import Self from langgraph.checkpoint.base import ( BaseCheckpointSaver, @@ -365,21 +363,3 @@ def interrupt(value: Any) -> Any: ), ) ) - - -P = ParamSpec("P") -T = TypeVar("T") - - -def call( - func: str | Callable[P, T], - *args: P.args, - retry: Optional[RetryPolicy] = None, -) -> concurrent.futures.Future[T]: - from langgraph.constants import CONFIG_KEY_CALL - from langgraph.utils.config import get_configurable - - conf = get_configurable() - impl = conf[CONFIG_KEY_CALL] - fut = impl(func, *args, retry=retry) - return fut From 2b77fdabee4b8de8bdcf88c304664a1bcbe978b7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Dec 2024 14:01:26 -0800 Subject: [PATCH 19/26] Lint --- libs/langgraph/langgraph/func/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index d608f41f7..2dda24754 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -8,12 +8,13 @@ from typing import ( Awaitable, Callable, Optional, - ParamSpec, TypeVar, Union, overload, ) +from typing_extensions import ParamSpec + from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import BaseCheckpointSaver From 007d7e72b124066dc1f96d2b203b0002b13d8f95 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Dec 2024 15:04:55 -0800 Subject: [PATCH 20/26] Add test for cancellation --- libs/langgraph/langgraph/pregel/algo.py | 3 +- libs/langgraph/langgraph/pregel/runner.py | 24 +++++---- libs/langgraph/tests/test_pregel_async.py | 66 +++++++++++++++++++++++ 3 files changed, 83 insertions(+), 10 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 4e17dd39d..b41424fe9 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -43,6 +43,7 @@ from langgraph.constants import ( CONFIG_KEY_TASK_ID, CONFIG_KEY_WRITES, EMPTY_SEQ, + ERROR, INTERRUPT, NO_WRITES, NS_END, @@ -270,7 +271,7 @@ def apply_writes( pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list) for task in tasks: for chan, val in task.writes: - if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN): + if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR): pass elif chan == TASKS: # TODO: remove branch in 1.0 checkpoint["pending_sends"].append(val) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 6e68df66a..41c9e7e69 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -441,7 +441,12 @@ class PregelRunner: ) -> None: if fut is not None: exception = _exception(fut) - if exception: + if isinstance(exception, asyncio.CancelledError): + # for cancelled tasks, also save error in task, + # so loop can finish super-step + task.writes.append((ERROR, exception)) + self.put_writes(task.id, task.writes) + elif exception: if isinstance(exception, GraphInterrupt): # save interrupt to checkpointer if interrupts := [(INTERRUPT, i) for i in exception.args[0]]: @@ -472,11 +477,12 @@ def _should_stop_others( GraphInterrupts are not considered failures.""" for fut in done: if fut.cancelled(): - return True - if exc := fut.exception(): - return not isinstance(exc, GraphBubbleUp) - else: - return False + continue + elif exc := fut.exception(): + if not isinstance(exc, GraphBubbleUp): + return True + + return False def _exception( @@ -502,7 +508,9 @@ def _panic_or_proceed( done: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set() inflight: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set() for fut in futs: - if fut.done(): + if fut.cancelled(): + continue + elif fut.done(): done.add(fut) else: inflight.add(fut) @@ -515,8 +523,6 @@ def _panic_or_proceed( # raise the exception if panic: raise exc - else: - return if inflight: # if we got here means we timed out while inflight: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index c4e5174f6..9f568d5c3 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2648,6 +2648,10 @@ async def test_send_sequences(checkpointer_name: str) -> None: ] +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_task(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -2690,6 +2694,64 @@ async def test_imp_task(checkpointer_name: str) -> None: assert mapper_calls == 2 +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_imp_task_cancel(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + mapper_calls = 0 + mapper_cancels = 0 + + @task() + async def mapper(input: int) -> str: + nonlocal mapper_calls, mapper_cancels + mapper_calls += 1 + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + mapper_cancels += 1 + raise + return str(input) * 2 + + @entrypoint(checkpointer=checkpointer) + async def graph(input: list[int]) -> list[str]: + futures = [mapper(i) for i in input] + await asyncio.sleep(0.1) + futures.pop().cancel() # cancel one + mapped = await asyncio.gather(*futures) + answer = interrupt("question") + return [m + answer for m in mapped] + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c async for c in graph.astream([0, 1], thread1)] == [ + {"mapper": "00"}, + { + "__interrupt__": ( + Interrupt( + value="question", + resumable=True, + ns=[AnyStr("graph:")], + when="during", + ), + ) + }, + ] + assert mapper_calls == 2 + assert mapper_cancels == 1 + + assert await graph.ainvoke(Command(resume="answer"), thread1) == [ + "00answer", + ] + assert mapper_calls == 3 + assert mapper_cancels == 2 + + +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_sync_from_async(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -2722,6 +2784,10 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None: ] +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_stream_order(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: From 4e8f4ce440b31336ab48d63ee4522c5c7ef5ebaf Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Dec 2024 15:33:41 -0800 Subject: [PATCH 21/26] Update --- libs/scheduler-kafka/tests/test_subgraph.py | 6 ++++++ libs/scheduler-kafka/tests/test_subgraph_sync.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index 4ab92676c..1e1f1e396 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -191,6 +191,7 @@ async def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, @@ -257,6 +258,7 @@ async def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, @@ -353,6 +355,7 @@ async def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, @@ -459,6 +462,7 @@ async def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, @@ -520,6 +524,7 @@ async def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, @@ -637,6 +642,7 @@ async def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index 5fa43998a..210312b3b 100644 --- a/libs/scheduler-kafka/tests/test_subgraph_sync.py +++ b/libs/scheduler-kafka/tests/test_subgraph_sync.py @@ -190,6 +190,7 @@ def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, @@ -256,6 +257,7 @@ def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_store": None, "__pregel_dedupe_tasks": True, @@ -352,6 +354,7 @@ def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_store": None, @@ -457,6 +460,7 @@ def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_store": None, @@ -518,6 +522,7 @@ def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_store": None, @@ -635,6 +640,7 @@ def test_subgraph_w_interrupt( "__pregel_delegate": False, "__pregel_read": None, "__pregel_send": None, + "__pregel_call": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, From eb593d47ddb4226d40fbd1ff3630a85511264ff7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Dec 2024 16:37:02 -0800 Subject: [PATCH 22/26] Fix writes for task being saved against next checkpoint id --- libs/langgraph/langgraph/pregel/loop.py | 9 ++++----- libs/langgraph/langgraph/pregel/runner.py | 4 ++++ 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index cb141c4a1..678e355ab 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -290,16 +290,15 @@ class PregelLoop(LoopProtocol): if self.checkpointer_put_writes is not None: self.submit( self.checkpointer_put_writes, - { - **self.checkpoint_config, - CONF: { - **self.checkpoint_config[CONF], + patch_configurable( + self.checkpoint_config, + { CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( CONFIG_KEY_CHECKPOINT_NS, "" ), CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"], }, - }, + ), writes, task_id, ) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 41c9e7e69..5ba0209c8 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -423,6 +423,10 @@ class PregelRunner: break # give control back to the caller yield + # wait for pending done callbacks + # if a 2nd future finishes while `wait` is returning, it's possible + # that done callbacks for the 2nd future aren't called until next tick + await asyncio.sleep(0) # cancel waiter task for fut in futures: fut.cancel() From e1f65012e6a76b0cf1f40deecb92e1ee74b7ce4e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Dec 2024 17:11:24 -0800 Subject: [PATCH 23/26] Fix --- libs/langgraph/tests/test_pregel_async.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 9f568d5c3..f43e76d31 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9809,14 +9809,14 @@ async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None ), (FloatBetween(0.2, 0.4), ((), {"outer_1": {"my_key": " and parallel"}})), ( - FloatBetween(0.5, 0.7), + FloatBetween(0.5, 0.8), ( (AnyStr("inner:"),), {"inner_2": {"my_key": " and there", "my_other_key": "got here"}}, ), ), - (FloatBetween(0.5, 0.7), ((), {"inner": {"my_key": "got here and there"}})), - (FloatBetween(0.5, 0.7), ((), {"outer_2": {"my_key": " and back again"}})), + (FloatBetween(0.5, 0.8), ((), {"inner": {"my_key": "got here and there"}})), + (FloatBetween(0.5, 0.8), ((), {"outer_2": {"my_key": " and back again"}})), ] From 9733db03c537a82ce5b2a27d607962589a1b35b9 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 4 Dec 2024 12:04:12 -0800 Subject: [PATCH 24/26] Wait until next tick to start send task --- libs/langgraph/langgraph/pregel/executor.py | 23 ++++++++++++++++++++- libs/langgraph/langgraph/pregel/runner.py | 5 +++++ libs/langgraph/tests/test_pregel.py | 1 + 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 70aea29e3..46a4e6036 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -1,6 +1,7 @@ import asyncio import concurrent.futures import sys +import time from contextlib import ExitStack from contextvars import copy_context from types import TracebackType @@ -34,6 +35,7 @@ class Submit(Protocol[P, T]): __name__: Optional[str] = None, __cancel_on_exit__: bool = False, __reraise_on_exit__: bool = True, + __next_tick__: bool = False, **kwargs: P.kwargs, ) -> concurrent.futures.Future[T]: ... @@ -58,9 +60,13 @@ class BackgroundExecutor(ContextManager): __name__: Optional[str] = None, # currently not used in sync version __cancel_on_exit__: bool = False, # for sync, can cancel only if not started __reraise_on_exit__: bool = True, + __next_tick__: bool = False, **kwargs: P.kwargs, ) -> concurrent.futures.Future[T]: - task = self.executor.submit(fn, *args, **kwargs) + if __next_tick__: + task = self.executor.submit(next_tick, fn, *args, **kwargs) + else: + task = self.executor.submit(fn, *args, **kwargs) self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__) task.add_done_callback(self.done) return task @@ -137,11 +143,14 @@ class AsyncBackgroundExecutor(AsyncContextManager): __name__: Optional[str] = None, __cancel_on_exit__: bool = False, __reraise_on_exit__: bool = True, + __next_tick__: bool = False, **kwargs: P.kwargs, ) -> asyncio.Task[T]: coro = cast(Coroutine[None, None, T], fn(*args, **kwargs)) if self.semaphore: coro = gated(self.semaphore, coro) + if __next_tick__: + coro = anext_tick(coro) if self.context_not_supported: task = self.loop.create_task(coro, name=__name__) else: @@ -197,3 +206,15 @@ async def gated(semaphore: asyncio.Semaphore, coro: Coroutine[None, None, T]) -> """A coroutine that waits for a semaphore before running another coroutine.""" async with semaphore: return await coro + + +def next_tick(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: + """A function that yields control to other threads before running another function.""" + time.sleep(0) + return fn(*args, **kwargs) + + +async def anext_tick(coro: Coroutine[None, None, T]) -> T: + """A coroutine that yields control to event loop before running another coroutine.""" + await asyncio.sleep(0) + return await coro diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 5ba0209c8..e60ae599b 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -133,6 +133,7 @@ class PregelRunner: CONFIG_KEY_CALL: partial(call, next_task), }, __reraise_on_exit__=reraise, + __next_tick__=True, ) fut.add_done_callback(partial(self.commit, next_task)) futures[fut] = next_task @@ -228,6 +229,10 @@ class PregelRunner: break # give control back to the caller yield + # wait for pending done callbacks + # if a 2nd future finishes while `wait` is returning, it's possible + # that done callbacks for the 2nd future aren't called until next tick + time.sleep(0) # panic on failure or timeout _panic_or_proceed( done_futures.union(f for f, t in futures.items() if t is not None), diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index cb7998430..9948be87e 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5837,6 +5837,7 @@ def test_state_graph_packets( @tool() def search_api(query: str) -> str: """Searches the API for the query.""" + time.sleep(0.1) return f"result for {query}" tools = [search_api] From de86a46b3d17091f2862cab997315d0effb3bdb8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 4 Dec 2024 12:05:32 -0800 Subject: [PATCH 25/26] Comment --- libs/langgraph/langgraph/pregel/runner.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index e60ae599b..e680518a5 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -133,6 +133,8 @@ class PregelRunner: CONFIG_KEY_CALL: partial(call, next_task), }, __reraise_on_exit__=reraise, + # starting a new task in the next tick ensures + # updates from this tick are committed/streamed first __next_tick__=True, ) fut.add_done_callback(partial(self.commit, next_task)) @@ -319,6 +321,9 @@ class PregelRunner: __name__=t.name, __cancel_on_exit__=True, __reraise_on_exit__=reraise, + # starting a new task in the next tick ensures + # updates from this tick are committed/streamed first + __next_tick__=True, ), ) fut.add_done_callback(partial(self.commit, next_task)) From 2fa2469967148e9cd10a34bfeed46adaad265786 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 4 Dec 2024 15:18:15 -0800 Subject: [PATCH 26/26] Update --- libs/langgraph/langgraph/pregel/algo.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index b41424fe9..3adea073a 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -553,6 +553,13 @@ def prepare_single_task( }, CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, + CONFIG_KEY_WRITES: [ + w + for w in pending_writes + + configurable.get(CONFIG_KEY_WRITES, []) + if w[0] in (NULL_TASK_ID, task_id) + ], + CONFIG_KEY_SCRATCHPAD: {}, }, ), triggers,