mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-21 17:18:09 +02:00
Methods are weak
This commit is contained in:
William Fu-Hinthorn
@@ -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
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user