From 3d48526c16b62d88c4a7b21d2c76baeea44b1312 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Tue, 18 Mar 2025 15:14:58 -0700 Subject: [PATCH] Methods are weak --- libs/langgraph/Makefile | 4 +- libs/langgraph/langgraph/pregel/__init__.py | 65 +-- libs/langgraph/langgraph/pregel/runner.py | 428 ++++++++++++-------- 3 files changed, 300 insertions(+), 197 deletions(-) diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index 8974fcd32..b1946b045 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -58,9 +58,11 @@ WORKERS ?= auto XDIST_ARGS := $(if $(WORKERS),-n $(WORKERS) --dist worksteal,) MAXFAIL ?= MAXFAIL_ARGS := $(if $(MAXFAIL),--maxfail $(MAXFAIL),) +# Add an '-x' if xdist is enabled +XDIST_ARGS := $(if $(WORKERS),-x $(XDIST_ARGS),) test_watch: - make start-postgres && poetry run ptw . -- --ff -vv -x $(XDIST_ARGS) $(MAXFAIL_ARGS) --snapshot-update --tb short $(TEST); \ + make start-postgres && poetry run ptw . -- --ff -vv $(XDIST_ARGS) $(MAXFAIL_ARGS) --snapshot-update --tb short $(TEST); \ EXIT_CODE=$$?; \ make stop-postgres; \ exit $$EXIT_CODE diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 362c1459c..f6479422c 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -4,6 +4,7 @@ import asyncio import concurrent import concurrent.futures import queue +import weakref from collections import deque from functools import partial from typing import ( @@ -797,9 +798,11 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, + checkpointer=( + self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None + ), manager=None, ) # get the subgraphs @@ -911,9 +914,11 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, + checkpointer=( + self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None + ), manager=None, ) # get the subgraphs @@ -1225,9 +1230,11 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, + checkpointer=( + self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None + ), manager=None, ) # apply null writes @@ -1321,9 +1328,11 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, + checkpointer=( + self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None + ), manager=None, ) # apply null writes @@ -1508,9 +1517,11 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, + checkpointer=( + self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None + ), manager=None, ) # apply null writes @@ -1604,9 +1615,11 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer - if isinstance(self.checkpointer, BaseCheckpointSaver) - else None, + checkpointer=( + self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None + ), manager=None, ) # apply null writes @@ -1989,9 +2002,11 @@ class Pregel(PregelProtocol): ) as loop: # create runner runner = PregelRunner( - submit=config[CONF].get(CONFIG_KEY_RUNNER_SUBMIT, loop.submit), - put_writes=loop.put_writes, - schedule_task=loop.accept_push, + submit=config[CONF].get( + CONFIG_KEY_RUNNER_SUBMIT, weakref.WeakMethod(loop.submit) + ), + put_writes=weakref.WeakMethod(loop.put_writes), + schedule_task=weakref.WeakMethod(loop.accept_push), node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), ) # enable subgraph streaming @@ -2280,9 +2295,11 @@ class Pregel(PregelProtocol): ) as loop: # create runner runner = PregelRunner( - submit=config[CONF].get(CONFIG_KEY_RUNNER_SUBMIT, loop.submit), - put_writes=loop.put_writes, - schedule_task=loop.accept_push, + submit=config[CONF].get( + CONFIG_KEY_RUNNER_SUBMIT, weakref.WeakMethod(loop.submit) + ), + put_writes=weakref.WeakMethod(loop.put_writes), + schedule_task=weakref.WeakMethod(loop.accept_push), use_astream=do_stream is not None, node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), ) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 4a0b226e8..faf45e3d9 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -2,6 +2,7 @@ import asyncio import concurrent.futures import threading import time +import weakref from functools import partial from typing import ( Any, @@ -46,7 +47,9 @@ E = TypeVar("E", threading.Event, asyncio.Event) class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]): event: E - callback: Callable[[PregelExecutableTask, Optional[BaseException]], None] + callback: weakref.ref[ + Callable[[PregelExecutableTask, Optional[BaseException]], None] + ] counter: int done: set[F] lock: threading.Lock @@ -54,7 +57,9 @@ class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]): def __init__( self, event: E, - callback: Callable[[PregelExecutableTask, Optional[BaseException]], None], + callback: weakref.ref[ + Callable[[PregelExecutableTask, Optional[BaseException]], None] + ], future_type: Type[F], # used for generic typing, newer py supports FutureDict[...](...) ) -> None: @@ -83,7 +88,7 @@ class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]): fut: F, ) -> None: try: - self.callback(task, _exception(fut)) + self.callback()(task, _exception(fut)) # type: ignore[misc] finally: with self.lock: self.done.add(fut) @@ -100,10 +105,13 @@ class PregelRunner: def __init__( self, *, - submit: Submit, - put_writes: Callable[[str, Sequence[tuple[str, Any]]], None], - schedule_task: Callable[ - [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] + submit: weakref.ref[Submit], + put_writes: weakref.ref[Callable[[str, Sequence[tuple[str, Any]]], None]], + schedule_task: weakref.ref[ + Callable[ + [PregelExecutableTask, int, Optional[Call]], + Optional[PregelExecutableTask], + ] ], use_astream: bool = False, node_finished: Optional[Callable[[str], None]] = None, @@ -123,74 +131,9 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: - def call( - task: PregelExecutableTask, - func: Callable[[Any], Union[Awaitable[Any], Any]], - input: Any, - *, - retry: Optional[RetryPolicy] = None, - callbacks: Callbacks = None, - ) -> concurrent.futures.Future[Any]: - if asyncio.iscoroutinefunction(func): - raise RuntimeError("In an sync context async tasks cannot be called") - - fut: Optional[concurrent.futures.Future] = None - # schedule PUSH tasks, collect futures - scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] - # schedule the next task, if the callback returns one - if next_task := self.schedule_task( - task, - scratchpad.call_counter(), - Call(func, input, retry=retry, callbacks=callbacks), - ): - if fut := next( - ( - f - for f, t in futures.items() - if t is not None and t == next_task.id - ), - None, - ): - # if the parent task was retried, - # the next task might already be running - pass - elif next_task.writes: - # if it already ran, return the result - fut = concurrent.futures.Future() - ret = next((v for c, v in next_task.writes if c == RETURN), MISSING) - if ret is not MISSING: - fut.set_result(ret) - elif exc := next( - (v for c, v in next_task.writes if c == ERROR), None - ): - fut.set_exception( - exc if isinstance(exc, BaseException) else Exception(exc) - ) - else: - fut.set_result(None) - else: - # schedule the next task - fut = self.submit( - run_with_retry, - next_task, - retry_policy, - configurable={ - 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, - ) - futures[fut] = next_task - - # return a chained future to ensure commit() callback is called - # before the returned future is resolved, to ensure stream order etc - return chain_future(fut, concurrent.futures.Future()) - tasks = tuple(tasks) futures = FuturesDict( - callback=self.commit, + callback=weakref.WeakMethod(self.commit), event=threading.Event(), future_type=concurrent.futures.Future, ) @@ -204,7 +147,15 @@ class PregelRunner: t, retry_policy, configurable={ - CONFIG_KEY_CALL: partial(call, t), + CONFIG_KEY_CALL: partial( + _call, + t, + retry=retry_policy, + futures=weakref.ref(futures), + schedule_task=self.schedule_task, + submit=self.submit, + reraise=reraise, + ), }, ) self.commit(t, None) @@ -227,12 +178,20 @@ class PregelRunner: # schedule tasks for t in tasks: if not t.writes: - fut = self.submit( + fut = self.submit()( # type: ignore[misc] run_with_retry, t, retry_policy, configurable={ - CONFIG_KEY_CALL: partial(call, t), + CONFIG_KEY_CALL: partial( + _call, + t, + retry=retry_policy, + futures=weakref.ref(futures), + schedule_task=self.schedule_task, + submit=self.submit, + reraise=reraise, + ), }, __reraise_on_exit__=reraise, ) @@ -284,98 +243,10 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, ) -> AsyncIterator[None]: - def call( - task: PregelExecutableTask, - func: Callable[[Any], Union[Awaitable[Any], Any]], - input: Any, - *, - retry: Optional[RetryPolicy] = None, - callbacks: Callbacks = None, - ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]: - fut: Optional[asyncio.Future] = None - # schedule PUSH tasks, collect futures - scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] - # schedule the next task, if the callback returns one - if next_task := self.schedule_task( - task, - scratchpad.call_counter(), - Call(func, input, retry=retry, callbacks=callbacks), - ): - # if the parent task was retried, - # the next task might already be running - if fut := next( - ( - f - for f, t in futures.items() - if t is not None and t == next_task.id - ), - None, - ): - # if the parent task was retried, - # the next task might already be running - pass - elif next_task.writes: - # if it already ran, return the result - fut = asyncio.Future(loop=loop) - ret = next((v for c, v in next_task.writes if c == RETURN), MISSING) - if ret is not MISSING: - fut.set_result(ret) - elif exc := next( - (v for c, v in next_task.writes if c == ERROR), None - ): - fut.set_exception( - exc if isinstance(exc, BaseException) else Exception(exc) - ) - else: - fut.set_result(None) - 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_CALL: partial(call, next_task), - }, - __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, - ), - ) - futures[fut] = next_task - # return a chained future to ensure commit() callback is called - # before the returned future is resolved, to ensure stream order etc - try: - in_async = asyncio.current_task() is not None - except RuntimeError: - in_async = False - # if in async context return an async future - # otherwise return a chained sync future - if in_async: - if isinstance(fut, asyncio.Task): - sfut: Union[asyncio.Future[Any], concurrent.futures.Future[Any]] = ( - asyncio.Future(loop=loop) - ) - loop.call_soon_threadsafe(chain_future, fut, sfut) - return sfut - else: - # already wrapped in a future - return fut - else: - sfut = concurrent.futures.Future() - loop.call_soon_threadsafe(chain_future, fut, sfut) - return sfut - loop = asyncio.get_event_loop() tasks = tuple(tasks) futures = FuturesDict( - callback=self.commit, + callback=weakref.WeakMethod(self.commit), event=asyncio.Event(), future_type=asyncio.Future, ) @@ -390,7 +261,17 @@ class PregelRunner: retry_policy, stream=self.use_astream, configurable={ - CONFIG_KEY_CALL: partial(call, t), + CONFIG_KEY_CALL: partial( + _acall, + t, + stream=self.use_astream, + retry=retry_policy, + futures=weakref.ref(futures), + schedule_task=self.schedule_task, + submit=self.submit, + reraise=reraise, + loop=loop, + ), }, ) self.commit(t, None) @@ -415,13 +296,23 @@ class PregelRunner: if not t.writes: fut = cast( asyncio.Future, - self.submit( + self.submit()( # type: ignore[misc] arun_with_retry, t, retry_policy, stream=self.use_astream, configurable={ - CONFIG_KEY_CALL: partial(call, t), + CONFIG_KEY_CALL: partial( + _acall, + t, + retry=retry_policy, + stream=self.use_astream, + futures=weakref.ref(futures), + schedule_task=self.schedule_task, + submit=self.submit, + reraise=reraise, + loop=loop, + ), }, __name__=t.name, __cancel_on_exit__=True, @@ -481,7 +372,7 @@ class PregelRunner: # 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) + self.put_writes()(task.id, task.writes) # type: ignore[misc] elif exception: if isinstance(exception, GraphInterrupt): # save interrupt to checkpointer @@ -489,12 +380,12 @@ class PregelRunner: writes = [(INTERRUPT, exception.args[0])] if resumes := [w for w in task.writes if w[0] == RESUME]: writes.extend(resumes) - self.put_writes(task.id, writes) + self.put_writes()(task.id, writes) # type: ignore[misc] elif isinstance(exception, GraphBubbleUp): raise exception else: # save error to checkpointer - self.put_writes(task.id, [(ERROR, exception)]) + self.put_writes()(task.id, [(ERROR, exception)]) # type: ignore[misc] else: if self.node_finished and ( task.config is None or TAG_HIDDEN not in task.config.get("tags", []) @@ -504,7 +395,7 @@ class PregelRunner: # add no writes marker task.writes.append((NO_WRITES, None)) # save task writes to checkpointer - self.put_writes(task.id, task.writes) + self.put_writes()(task.id, task.writes) # type: ignore[misc] def _should_stop_others( @@ -575,3 +466,196 @@ def _panic_or_proceed( inflight.pop().cancel() # raise timeout error raise timeout_exc_cls("Timed out") + + +def _call( + task: PregelExecutableTask, + func: Callable[[Any], Union[Awaitable[Any], Any]], + input: Any, + *, + retry: Optional[RetryPolicy] = None, + callbacks: Callbacks = None, + futures: weakref.ref[FuturesDict], + schedule_task: weakref.ref[ + Callable[ + [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] + ] + ], + submit: weakref.ref[Submit], + reraise: bool, +) -> concurrent.futures.Future[Any]: + if asyncio.iscoroutinefunction(func): + raise RuntimeError("In an sync context async tasks cannot be called") + + fut: Optional[concurrent.futures.Future] = None + # schedule PUSH tasks, collect futures + scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] + # schedule the next task, if the callback returns one + if next_task := schedule_task()( # type: ignore[misc] + task, + scratchpad.call_counter(), + Call(func, input, retry=retry, callbacks=callbacks), + ): + if fut := next( + ( + f + for f, t in futures().items() # type: ignore[union-attr] + if t is not None and t == next_task.id + ), + None, + ): + # if the parent task was retried, + # the next task might already be running + pass + elif next_task.writes: + # if it already ran, return the result + fut = concurrent.futures.Future() + ret = next((v for c, v in next_task.writes if c == RETURN), MISSING) + if ret is not MISSING: + fut.set_result(ret) + elif exc := next((v for c, v in next_task.writes if c == ERROR), None): + fut.set_exception( + exc if isinstance(exc, BaseException) else Exception(exc) + ) + else: + fut.set_result(None) + else: + # schedule the next task + fut = submit()( # type: ignore[misc] + run_with_retry, + next_task, + retry, + configurable={ + CONFIG_KEY_CALL: partial( + _call, + next_task, + futures=futures, + retry=retry, + callbacks=callbacks, + schedule_task=schedule_task, + submit=submit, + reraise=reraise, + ), + }, + __reraise_on_exit__=reraise, + # starting a new task in the next tick ensures + # updates from this tick are committed/streamed first + __next_tick__=True, + ) + futures()[fut] = next_task # type: ignore[index] + fut = cast(Union[asyncio.Future, concurrent.futures.Future], fut) + # return a chained future to ensure commit() callback is called + # before the returned future is resolved, to ensure stream order etc + return chain_future(fut, concurrent.futures.Future()) + + +def _acall( + task: PregelExecutableTask, + func: Callable[[Any], Union[Awaitable[Any], Any]], + input: Any, + *, + retry: Optional[RetryPolicy] = None, + callbacks: Callbacks = None, + # injected dependencies + futures: weakref.ref[FuturesDict], + schedule_task: weakref.ref[ + Callable[ + [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] + ] + ], + submit: weakref.ref[Submit], + loop: asyncio.AbstractEventLoop, + reraise: bool = False, + stream: bool = False, +) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]: + fut: Optional[asyncio.Future] = None + # schedule PUSH tasks, collect futures + scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] + # schedule the next task, if the callback returns one + if next_task := schedule_task()( # type: ignore[misc] + task, + scratchpad.call_counter(), + Call(func, input, retry=retry, callbacks=callbacks), + ): + if fut := next( + ( + f + for f, t in futures().items() # type: ignore[union-attr] + if t is not None and t == next_task.id + ), + None, + ): + # if the parent task was retried, + # the next task might already be running + + pass + elif next_task.writes: + # if it already ran, return the result + fut = asyncio.Future(loop=loop) + ret = next((v for c, v in next_task.writes if c == RETURN), MISSING) + if ret is not MISSING: + fut.set_result(ret) + elif exc := next((v for c, v in next_task.writes if c == ERROR), None): + fut.set_exception( + exc if isinstance(exc, BaseException) else Exception(exc) + ) + else: + fut.set_result(None) + futures()[fut] = next_task # type: ignore[index] + else: + # schedule the next task + fut = cast( + asyncio.Future, + submit()( # type: ignore[misc] + arun_with_retry, + next_task, + retry, + stream=stream, + configurable={ + CONFIG_KEY_CALL: partial( + _acall, + next_task, + stream=stream, + futures=futures, + schedule_task=schedule_task, + submit=submit, + loop=loop, + reraise=reraise, + ), + }, + __name__=task.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, + ), + ) + futures()[fut] = next_task # type: ignore[index] + + fut = cast(Union[asyncio.Future, concurrent.futures.Future], fut) + import sys + + print(f"FOO: {fut} {loop}", flush=True, file=sys.stderr) + # return a chained future to ensure commit() callback is called + # before the returned future is resolved, to ensure stream order etc + try: + in_async = asyncio.current_task() is not None + except RuntimeError: + in_async = False + # if in async context return an async future + # otherwise return a chained sync future + if in_async: + if isinstance(fut, asyncio.Task): + sfut: Union[asyncio.Future[Any], concurrent.futures.Future[Any]] = ( + asyncio.Future(loop=loop) + ) + loop.call_soon_threadsafe(chain_future, fut, sfut) + return sfut + else: + # already wrapped in a future + return fut + else: + sfut = concurrent.futures.Future() + loop.call_soon_threadsafe(chain_future, fut, sfut) + return sfut