diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index fde60fedd..6d8e7adb1 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2509,9 +2509,10 @@ class Pregel(PregelProtocol): # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps. while loop.tick(input_keys=self.input_channels): - loop.match_cached_writes() + for task in loop.match_cached_writes(): + loop.output_writes(task.id, task.writes, cached=True) for _ in runner.tick( - loop.tasks.values(), + [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, retry_policy=self.retry_policy, get_waiter=get_waiter, @@ -2811,9 +2812,10 @@ class Pregel(PregelProtocol): # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps while loop.tick(input_keys=self.input_channels): - await loop.amatch_cached_writes() + for task in await loop.amatch_cached_writes(): + loop.output_writes(task.id, task.writes, cached=True) async for _ in runner.atick( - loop.tasks.values(), + [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, retry_policy=self.retry_policy, get_waiter=get_waiter, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index e46a7fcc6..0bc710e04 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -349,7 +349,7 @@ class PregelLoop(LoopProtocol): ) # output writes if hasattr(self, "tasks"): - self._output_writes(task_id, writes) + self.output_writes(task_id, writes) def _put_pending_writes(self) -> None: if self.checkpointer_put_writes is None: @@ -612,14 +612,14 @@ class PregelLoop(LoopProtocol): # print output for any tasks we applied previous writes to for task in self.tasks.values(): if task.writes: - self._output_writes(task.id, task.writes, cached=True) + self.output_writes(task.id, task.writes, cached=True) return True - def match_cached_writes(self) -> None: + def match_cached_writes(self) -> list[PregelExecutableTask]: raise NotImplementedError - async def amatch_cached_writes(self) -> None: + async def amatch_cached_writes(self) -> list[PregelExecutableTask]: raise NotImplementedError # private @@ -921,7 +921,7 @@ class PregelLoop(LoopProtocol): for v in values(*args, **kwargs): self.stream((self.checkpoint_ns, mode, v)) - def _output_writes( + def output_writes( self, task_id: str, writes: WritesT, *, cached: bool = False ) -> None: if task := self.tasks.get(task_id): @@ -1048,16 +1048,20 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): return self.submit(cast(WritableManagedValue, managed_value).update, values) - def match_cached_writes(self) -> None: + def match_cached_writes(self) -> list[PregelExecutableTask]: if self.cache is None: return + matched: list[PregelExecutableTask] = [] if cached := { t.cache_key.key: t for t in self.tasks.values() if t.cache_key and not t.cache_key.refresh and not t.writes }: for key, values in self.cache.get(tuple(cached)).items(): - cached[key].writes.extend(values) + task = cached[key] + task.writes.extend(values) + matched.append(task) + return matched def put_writes(self, task_id: str, writes: WritesT) -> None: """Put writes for a task, to be read by the next tick.""" @@ -1227,13 +1231,17 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): async def amatch_cached_writes(self) -> None: if self.cache is None: return + matched: list[PregelExecutableTask] = [] if cached := { t.cache_key.key: t for t in self.tasks.values() if t.cache_key and not t.cache_key.refresh and not t.writes }: for key, values in (await self.cache.aget(tuple(cached))).items(): - cached[key].writes.extend(values) + task = cached[key] + task.writes.extend(values) + matched.append(task) + return matched def put_writes(self, task_id: str, writes: WritesT) -> None: """Put writes for a task, to be read by the next tick.""" diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 42c83c7ca..dcf32f96e 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -198,26 +198,25 @@ class PregelRunner: futures[get_waiter()] = None # schedule tasks for t in tasks: - if not t.writes: - fut = self.submit()( # type: ignore[misc] - run_with_retry, - t, - retry_policy, - configurable={ - CONFIG_KEY_CALL: partial( - _call, - weakref.ref(t), - retry=retry_policy, - futures=weakref.ref(futures), - schedule_task=self.schedule_task, - match_cached_writes=match_cached_writes, - submit=self.submit, - reraise=reraise, - ), - }, - __reraise_on_exit__=reraise, - ) - futures[fut] = t + fut = self.submit()( # type: ignore[misc] + run_with_retry, + t, + retry_policy, + configurable={ + CONFIG_KEY_CALL: partial( + _call, + weakref.ref(t), + retry=retry_policy, + futures=weakref.ref(futures), + schedule_task=self.schedule_task, + match_cached_writes=match_cached_writes, + submit=self.submit, + reraise=reraise, + ), + }, + __reraise_on_exit__=reraise, + ) + 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 @@ -332,33 +331,32 @@ class PregelRunner: futures[get_waiter()] = None # schedule tasks for t in tasks: - if not t.writes: - fut = cast( - asyncio.Future, - self.submit()( # type: ignore[misc] - arun_with_retry, - t, - retry_policy, - stream=self.use_astream, - configurable={ - CONFIG_KEY_CALL: partial( - _acall, - weakref.ref(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, - __reraise_on_exit__=reraise, - ), - ) - futures[fut] = t + fut = cast( + asyncio.Future, + self.submit()( # type: ignore[misc] + arun_with_retry, + t, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_CALL: partial( + _acall, + weakref.ref(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, + __reraise_on_exit__=reraise, + ), + ) + 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