WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-10-01 22:15:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix reference cycle btwn PregelLoop and PregelRunner (#3907)

Browse Source
This commit is contained in:
William FH
2025-03-18 17:37:01 -07:00
committed by GitHub
parent 2eba27b01f 12c7ddf3a3
commit 48de4a7234
4 changed files with 306 additions and 261 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/Makefile
+3 -1
View File
@@ -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
libs/langgraph/langgraph/pregel/__init__.py
+41 -24
View File
@@ -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),
)
libs/langgraph/langgraph/pregel/runner.py
+253 -230
View File
@@ -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
libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
+9 -6
View File
@@ -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