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 e54640ce4..67f9c7b3e 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, @@ -25,12 +26,10 @@ from langgraph.constants import ( CONF, CONFIG_KEY_CALL, CONFIG_KEY_SCRATCHPAD, - CONFIG_KEY_SEND, ERROR, INTERRUPT, MISSING, NO_WRITES, - PUSH, RESUME, RETURN, TAG_HIDDEN, @@ -48,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 @@ -56,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: @@ -85,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) @@ -102,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, @@ -125,99 +131,9 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: - def writer( - task: PregelExecutableTask, - writes: Sequence[tuple[str, Any]], - *, - calls: Optional[Sequence[Call]] = None, - ) -> Sequence[Optional[concurrent.futures.Future]]: - if all(w[0] != PUSH for w in writes): - return task.config[CONF][CONFIG_KEY_SEND](writes) - - # schedule PUSH tasks, collect futures - scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] - rtn: dict[int, Optional[concurrent.futures.Future]] = {} - for idx, w in enumerate(writes): - # bail if not a PUSH write - if w[0] != PUSH: - continue - # schedule the next task, if the callback returns one - wcall = calls[idx] if calls else None - if next_task := self.schedule_task( - task, scratchpad.call_counter(), wcall - ): - 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 - rtn[idx] = fut - 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) - rtn[idx] = 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, - # starting a new task in the next tick ensures - # updates from this tick are committed/streamed first - __next_tick__=True, - ) - futures[fut] = next_task - rtn[idx] = fut - return [rtn.get(i) for i in range(len(writes))] - - 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,) = writer( - task, - [(PUSH, None)], - calls=[Call(func, input, retry=retry, callbacks=callbacks)], - ) - assert fut is not None, "writer did not return a future for call" - # 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, ) @@ -231,8 +147,15 @@ class PregelRunner: t, retry_policy, configurable={ - CONFIG_KEY_SEND: partial(writer, t), - 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) @@ -255,13 +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_SEND: partial(writer, t), - 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, ) @@ -313,125 +243,10 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, ) -> AsyncIterator[None]: - def writer( - task: PregelExecutableTask, - writes: Sequence[tuple[str, Any]], - *, - calls: Optional[Sequence[Call]] = None, - ) -> Sequence[Optional[asyncio.Future]]: - if all(w[0] != PUSH for w in writes): - return task.config[CONF][CONFIG_KEY_SEND](writes) - - # schedule PUSH tasks, collect futures - scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] - rtn: dict[int, Optional[asyncio.Future]] = {} - for idx, w in enumerate(writes): - # bail if not a PUSH write - if w[0] != PUSH: - continue - # schedule the next task, if the callback returns one - wcall = calls[idx] if calls is not None else None - if next_task := self.schedule_task( - task, scratchpad.call_counter(), wcall - ): - # if the parent task was retried, - # the next task might already be running - if fut := next( - ( - 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 - rtn[idx] = fut - 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) - rtn[idx] = 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, - # starting a new task in the next tick ensures - # updates from this tick are committed/streamed first - __next_tick__=True, - ), - ) - futures[fut] = next_task - rtn[idx] = fut - return [rtn.get(i) for i in range(len(writes))] - - 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,) = writer( - task, - [(PUSH, None)], - calls=[Call(func, input, retry=retry, callbacks=callbacks)], - ) - assert fut is not None, "writer did not return a future for call" - # 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, ) @@ -446,8 +261,17 @@ class PregelRunner: retry_policy, stream=self.use_astream, configurable={ - CONFIG_KEY_SEND: partial(writer, t), - 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) @@ -472,14 +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_SEND: partial(writer, t), - 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, @@ -539,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 @@ -547,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", []) @@ -562,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( @@ -633,3 +466,193 @@ 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) + # 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 diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index 2d442e131..6564b28a0 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -1,6 +1,7 @@ import asyncio import binascii import concurrent.futures +import weakref from collections.abc import Sequence from contextlib import ( AbstractAsyncContextManager, @@ -216,10 +217,11 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): checkpoint_null_version=checkpoint_null_version(saved.checkpoint), ): # execute task, saving writes + put_writes = partial(self._put_writes, submit, msg["config"]) runner = PregelRunner( - submit=submit, - put_writes=partial(self._put_writes, submit, msg["config"]), - schedule_task=self._schedule_task, + submit=weakref.ref(submit), + put_writes=weakref.ref(put_writes), + schedule_task=weakref.WeakMethod(self._schedule_task), ) async for _ in runner.atick([task], reraise=False): pass @@ -432,10 +434,11 @@ class KafkaExecutor(AbstractContextManager): checkpoint_null_version=checkpoint_null_version(saved.checkpoint), ): # execute task, saving writes + put_writes = partial(self._put_writes, submit, msg["config"]) runner = PregelRunner( - submit=submit, - put_writes=partial(self._put_writes, submit, msg["config"]), - schedule_task=self._schedule_task, + submit=weakref.ref(submit), + put_writes=weakref.ref(put_writes), + schedule_task=weakref.WeakMethod(self._schedule_task), ) for _ in runner.tick([task], reraise=False): pass