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

Merge pull request #244 from langchain-ai/nc/30mar/pending-writes-by-channel

Browse Source 
Collect pending writes for each node separately
This commit is contained in:
Nuno Campos
2024-04-01 16:13:32 -07:00
committed by GitHub
parent 90e39a6cc8 885e5884c5
commit 6de5a6de28
2 changed files with 69 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/pregel/__init__.py
+64 -22
View File
@@ -10,6 +10,7 @@ from typing import (
Awaitable,
Callable,
Iterator,
Literal,
Mapping,
NamedTuple,
Optional,
@@ -284,7 +285,7 @@ class Pregel(
config = saved.config if saved else config
with ChannelsManager(self.channels, checkpoint) as channels:
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, update_seen=False
checkpoint, self.nodes, channels, for_execution=False
)
values = {
k: _read_channel(channels, k)
@@ -308,7 +309,7 @@ class Pregel(
config = saved.config if saved else config
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, update_seen=False
checkpoint, self.nodes, channels, for_execution=False
)
values = {
k: _read_channel(channels, k)
@@ -330,7 +331,7 @@ class Pregel(
for config, checkpoint, parent_config in self.checkpointer.list(config):
with ChannelsManager(self.channels, checkpoint) as channels:
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, update_seen=False
checkpoint, self.nodes, channels, for_execution=False
)
values = {
k: _read_channel(channels, k)
@@ -355,7 +356,7 @@ class Pregel(
async for config, checkpoint, parent_config in self.checkpointer.alist(config):
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
_, next_tasks = _prepare_next_tasks(
checkpoint, self.nodes, channels, update_seen=False
checkpoint, self.nodes, channels, for_execution=False
)
values = {
k: _read_channel(channels, k)
@@ -488,7 +489,9 @@ class Pregel(
w for c in input for w in map_input(input_keys, c)
):
# discard any unfinished tasks from previous checkpoint
checkpoint, _ = _prepare_next_tasks(checkpoint, processes, channels)
checkpoint, _ = _prepare_next_tasks(
checkpoint, processes, channels, for_execution=True
)
# apply input writes
_apply_writes(
checkpoint,
@@ -514,7 +517,7 @@ class Pregel(
# with channel updates applied only at the transition between steps
for step in range(config["recursion_limit"] + 1):
checkpoint, next_tasks = _prepare_next_tasks(
checkpoint, processes, channels
checkpoint, processes, channels, for_execution=True
)
# if no more tasks, we're done
@@ -533,9 +536,6 @@ class Pregel(
if debug:
print_step_start(step, next_tasks)
# collect all writes to channels, without applying them yet
pending_writes = deque[tuple[str, Any]]()
# prepare tasks with config
tasks_w_config = [
(
@@ -547,12 +547,12 @@ class Pregel(
callbacks=run_manager.get_child(f"graph:step:{step}"),
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: pending_writes.extend,
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_READ: read,
},
),
)
for proc, input, name in next_tasks
for proc, input, name, writes in next_tasks
]
futures = [
@@ -571,6 +571,11 @@ class Pregel(
# panic on failure or timeout
_panic_or_proceed(done, inflight, step)
# combine pending writes from all tasks
pending_writes = deque[tuple[str, Any]]()
for _, _, _, writes in next_tasks:
pending_writes.extend(writes)
# apply writes to channels
_apply_writes(
checkpoint, channels, pending_writes, config, step + 1
@@ -671,7 +676,9 @@ class Pregel(
[w async for c in input for w in map_input(input_keys, c)]
):
# discard any unfinished tasks from previous checkpoint
checkpoint, _ = _prepare_next_tasks(checkpoint, processes, channels)
checkpoint, _ = _prepare_next_tasks(
checkpoint, processes, channels, for_execution=True
)
# apply input writes
_apply_writes(
checkpoint,
@@ -697,7 +704,7 @@ class Pregel(
# channel updates being applied only at the transition between steps
for step in range(config["recursion_limit"] + 1):
checkpoint, next_tasks = _prepare_next_tasks(
checkpoint, processes, channels
checkpoint, processes, channels, for_execution=True
)
# if no more tasks, we're done
@@ -713,9 +720,6 @@ class Pregel(
if debug:
print_step_start(step, next_tasks)
# collect all writes to channels, without applying them yet
pending_writes = deque[tuple[str, Any]]()
# prepare tasks with config
tasks_w_config = [
(
@@ -727,12 +731,12 @@ class Pregel(
callbacks=run_manager.get_child(f"graph:step:{step}"),
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: pending_writes.extend,
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_READ: read,
},
),
)
for proc, input, name in next_tasks
for proc, input, name, writes in next_tasks
]
futures = (
@@ -758,6 +762,11 @@ class Pregel(
# panic on failure or timeout
_panic_or_proceed(done, inflight, step)
# combine pending writes from all tasks
pending_writes = deque[tuple[str, Any]]()
for _, _, _, writes in next_tasks:
pending_writes.extend(writes)
# apply writes to channels
_apply_writes(
checkpoint, channels, pending_writes, config, step + 1
@@ -1066,14 +1075,44 @@ def _apply_writes(
channels[chan].update([])
@overload
def _prepare_next_tasks(
checkpoint: Checkpoint,
processes: Mapping[str, ChannelInvoke],
channels: Mapping[str, BaseChannel],
update_seen: bool = True,
for_execution: Literal[False],
) -> tuple[Checkpoint, list[tuple[Runnable, Any, str]]]:
...
@overload
def _prepare_next_tasks(
checkpoint: Checkpoint,
processes: Mapping[str, ChannelInvoke],
channels: Mapping[str, BaseChannel],
for_execution: Literal[True],
) -> tuple[Checkpoint, list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]]]:
...
def _prepare_next_tasks(
checkpoint: Checkpoint,
processes: Mapping[str, ChannelInvoke],
channels: Mapping[str, BaseChannel],
*,
for_execution: bool,
) -> tuple[
Checkpoint,
Union[
list[tuple[Runnable, Any, str]],
list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]],
],
]:
checkpoint = copy_checkpoint(checkpoint)
tasks: list[tuple[Runnable, Any, str]] = []
tasks: Union[
list[tuple[Runnable, Any, str]],
list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]],
] = []
# Check if any processes should be run in next step
# If so, prepare the values to be passed to them
for name, proc in processes.items():
@@ -1106,7 +1145,7 @@ def _prepare_next_tasks(
val = val[None]
# update seen versions
if update_seen:
if for_execution:
seen.update(
{
chan: checkpoint["channel_versions"][chan]
@@ -1116,7 +1155,10 @@ def _prepare_next_tasks(
# skip if condition is not met
if proc.when is None or proc.when(val):
tasks.append((proc, val, name))
if for_execution:
tasks.append((proc, val, name, deque()))
else:
tasks.append((proc, val, name))
return checkpoint, tasks
langgraph/pregel/debug.py
+5 -2
View File
@@ -1,3 +1,4 @@
from collections import deque
from pprint import pformat
from typing import Any, Iterator, Mapping
@@ -7,14 +8,16 @@ from langchain_core.utils.input import get_bolded_text, get_colored_text
from langgraph.channels.base import BaseChannel, EmptyChannelError
def print_step_start(step: int, next_tasks: list[tuple[Runnable, Any, str]]) -> None:
def print_step_start(
step: int, next_tasks: list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]]
) -> None:
n_tasks = len(next_tasks)
print(
f"{get_colored_text('[langgraph/step]', color='blue')} "
+ get_bolded_text(
f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Next tasks:\n"
)
+ "\n".join(f"- {name}({pformat(val)})" for _, val, name in next_tasks)
+ "\n".join(f"- {name}({pformat(val)})" for _, val, name, _ in next_tasks)
)