WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-25 17:12:26 +02:00
Code Issues Packages Projects Releases Wiki Activity

Output cached writes

Browse Source
This commit is contained in:
Nuno Campos
2025-05-08 16:49:01 -07:00
parent 64491a2b29
commit 09fdc14d0a
3 changed files with 67 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+6 -4
View File
@@ -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,
libs/langgraph/langgraph/pregel/loop.py
+16 -8
View File
@@ -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."""
libs/langgraph/langgraph/pregel/runner.py
+45 -47
View File
@@ -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