WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-12 04:37:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

Execute Sends in same super step that triggered them

Browse Source 
- Keep old code path for compatibility with existing checkpoints
- Keep a similar order of application of updates, in some cases there will be no visible change
- Update task path for Sends to contain the path of all the parent tasks (multiple parents when a Send task creates another Send)
- That lineage path is used to ensure order of application of updates respects their logical lineage (ie updates from parents always applied before their child tasks)
- Move Interrupt writes to use negative indexes, which allow replacing/shadowing (when task is re-run it may interrupt again, or succeed)
- Runner will now attempt to schedule new Send tasks as soon as the write is received (ie while the originating node is still running)
- Update kafka scheduler to support new Send behavior
This commit is contained in:
Nuno Campos
2024-11-11 14:01:16 -08:00
parent 75c502cd93
commit 89a0859928
19 changed files with 2143 additions and 532 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+2 -1
View File
@@ -24,6 +24,7 @@ from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import (
ERROR,
INTERRUPT,
SCHEDULED,
ChannelProtocol,
SendProtocol,
@@ -449,4 +450,4 @@ Special writes (e.g. errors) map to negative indices, to avoid those writes from
conflicting with regular writes.
Each Checkpointer implementation should use this mapping in put_writes.
"""
WRITES_IDX_MAP = {ERROR: -1, SCHEDULED: -2}
WRITES_IDX_MAP = {ERROR: -1, SCHEDULED: -2, INTERRUPT: -3}
libs/checkpoint/langgraph/checkpoint/serde/types.py
+1
View File
@@ -12,6 +12,7 @@ from typing_extensions import Self
ERROR = "__error__"
SCHEDULED = "__scheduled__"
INTERRUPT = "__interrupt__"
TASKS = "__pregel_tasks"
Value = TypeVar("Value", covariant=True)
libs/langgraph/langgraph/pregel/__init__.py
+46 -17
View File
@@ -69,6 +69,7 @@ from langgraph.constants import (
INTERRUPT,
NS_END,
NS_SEP,
PUSH,
SCHEDULED,
)
from langgraph.errors import (
@@ -98,7 +99,13 @@ from langgraph.pregel.utils import find_subgraph_pregel, get_new_channel_version
from langgraph.pregel.validate import validate_graph, validate_keys
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.types import All, Checkpointer, LoopProtocol, StateSnapshot, StreamMode
from langgraph.types import (
All,
Checkpointer,
LoopProtocol,
StateSnapshot,
StreamMode,
)
from langgraph.utils.config import (
ensure_config,
merge_configs,
@@ -468,6 +475,7 @@ class Pregel(PregelProtocol):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
saved.checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
managed,
@@ -570,6 +578,7 @@ class Pregel(PregelProtocol):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
saved.checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
managed,
@@ -922,6 +931,7 @@ class Pregel(PregelProtocol):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
checkpoint,
saved.pending_writes,
self.nodes,
channels,
managed,
@@ -1001,8 +1011,14 @@ class Pregel(PregelProtocol):
),
)
# save task writes
if saved:
checkpointer.put_writes(checkpoint_config, task.writes, task_id)
# channel writes are saved to current checkpoint
# push writes are saved to next checkpoint
channel_writes, push_writes = (
[w for w in task.writes if w[0] != PUSH],
[w for w in task.writes if w[0] == PUSH],
)
if saved and channel_writes:
checkpointer.put_writes(checkpoint_config, channel_writes, task_id)
# apply to checkpoint and save
mv_writes = apply_writes(
checkpoint, channels, [task], checkpointer.get_next_version
@@ -1023,6 +1039,8 @@ class Pregel(PregelProtocol):
checkpoint_previous_versions, checkpoint["channel_versions"]
),
)
if push_writes:
checkpointer.put_writes(next_config, push_writes, task_id)
return patch_checkpoint_map(next_config, saved.metadata if saved else None)
async def aupdate_state(
@@ -1132,6 +1150,7 @@ class Pregel(PregelProtocol):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
checkpoint,
saved.pending_writes,
self.nodes,
channels,
managed,
@@ -1208,14 +1227,23 @@ class Pregel(PregelProtocol):
),
)
# save task writes
if saved:
await checkpointer.aput_writes(checkpoint_config, writes, task_id)
# channel writes are saved to current checkpoint
# push writes are saved to next checkpoint
channel_writes, push_writes = (
[w for w in task.writes if w[0] != PUSH],
[w for w in task.writes if w[0] == PUSH],
)
if saved and channel_writes:
await checkpointer.aput_writes(
checkpoint_config, channel_writes, task_id
)
# apply to checkpoint and save
mv_writes = apply_writes(
checkpoint, channels, [task], checkpointer.get_next_version
)
assert not mv_writes, "Can't write to SharedValues from update_state"
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
# save checkpoint, after applying writes
next_config = await checkpointer.aput(
checkpoint_config,
checkpoint,
@@ -1230,6 +1258,9 @@ class Pregel(PregelProtocol):
checkpoint_previous_versions, checkpoint["channel_versions"]
),
)
# save push writes
if push_writes:
await checkpointer.aput_writes(next_config, push_writes, task_id)
return patch_checkpoint_map(next_config, saved.metadata if saved else None)
def _defaults(
@@ -1432,12 +1463,16 @@ class Pregel(PregelProtocol):
specs=self.channels,
output_keys=output_keys,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
debug=debug,
) as loop:
# create runner
runner = PregelRunner(
submit=loop.submit,
put_writes=loop.put_writes,
schedule_task=loop.accept_push,
node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED),
)
# enable subgraph streaming
@@ -1468,12 +1503,7 @@ class Pregel(PregelProtocol):
# channel updates from step N are only visible in step N+1
# 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,
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
):
while loop.tick(input_keys=self.input_channels):
for _ in runner.tick(
loop.tasks.values(),
timeout=self.step_timeout,
@@ -1654,11 +1684,15 @@ class Pregel(PregelProtocol):
specs=self.channels,
output_keys=output_keys,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
) as loop:
# create runner
runner = PregelRunner(
submit=loop.submit,
put_writes=loop.put_writes,
schedule_task=loop.accept_push,
use_astream=do_stream is not None,
node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED),
)
@@ -1678,12 +1712,7 @@ class Pregel(PregelProtocol):
# channel updates from step N are only visible in step N+1
# 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,
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
):
while loop.tick(input_keys=self.input_channels):
async for _ in runner.atick(
loop.tasks.values(),
timeout=self.step_timeout,
libs/langgraph/langgraph/pregel/algo.py
+160 -37
View File
@@ -25,6 +25,7 @@ from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
PendingWrite,
V,
copy_checkpoint,
)
@@ -68,7 +69,7 @@ class WritesProtocol(Protocol):
Implemented by PregelTaskWrites and PregelExecutableTask."""
@property
def path(self) -> tuple[Union[str, int], ...]: ...
def path(self) -> tuple[Union[str, int, tuple], ...]: ...
@property
def name(self) -> str: ...
@@ -84,7 +85,7 @@ class PregelTaskWrites(NamedTuple):
"""Simplest implementation of WritesProtocol, for usage with writes that
don't originate from a runnable task, eg. graph input, update_state, etc."""
path: tuple[Union[str, int], ...]
path: tuple[Union[str, int, tuple], ...]
name: str
writes: Sequence[tuple[str, Any]]
triggers: Sequence[str]
@@ -172,7 +173,7 @@ def local_write(
"""Function injected under CONFIG_KEY_SEND in task config, to write to channels.
Validates writes and forwards them to `commit` function."""
for chan, value in writes:
if chan == TASKS:
if chan == PUSH:
if not isinstance(value, Send):
raise InvalidUpdateError(f"Expected Send, got {value}")
if value.node not in process_keys:
@@ -194,8 +195,10 @@ def apply_writes(
"""Apply writes from a set of tasks (usually the tasks from a Pregel step)
to the checkpoint and channels, and return managed values writes to be applied
externally."""
# sort tasks on path
tasks = sorted(tasks, key=lambda t: t.path)
# sort tasks on path, to ensure deterministic order for update application
# any path parts after the 3rd are ignored for sorting
# (we use them for eg. task ids which aren't good for sorting)
tasks = sorted(tasks, key=lambda t: t.path[:3])
# update seen versions
for task in tasks:
@@ -237,8 +240,10 @@ def apply_writes(
for chan, val in task.writes:
if chan == NO_WRITES:
pass
elif chan == TASKS:
elif chan == TASKS: # TODO: remove branch in 1.0
checkpoint["pending_sends"].append(val)
elif chan == PUSH:
pass
elif chan in channels:
pending_writes_by_channel[chan].append(val)
else:
@@ -277,6 +282,7 @@ def apply_writes(
@overload
def prepare_next_tasks(
checkpoint: Checkpoint,
pending_writes: Sequence[PendingWrite],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
@@ -293,6 +299,7 @@ def prepare_next_tasks(
@overload
def prepare_next_tasks(
checkpoint: Checkpoint,
pending_writes: Sequence[PendingWrite],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
@@ -308,6 +315,7 @@ def prepare_next_tasks(
def prepare_next_tasks(
checkpoint: Checkpoint,
pending_writes: Sequence[PendingWrite],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
@@ -322,13 +330,14 @@ def prepare_next_tasks(
"""Prepare the set of tasks that will make up the next Pregel step.
This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered
by edges)."""
tasks: dict[str, Union[PregelTask, PregelExecutableTask]] = {}
# Consume pending packets
for idx, _ in enumerate(checkpoint["pending_sends"]):
tasks: list[Union[PregelTask, PregelExecutableTask]] = []
# Consume pending_sends from previous step (legacy version of Send)
for idx, _ in enumerate(checkpoint["pending_sends"]): # TODO: remove branch in 1.0
if task := prepare_single_task(
(PUSH, idx),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
@@ -339,7 +348,7 @@ def prepare_next_tasks(
checkpointer=checkpointer,
manager=manager,
):
tasks[task.id] = task
tasks.append(task)
# Check if any processes should be run in next step
# If so, prepare the values to be passed to them
for name in processes:
@@ -347,6 +356,7 @@ def prepare_next_tasks(
(PULL, name),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
@@ -357,15 +367,74 @@ def prepare_next_tasks(
checkpointer=checkpointer,
manager=manager,
):
tasks[task.id] = task
return tasks
tasks.append(task)
# Consume pending Sends from this step (new version of Send)
if any(c == PUSH for _, c, _ in pending_writes):
# group writes by task id
grouped_by_task = defaultdict(list)
for tid, c, _ in pending_writes:
grouped_by_task[tid].append(c)
# prepare send tasks from grouped writes
# 1. start from sends originating from existing tasks
tidx = 0
while tidx < len(tasks):
task = tasks[tidx]
if twrites := grouped_by_task.pop(task.id, None):
for idx, c in enumerate(twrites):
if c != PUSH:
continue
if next_task := prepare_single_task(
(PUSH, task.path, idx, task.id),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
config=config,
step=step,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
manager=manager,
):
tasks.append(next_task)
tidx += 1
# key tasks by id
task_map = {t.id: t for t in tasks}
# 2. create new tasks for remaining sends (eg. from update_state)
for tid, writes in grouped_by_task.items():
task = task_map.get(tid)
for idx, c in enumerate(writes):
if c != PUSH:
continue
if next_task := prepare_single_task(
(PUSH, task.path if task else (), idx, tid),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
config=config,
step=step,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
manager=manager,
):
task_map[next_task.id] = next_task
else:
task_map = {t.id: t for t in tasks}
return task_map
def prepare_single_task(
task_path: tuple[str, Union[int, str]],
task_path: tuple[Union[str, int, tuple], ...],
task_id_checksum: Optional[str],
*,
checkpoint: Checkpoint,
pending_writes: Sequence[PendingWrite],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
@@ -383,31 +452,75 @@ def prepare_single_task(
parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
if task_path[0] == PUSH:
idx = int(task_path[1])
if idx >= len(checkpoint["pending_sends"]):
return
packet = checkpoint["pending_sends"][idx]
if not isinstance(packet, Send):
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending sends"
if len(task_path) == 2: # TODO: remove branch in 1.0
# legacy SEND tasks, executed in superstep n+1
# (PUSH, idx of pending send)
idx = cast(int, task_path[1])
if idx >= len(checkpoint["pending_sends"]):
return
packet = checkpoint["pending_sends"][idx]
if not isinstance(packet, Send):
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending sends"
)
return
if packet.node not in processes:
logger.warning(
f"Ignoring unknown node name {packet.node} in pending sends"
)
return
# create task id
triggers = [PUSH]
checkpoint_ns = (
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = _uuid5_str(
checkpoint_id,
checkpoint_ns,
str(step),
packet.node,
PUSH,
str(idx),
)
elif len(task_path) == 4:
# new PUSH tasks, executed in superstep n
# (PUSH, parent task path, idx of PUSH write, id of parent task)
task_path_t = cast(tuple[str, tuple, int, str], task_path)
writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]]
if task_path_t[2] >= len(writes_for_path):
logger.warning(
f"Ignoring invalid write index {task_path[2]} in pending writes"
)
return
packet = writes_for_path[task_path_t[2]][2]
if not isinstance(packet, Send):
print("packet", task_path_t, writes_for_path)
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending writes"
)
return
if packet.node not in processes:
logger.warning(
f"Ignoring unknown node name {packet.node} in pending writes"
)
return
# create task id
triggers = [PUSH]
checkpoint_ns = (
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = _uuid5_str(
checkpoint_id,
checkpoint_ns,
str(step),
packet.node,
PUSH,
_tuple_str(task_path[1]),
str(task_path[2]),
)
else:
logger.warning(f"Ignoring invalid PUSH task path {task_path}")
return
if packet.node not in processes:
logger.warning(f"Ignoring unknown node name {packet.node} in pending sends")
return
# create task id
triggers = [PUSH]
checkpoint_ns = (
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = _uuid5_str(
checkpoint_id,
checkpoint_ns,
str(step),
packet.node,
PUSH,
str(idx),
)
task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
metadata = {
"langgraph_step": step,
@@ -417,7 +530,7 @@ def prepare_single_task(
"langgraph_checkpoint_ns": task_checkpoint_ns,
}
if task_id_checksum is not None:
assert task_id == task_id_checksum
assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}"
if for_execution:
proc = processes[packet.node]
if node := proc.node:
@@ -481,6 +594,7 @@ def prepare_single_task(
else:
return PregelTask(task_id, packet.node, task_path)
elif task_path[0] == PULL:
# (PULL, node name)
name = cast(str, task_path[1])
if name not in processes:
return
@@ -642,3 +756,12 @@ def _uuid5_str(namespace: bytes, *parts: str) -> str:
sha.update(b"".join(p.encode() for p in parts))
hex = sha.hexdigest()
return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}"
def _tuple_str(tup: Union[str, int, tuple]) -> str:
"""Generate a string representation of a tuple."""
return (
f"({', '.join(_tuple_str(x) for x in tup)})"
if isinstance(tup, (tuple, list))
else str(tup)
)
libs/langgraph/langgraph/pregel/debug.py
+1 -1
View File
@@ -208,7 +208,7 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None:
print(
f"{get_colored_text(f'[{step}:tasks]', color='blue')} "
+ get_bolded_text(
f"Starting step {step} with {n_tasks} task{'s' if n_tasks != 1 else ''}:\n"
f"Starting {n_tasks} task{'s' if n_tasks != 1 else ''} for step {step}:\n"
)
+ "\n".join(
f"- {get_colored_text(task.name, 'green')} -> {pformat(task.input)}"
libs/langgraph/langgraph/pregel/loop.py
+109 -13
View File
@@ -52,6 +52,7 @@ from langgraph.constants import (
INPUT,
INTERRUPT,
NS_SEP,
PUSH,
SCHEDULED,
TAG_HIDDEN,
)
@@ -74,6 +75,7 @@ from langgraph.pregel.algo import (
apply_writes,
increment,
prepare_next_tasks,
prepare_single_task,
should_interrupt,
)
from langgraph.pregel.debug import (
@@ -130,6 +132,9 @@ class PregelLoop(LoopProtocol):
stream_keys: Union[str, Sequence[str]]
skip_done_tasks: bool
is_nested: bool
manager: Union[None, AsyncParentRunManager, ParentRunManager]
interrupt_after: Union[All, Sequence[str]]
interrupt_before: Union[All, Sequence[str]]
checkpointer_get_next_version: GetNextVersion
checkpointer_put_writes: Optional[
@@ -162,6 +167,7 @@ class PregelLoop(LoopProtocol):
"pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
]
tasks: dict[str, PregelExecutableTask]
to_interrupt: list[PregelExecutableTask]
output: Union[None, dict[str, Any], Any] = None
# public
@@ -178,6 +184,9 @@ class PregelLoop(LoopProtocol):
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]],
stream_keys: Union[str, Sequence[str]],
interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
check_subgraphs: bool = True,
debug: bool = False,
) -> None:
@@ -194,6 +203,9 @@ class PregelLoop(LoopProtocol):
self.specs = specs
self.output_keys = output_keys
self.stream_keys = stream_keys
self.interrupt_after = interrupt_after
self.interrupt_before = interrupt_before
self.manager = manager
self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {})
self.skip_done_tasks = (
CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
@@ -263,13 +275,57 @@ class PregelLoop(LoopProtocol):
# output writes
self._output_writes(task_id, writes)
def accept_push(
self, task: PregelExecutableTask, write_idx: int
) -> Optional[PregelExecutableTask]:
"""Accept a PUSH from a task, potentially returning a new task to start."""
# don't start if an earlier PUSH has already triggered an interrupt
if self.to_interrupt:
return
# don't start if we should interrupt *after* the original task
if should_interrupt(self.checkpoint, self.interrupt_after, [task]):
self.to_interrupt.append(task)
return
if pushed := cast(
Optional[PregelExecutableTask],
prepare_single_task(
(PUSH, task.path, write_idx, task.id),
None,
checkpoint=self.checkpoint,
pending_writes=[(task.id, *w) for w in task.writes],
processes=self.nodes,
channels=self.channels,
managed=self.managed,
config=self.config,
step=self.step,
for_execution=True,
store=self.store,
checkpointer=self.checkpointer,
manager=self.manager,
),
):
# don't start if we should interrupt *before* the new task
if should_interrupt(self.checkpoint, self.interrupt_before, [pushed]):
self.to_interrupt.append(pushed)
return
# produce debug output
self._emit("debug", map_debug_tasks, self.step, [pushed])
# debug flag
if self.debug:
print_step_tasks(self.step, [pushed])
# save the new task
self.tasks[pushed.id] = pushed
# match any pending writes to the new task
if self.skip_done_tasks:
self._match_writes({pushed.id: pushed})
# return the new task, to be started, if not run before
if not pushed.writes:
return pushed
def tick(
self,
*,
input_keys: Union[str, Sequence[str]],
interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
) -> bool:
"""Execute a single iteration of the Pregel loop.
Returns True if more iterations are needed."""
@@ -278,6 +334,10 @@ class PregelLoop(LoopProtocol):
if self.input not in (INPUT_DONE, INPUT_RESUMING):
self._first(input_keys=input_keys)
elif self.to_interrupt:
# if we need to interrupt, do so
self.status = "interrupt_before"
raise GraphInterrupt()
elif all(task.writes for task in self.tasks.values()):
writes = [w for t in self.tasks.values() for w in t.writes]
# debug flag
@@ -322,7 +382,9 @@ class PregelLoop(LoopProtocol):
}
)
# after execution, check if we should interrupt
if should_interrupt(self.checkpoint, interrupt_after, self.tasks.values()):
if should_interrupt(
self.checkpoint, self.interrupt_after, self.tasks.values()
):
self.status = "interrupt_after"
raise GraphInterrupt()
else:
@@ -336,16 +398,18 @@ class PregelLoop(LoopProtocol):
# prepare next tasks
self.tasks = prepare_next_tasks(
self.checkpoint,
self.checkpoint_pending_writes,
self.nodes,
self.channels,
self.managed,
self.config,
self.step,
for_execution=True,
manager=manager,
manager=self.manager,
store=self.store,
checkpointer=self.checkpointer,
)
self.to_interrupt = []
# produce debug output
if self._checkpointer_put_after_previous is not None:
@@ -387,15 +451,12 @@ class PregelLoop(LoopProtocol):
# if all tasks have finished, re-tick
if all(task.writes for task in self.tasks.values()):
return self.tick(
input_keys=input_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
manager=manager,
)
return self.tick(input_keys=input_keys)
# before execution, check if we should interrupt
if should_interrupt(self.checkpoint, interrupt_before, self.tasks.values()):
if should_interrupt(
self.checkpoint, self.interrupt_before, self.tasks.values()
):
self.status = "interrupt_before"
raise GraphInterrupt()
@@ -464,6 +525,7 @@ class PregelLoop(LoopProtocol):
# discard any unfinished tasks from previous checkpoint
discard_tasks = prepare_next_tasks(
self.checkpoint,
self.checkpoint_pending_writes,
self.nodes,
self.channels,
self.managed,
@@ -577,11 +639,33 @@ class PregelLoop(LoopProtocol):
# save final output
self.output = read_channels(self.channels, self.output_keys)
if suppress:
# suppress interrupt
# emit one last "values" event, with pending writes applied
if (
hasattr(self, "tasks")
and self.checkpoint_pending_writes
and any(task.writes for task in self.tasks.values())
):
mv_writes = apply_writes(
self.checkpoint,
self.channels,
self.tasks.values(),
self.checkpointer_get_next_version,
)
for key, values in mv_writes.items():
self._update_mv(key, values)
self._emit(
"values",
map_output_values,
self.output_keys,
[w for t in self.tasks.values() for w in t.writes],
self.channels,
)
# emit INTERRUPT event
self._emit(
"updates",
lambda: iter([{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]),
)
# suppress interrupt
return True
def _emit(
@@ -635,6 +719,9 @@ class SyncPregelLoop(PregelLoop, ContextManager):
checkpointer: Optional[BaseCheckpointSaver],
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
check_subgraphs: bool = True,
@@ -650,7 +737,10 @@ class SyncPregelLoop(PregelLoop, ContextManager):
specs=specs,
output_keys=output_keys,
stream_keys=stream_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
check_subgraphs=check_subgraphs,
manager=manager,
debug=debug,
)
self.stack = ExitStack()
@@ -761,6 +851,9 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
checkpointer: Optional[BaseCheckpointSaver],
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
check_subgraphs: bool = True,
@@ -776,7 +869,10 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
specs=specs,
output_keys=output_keys,
stream_keys=stream_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
check_subgraphs=check_subgraphs,
manager=manager,
debug=debug,
)
self.stack = AsyncExitStack()
libs/langgraph/langgraph/pregel/retry.py
+18 -2
View File
@@ -2,9 +2,15 @@ import asyncio
import logging
import random
import time
from typing import Optional, Sequence
from functools import partial
from typing import Any, Callable, Optional, Sequence
from langgraph.constants import CONF, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_RESUMING
from langgraph.constants import (
CONF,
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_RESUMING,
CONFIG_KEY_SEND,
)
from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphInterrupt
from langgraph.types import PregelExecutableTask, RetryPolicy
from langgraph.utils.config import patch_configurable
@@ -15,12 +21,17 @@ logger = logging.getLogger(__name__)
def run_with_retry(
task: PregelExecutableTask,
retry_policy: Optional[RetryPolicy],
writer: Optional[
Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None]
] = None,
) -> None:
"""Run a task with retries."""
retry_policy = task.retry_policy or retry_policy
interval = retry_policy.initial_interval if retry_policy else 0
attempts = 0
config = task.config
if writer is not None:
config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)})
while True:
try:
# clear any writes from previous attempts
@@ -84,12 +95,17 @@ async def arun_with_retry(
task: PregelExecutableTask,
retry_policy: Optional[RetryPolicy],
stream: bool = False,
writer: Optional[
Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None]
] = None,
) -> None:
"""Run a task asynchronously with retries."""
retry_policy = task.retry_policy or retry_policy
interval = retry_policy.initial_interval if retry_policy else 0
attempts = 0
config = task.config
if writer is not None:
config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)})
while True:
try:
# clear any writes from previous attempts
libs/langgraph/langgraph/pregel/runner.py
+115 -27
View File
@@ -14,7 +14,15 @@ from typing import (
cast,
)
from langgraph.constants import ERROR, INTERRUPT, NO_WRITES, TAG_HIDDEN
from langgraph.constants import (
CONF,
CONFIG_KEY_SEND,
ERROR,
INTERRUPT,
NO_WRITES,
PUSH,
TAG_HIDDEN,
)
from langgraph.errors import GraphDelegate, GraphInterrupt
from langgraph.pregel.executor import Submit
from langgraph.pregel.retry import arun_with_retry, run_with_retry
@@ -31,6 +39,9 @@ class PregelRunner:
*,
submit: Submit,
put_writes: Callable[[str, Sequence[tuple[str, Any]]], None],
schedule_task: Callable[
[PregelExecutableTask, int], Optional[PregelExecutableTask]
],
use_astream: bool = False,
node_finished: Optional[Callable[[str], None]] = None,
) -> None:
@@ -38,6 +49,7 @@ class PregelRunner:
self.put_writes = put_writes
self.use_astream = use_astream
self.node_finished = node_finished
self.schedule_task = schedule_task
def tick(
self,
@@ -48,27 +60,58 @@ 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]]
) -> None:
prev_length = len(task.writes)
# delegate to the underlying writer
task.config[CONF][CONFIG_KEY_SEND](writes)
for idx, w in enumerate(task.writes):
# find the index for the newly inserted writes
if idx < prev_length:
continue
assert writes[idx - prev_length] is w
# bail if not a PUSH write
if w[0] != PUSH:
continue
# schedule the next task, if the callback returns one
if next_task := self.schedule_task(task, idx):
# if the parent task was retried,
# the next task might already be running
if any(
t == next_task.id for t in futures.values() if t is not None
):
continue
# schedule the next task
futures[
self.submit(
run_with_retry,
next_task,
retry_policy,
writer=writer,
__reraise_on_exit__=reraise,
)
] = next_task
tasks = tuple(tasks)
futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {}
# give control back to the caller
yield
# fast path if single task with no timeout and no waiter
if len(tasks) == 1 and timeout is None and get_waiter is None:
t = tasks[0]
try:
run_with_retry(t, retry_policy)
run_with_retry(t, retry_policy, writer=writer)
self.commit(t, None)
except Exception as exc:
self.commit(t, exc)
if reraise:
raise
return
if not futures: # maybe `t` schuduled another task
return
# add waiter task if requested
if get_waiter is not None:
futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {
get_waiter(): None
}
else:
futures = {}
futures[get_waiter()] = None
# 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
@@ -79,10 +122,11 @@ class PregelRunner:
run_with_retry,
t,
retry_policy,
writer=writer,
__reraise_on_exit__=reraise,
)
] = t
all_futures = futures.copy()
done_futures: set[concurrent.futures.Future] = set()
end_time = timeout + time.monotonic() if timeout else None
while len(futures) > (1 if get_waiter is not None else 0):
done, inflight = concurrent.futures.wait(
@@ -99,6 +143,8 @@ class PregelRunner:
if inflight and get_waiter is not None:
futures[get_waiter()] = None
else:
# store for panic check
done_futures.add(fut)
# task finished, commit writes
self.commit(task, _exception(fut))
else:
@@ -110,7 +156,10 @@ class PregelRunner:
# give control back to the caller
yield
# panic on failure or timeout
_panic_or_proceed(all_futures, panic=reraise)
_panic_or_proceed(
done_futures.union(f for f, t in futures.items() if t is not None),
panic=reraise,
)
async def atick(
self,
@@ -121,28 +170,67 @@ 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]]
) -> None:
prev_length = len(task.writes)
# delegate to the underlying writer
task.config[CONF][CONFIG_KEY_SEND](writes)
for idx, w in enumerate(task.writes):
# find the index for the newly inserted writes
if idx < prev_length:
continue
assert writes[idx - prev_length] is w
# bail if not a PUSH write
if w[0] != PUSH:
continue
# schedule the next task, if the callback returns one
if next_task := self.schedule_task(task, idx):
# if the parent task was retried,
# the next task might already be running
if any(
t == next_task.id for t in futures.values() if t is not None
):
continue
# schedule the next task
futures[
cast(
asyncio.Future,
self.submit(
arun_with_retry,
next_task,
retry_policy,
stream=self.use_astream,
writer=writer,
__name__=t.name,
__cancel_on_exit__=True,
__reraise_on_exit__=reraise,
),
)
] = next_task
loop = asyncio.get_event_loop()
tasks = tuple(tasks)
futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {}
# give control back to the caller
yield
# fast path if single task with no waiter and no timeout
if len(tasks) == 1 and get_waiter is None and timeout is None:
t = tasks[0]
try:
await arun_with_retry(t, retry_policy, stream=self.use_astream)
await arun_with_retry(
t, retry_policy, stream=self.use_astream, writer=writer
)
self.commit(t, None)
except Exception as exc:
self.commit(t, exc)
if reraise:
raise
return
if not futures: # maybe `t` schuduled another task
return
# add waiter task if requested
if get_waiter is not None:
futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {
get_waiter(): None
}
else:
futures = {}
futures[get_waiter()] = None
# 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
@@ -156,13 +244,14 @@ class PregelRunner:
t,
retry_policy,
stream=self.use_astream,
writer=writer,
__name__=t.name,
__cancel_on_exit__=True,
__reraise_on_exit__=reraise,
),
)
] = t
all_futures = futures.copy()
done_futures: set[asyncio.Future] = set()
end_time = timeout + loop.time() if timeout else None
while len(futures) > (1 if get_waiter is not None else 0):
done, inflight = await asyncio.wait(
@@ -179,6 +268,8 @@ class PregelRunner:
if inflight and get_waiter is not None:
futures[get_waiter()] = None
else:
# store for panic check
done_futures.add(fut)
# task finished, commit writes
self.commit(task, _exception(fut))
else:
@@ -194,7 +285,9 @@ class PregelRunner:
fut.cancel()
# panic on failure or timeout
_panic_or_proceed(
all_futures, timeout_exc_cls=asyncio.TimeoutError, panic=reraise
done_futures.union(f for f, t in futures.items() if t is not None),
timeout_exc_cls=asyncio.TimeoutError,
panic=reraise,
)
def commit(
@@ -250,10 +343,7 @@ def _exception(
def _panic_or_proceed(
futs: Union[
dict[concurrent.futures.Future, Optional[PregelExecutableTask]],
dict[asyncio.Future, Optional[PregelExecutableTask]],
],
futs: Union[set[concurrent.futures.Future], set[asyncio.Future]],
*,
timeout_exc_cls: Type[Exception] = TimeoutError,
panic: bool = True,
@@ -261,10 +351,8 @@ def _panic_or_proceed(
"""Cancel remaining tasks if any failed, re-raise exception if panic is True."""
done: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
inflight: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
for fut, val in futs.items():
if val is None:
continue
elif fut.done():
for fut in futs:
if fut.done():
done.add(fut)
else:
inflight.add(fut)
libs/langgraph/langgraph/pregel/write.py
+3 -3
View File
@@ -14,7 +14,7 @@ from typing import (
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONF, CONFIG_KEY_SEND, TASKS, Send
from langgraph.constants import CONF, CONFIG_KEY_SEND, PUSH, TASKS, Send
from langgraph.errors import InvalidUpdateError
from langgraph.utils.runnable import RunnableCallable
@@ -112,14 +112,14 @@ class ChannelWrite(RunnableCallable):
# validate
for w in writes:
if isinstance(w, ChannelWriteEntry):
if w.channel == TASKS:
if w.channel in (TASKS, PUSH):
raise InvalidUpdateError(
"Cannot write to the reserved channel TASKS"
)
if w.value is PASSTHROUGH:
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
# split packets and entries
sends = [(TASKS, packet) for packet in writes if isinstance(packet, Send)]
sends = [(PUSH, packet) for packet in writes if isinstance(packet, Send)]
entries = [write for write in writes if isinstance(write, ChannelWriteEntry)]
# process entries into values
values = [
libs/langgraph/langgraph/types.py
+2 -2
View File
@@ -110,7 +110,7 @@ class Interrupt:
class PregelTask(NamedTuple):
id: str
name: str
path: tuple[Union[str, int], ...]
path: tuple[Union[str, int, tuple], ...]
error: Optional[Exception] = None
interrupts: tuple[Interrupt, ...] = ()
state: Union[None, RunnableConfig, "StateSnapshot"] = None
@@ -127,7 +127,7 @@ class PregelExecutableTask(NamedTuple):
retry_policy: Optional[RetryPolicy]
cache_policy: Optional[CachePolicy]
id: str
path: tuple[Union[str, int], ...]
path: tuple[Union[str, int, tuple], ...]
scheduled: bool = False
libs/langgraph/tests/test_algo.py
+9 -1
View File
@@ -11,13 +11,21 @@ def test_prepare_next_tasks() -> None:
with ChannelsManager({}, checkpoint, config) as (channels, managed):
assert (
prepare_next_tasks(
checkpoint, processes, channels, managed, config, 0, for_execution=False
checkpoint,
{},
processes,
channels,
managed,
config,
0,
for_execution=False,
)
== {}
)
assert (
prepare_next_tasks(
checkpoint,
{},
processes,
channels,
managed,
libs/langgraph/tests/test_pregel.py
+624 -215
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_pregel_async.py
+603 -197
View File
File diff suppressed because it is too large Load Diff
libs/scheduler-kafka/Makefile
+1 -1
View File
@@ -19,7 +19,7 @@ test:
exit $$EXIT_CODE
test_watch:
make start-services && poetry run ptw . -- $(TEST_PATH); \
make start-services && poetry run ptw . -- -x $(TEST_PATH); \
EXIT_CODE=$$?; \
make stop-services; \
exit $$EXIT_CODE
libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
+21 -1
View File
@@ -38,7 +38,7 @@ from langgraph.scheduler.kafka.types import (
Sendable,
Topics,
)
from langgraph.types import LoopProtocol, RetryPolicy
from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy
from langgraph.utils.config import patch_configurable
@@ -198,6 +198,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
msg["task"]["path"],
msg["task"]["id"],
checkpoint=saved.checkpoint,
pending_writes=saved.pending_writes or [],
processes=graph.nodes,
channels=channels,
managed=managed,
@@ -211,6 +212,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
runner = PregelRunner(
submit=submit,
put_writes=partial(self._put_writes, submit, msg["config"]),
schedule_task=self._schedule_task,
)
async for _ in runner.atick([task], reraise=False):
pass
@@ -239,6 +241,14 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
)
await fut
def _schedule_task(
self,
task: PregelExecutableTask,
idx: int,
) -> None:
# will be scheduled by orchestrator when executor finishes
pass
def _put_writes(
self,
submit: Submit,
@@ -400,6 +410,7 @@ class KafkaExecutor(AbstractContextManager):
msg["task"]["path"],
msg["task"]["id"],
checkpoint=saved.checkpoint,
pending_writes=saved.pending_writes or [],
processes=graph.nodes,
channels=channels,
managed=managed,
@@ -412,6 +423,7 @@ class KafkaExecutor(AbstractContextManager):
runner = PregelRunner(
submit=submit,
put_writes=partial(self._put_writes, submit, msg["config"]),
schedule_task=self._schedule_task,
)
for _ in runner.tick([task], reraise=False):
pass
@@ -440,6 +452,14 @@ class KafkaExecutor(AbstractContextManager):
)
fut.result()
def _schedule_task(
self,
task: PregelExecutableTask,
idx: int,
) -> None:
# will be scheduled by orchestrator when executor finishes
pass
def _put_writes(
self,
submit: Submit,
libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py
+12 -12
View File
@@ -161,18 +161,18 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
specs=graph.channels,
output_keys=graph.output_channels,
stream_keys=graph.stream_channels,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
check_subgraphs=False,
) as loop:
if loop.tick(
input_keys=graph.input_channels,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
):
if loop.tick(input_keys=graph.input_channels):
# wait for checkpoint to be saved
if hasattr(loop, "_put_checkpoint_fut"):
await loop._put_checkpoint_fut
# schedule any new tasks
if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]:
if new_tasks := [
t for t in loop.tasks.values() if not t.scheduled and not t.writes
]:
# send messages to executor
futures = await asyncio.gather(
*(
@@ -351,18 +351,18 @@ class KafkaOrchestrator(AbstractContextManager):
specs=graph.channels,
output_keys=graph.output_channels,
stream_keys=graph.stream_channels,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
check_subgraphs=False,
) as loop:
if loop.tick(
input_keys=graph.input_channels,
interrupt_after=graph.interrupt_after_nodes,
interrupt_before=graph.interrupt_before_nodes,
):
if loop.tick(input_keys=graph.input_channels):
# wait for checkpoint to be saved
if hasattr(loop, "_put_checkpoint_fut"):
loop._put_checkpoint_fut.result()
# schedule any new tasks
if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]:
if new_tasks := [
t for t in loop.tasks.values() if not t.scheduled and not t.writes
]:
# send messages to executor
futures = [
self.producer.send(
libs/scheduler-kafka/langgraph/scheduler/kafka/types.py
+2 -2
View File
@@ -24,8 +24,8 @@ class MessageToOrchestrator(TypedDict):
class ExecutorTask(TypedDict):
id: str
path: tuple[str, ...]
id: Optional[str]
path: tuple[Union[str, int], ...]
class MessageToExecutor(TypedDict):
libs/scheduler-kafka/tests/test_push.py
+206
View File
@@ -0,0 +1,206 @@
import operator
from typing import (
Annotated,
Literal,
Union,
)
import pytest
from aiokafka import AIOKafkaProducer
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, StateGraph
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from langgraph.types import Control, Send
from tests.any import AnyDict
from tests.drain import drain_topics_async
pytestmark = pytest.mark.anyio
def mk_push_graph(
checkpointer: BaseCheckpointSaver,
) -> CompiledStateGraph:
# copied from test_send_dedupe_on_resume
class InterruptOnce:
ticks: int = 0
def __call__(self, state):
self.ticks += 1
if self.ticks == 1:
raise NodeInterrupt("Bahh")
return ["|".join(("flaky", str(state)))]
class Node:
def __init__(self, name: str):
self.name = name
self.ticks = 0
self.__name__ = name
def __call__(self, state):
self.ticks += 1
update = (
[self.name]
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, Control):
state.state = update
return state
else:
return update
def send_for_fun(state):
return [
Send("2", Control(send=Send("2", 3))),
Send("2", Control(send=Send("flaky", 4))),
"3.1",
]
def route_to_three(state) -> Literal["3"]:
return "3"
builder = StateGraph(Annotated[list, operator.add])
builder.add_node(Node("1"))
builder.add_node(Node("2"))
builder.add_node(Node("3"))
builder.add_node(Node("3.1"))
builder.add_node("flaky", InterruptOnce())
builder.add_edge(START, "1")
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
return builder.compile(checkpointer=checkpointer)
async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
input = ["0"]
config = {"configurable": {"thread_id": "1"}}
graph = mk_push_graph(acheckpointer)
graph_compare = mk_push_graph(acheckpointer)
# start a new run
async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
await producer.send_and_wait(
topics.orchestrator,
MessageToOrchestrator(input=input, config=config),
)
# drain topics
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check state
state = await graph.aget_state(config)
assert all(not t.error for t in state.tasks)
assert state.next == ("flaky",)
assert (
state.values
== await graph_compare.ainvoke(input, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|3",
]
)
# check history
history = [c async for c in graph.aget_state_history(config)]
assert len(history) == 2
# check messages
assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [
{
"config": {
"callbacks": None,
"configurable": {
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_ns": "",
"thread_id": "1",
},
"metadata": AnyDict(),
"recursion_limit": 25,
"tags": [],
},
"input": None,
"finally_send": None,
}
for c in reversed(history)
for _ in c.tasks
]
assert exec_msgs == [
{
"config": {
"callbacks": None,
"configurable": {
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_ns": "",
"thread_id": "1",
},
"metadata": AnyDict(),
"recursion_limit": 25,
"tags": [],
},
"task": {
"id": t.id,
"path": _convert_path(t.path),
},
"finally_send": None,
}
for c in reversed(history)
for t in c.tasks
]
# resume the thread
async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
await producer.send_and_wait(
topics.orchestrator,
MessageToOrchestrator(input=None, config=config),
)
orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
# check final state
state = await graph.aget_state(config)
assert state.next == ()
assert (
state.values
== await graph_compare.ainvoke(None, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
"3.1",
]
)
# check history
history = [c async for c in graph.aget_state_history(config)]
assert len(history) == 4
# check executions
# node "2" doesn't get called again, as we recover writes saved before
assert graph.builder.nodes["2"].runnable.func.ticks == 3
# node "flaky" gets called again, as it was interrupted
assert graph.builder.nodes["flaky"].runnable.func.ticks == 2
def _convert_path(
path: tuple[Union[str, int, tuple], ...],
) -> list[Union[str, int, list]]:
return list(_convert_path(p) if isinstance(p, tuple) else p for p in path)
libs/scheduler-kafka/tests/test_push_sync.py
+208
View File
@@ -0,0 +1,208 @@
import operator
from typing import (
Annotated,
Literal,
Union,
)
import pytest
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, StateGraph
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.default_sync import DefaultProducer
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from langgraph.types import Control, Send
from tests.any import AnyDict
from tests.drain import drain_topics
pytestmark = pytest.mark.anyio
def mk_push_graph(
checkpointer: BaseCheckpointSaver,
) -> CompiledStateGraph:
# copied from test_send_dedupe_on_resume
class InterruptOnce:
ticks: int = 0
def __call__(self, state):
self.ticks += 1
if self.ticks == 1:
raise NodeInterrupt("Bahh")
return ["|".join(("flaky", str(state)))]
class Node:
def __init__(self, name: str):
self.name = name
self.ticks = 0
self.__name__ = name
def __call__(self, state):
self.ticks += 1
update = (
[self.name]
if isinstance(state, list)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, Control):
state.state = update
return state
else:
return update
def send_for_fun(state):
return [
Send("2", Control(send=Send("2", 3))),
Send("2", Control(send=Send("flaky", 4))),
"3.1",
]
def route_to_three(state) -> Literal["3"]:
return "3"
builder = StateGraph(Annotated[list, operator.add])
builder.add_node(Node("1"))
builder.add_node(Node("2"))
builder.add_node(Node("3"))
builder.add_node(Node("3.1"))
builder.add_node("flaky", InterruptOnce())
builder.add_edge(START, "1")
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
return builder.compile(checkpointer=checkpointer)
def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
input = ["0"]
config = {"configurable": {"thread_id": "1"}}
graph = mk_push_graph(acheckpointer)
graph_compare = mk_push_graph(acheckpointer)
# start a new run
with DefaultProducer() as producer:
producer.send(
topics.orchestrator,
value=serde.dumps(MessageToOrchestrator(input=input, config=config)),
)
producer.flush()
# drain topics
orch_msgs, exec_msgs = drain_topics(topics, graph)
# check state
state = graph.get_state(config)
assert all(not t.error for t in state.tasks)
assert state.next == ("flaky",)
assert (
state.values
== graph_compare.invoke(input, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|3",
]
)
# check history
history = [c for c in graph.get_state_history(config)]
assert len(history) == 2
# check messages
assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [
{
"config": {
"callbacks": None,
"configurable": {
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_ns": "",
"thread_id": "1",
},
"metadata": AnyDict(),
"recursion_limit": 25,
"tags": [],
},
"input": None,
"finally_send": None,
}
for c in reversed(history)
for _ in c.tasks
]
assert exec_msgs == [
{
"config": {
"callbacks": None,
"configurable": {
"__pregel_ensure_latest": True,
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_ns": "",
"thread_id": "1",
},
"metadata": AnyDict(),
"recursion_limit": 25,
"tags": [],
},
"task": {
"id": t.id,
"path": _convert_path(t.path),
},
"finally_send": None,
}
for c in reversed(history)
for t in c.tasks
]
# resume the thread
with DefaultProducer() as producer:
producer.send(
topics.orchestrator,
value=serde.dumps(MessageToOrchestrator(input=None, config=config)),
)
producer.flush()
orch_msgs, exec_msgs = drain_topics(topics, graph)
# check final state
state = graph.get_state(config)
assert state.next == ()
assert (
state.values
== graph_compare.invoke(None, {"configurable": {"thread_id": "2"}})
== [
"0",
"1",
"2|Control(send=Send(node='2', arg=3))",
"2|Control(send=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
"3.1",
]
)
# check history
history = [c for c in graph.get_state_history(config)]
assert len(history) == 4
# check executions
# node "2" doesn't get called again, as we recover writes saved before
assert graph.builder.nodes["2"].runnable.func.ticks == 3
# node "flaky" gets called again, as it was interrupted
assert graph.builder.nodes["flaky"].runnable.func.ticks == 2
def _convert_path(
path: tuple[Union[str, int, tuple], ...],
) -> list[Union[str, int, list]]:
return list(_convert_path(p) if isinstance(p, tuple) else p for p in path)