Compare commits

...
Author SHA1 Message Date
Sydney Runkle 26d279a0ac rename 2026-03-05 09:30:02 -08:00
Sydney Runkle e850b21d08 continue 2026-03-05 09:10:43 -08:00
Sydney Runkle 6a92b7ff3c alt fix idea 2026-03-05 08:54:10 -08:00
Sydney Runkle 207dccf5b3 contextvars 2026-03-04 22:28:18 -08:00
Sydney Runkle 0623e4690c boom lint 2026-03-04 22:19:29 -08:00
Sydney Runkle 1366210740 better comments 2026-03-04 22:11:10 -08:00
Sydney Runkle b53c47675d refactor tests 2026-03-04 22:08:14 -08:00
Sydney Runkle 1aeafeeebd move tests 2026-03-04 21:45:16 -08:00
Sydney Runkle ba2b2f4a6f going crazy w/ tests 2026-03-04 21:26:51 -08:00
Sydney Runkle 61fb3563b4 maybe a fix 2026-03-04 18:18:06 -08:00
Sydney Runkle 63528f25af more tests 2026-03-04 17:36:15 -08:00
Sydney Runkle a59b3f1fee update 2026-03-04 14:26:15 -08:00
Sydney Runkle eeaac6d80d update comments 2026-03-04 13:26:31 -08:00
Sydney Runkle bb41c66547 lint 2026-03-04 13:22:48 -08:00
Sydney Runkle 52b586370d skip maybe 2026-03-04 13:20:52 -08:00
4 changed files with 3230 additions and 22 deletions
@@ -41,6 +41,8 @@ CONFIG_KEY_CACHE = sys.intern("__pregel_cache")
# holds a `BaseCache` made available to subgraphs
CONFIG_KEY_RESUMING = sys.intern("__pregel_resuming")
# holds a boolean indicating if subgraphs should resume from a previous checkpoint
CONFIG_KEY_REPLAYING = sys.intern("__pregel_replaying")
# holds a boolean indicating if subgraphs should replay (re-run tasks, drop cached RESUME writes)
CONFIG_KEY_TASK_ID = sys.intern("__pregel_task_id")
# holds the task ID for the current task
CONFIG_KEY_THREAD_ID = sys.intern("thread_id")
@@ -98,6 +100,7 @@ RESERVED = {
CONFIG_KEY_STREAM,
CONFIG_KEY_CHECKPOINT_MAP,
CONFIG_KEY_RESUMING,
CONFIG_KEY_REPLAYING,
CONFIG_KEY_TASK_ID,
CONFIG_KEY_CHECKPOINT_MAP,
CONFIG_KEY_CHECKPOINT_ID,
+78 -22
View File
@@ -42,6 +42,7 @@ from langgraph._internal._constants import (
CONFIG_KEY_CHECKPOINT_ID,
CONFIG_KEY_CHECKPOINT_MAP,
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_REPLAYING,
CONFIG_KEY_RESUME_MAP,
CONFIG_KEY_RESUMING,
CONFIG_KEY_SCRATCHPAD,
@@ -152,7 +153,7 @@ class PregelLoop:
input_keys: str | Sequence[str]
output_keys: str | Sequence[str]
stream_keys: str | Sequence[str]
skip_done_tasks: bool
is_replaying: bool
is_nested: bool
manager: None | AsyncParentRunManager | ParentRunManager
interrupt_after: All | Sequence[str]
@@ -244,7 +245,9 @@ class PregelLoop:
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]
self.is_replaying = CONFIG_KEY_CHECKPOINT_ID in config[
CONF
] or config[CONF].get(CONFIG_KEY_REPLAYING, False)
self._migrate_checkpoint = migrate_checkpoint
self.trigger_to_nodes = trigger_to_nodes
self.retry_policy = retry_policy
@@ -451,7 +454,7 @@ class PregelLoop:
# save the new task
self.tasks[pushed.id] = pushed
# match any pending writes to the new task
if self.skip_done_tasks:
if not self.is_replaying:
self._match_writes({pushed.id: pushed})
# return the new task, to be started if not run before
return pushed
@@ -515,7 +518,7 @@ class PregelLoop:
return False
# if there are pending writes from a previous loop, apply them
if self.skip_done_tasks and self.checkpoint_pending_writes:
if not self.is_replaying and self.checkpoint_pending_writes:
self._match_writes(self.tasks)
# before execution, check if we should interrupt
@@ -557,8 +560,8 @@ class PregelLoop:
)
# clear pending writes
self.checkpoint_pending_writes.clear()
# "not skip_done_tasks" only applies to first tick after resuming
self.skip_done_tasks = True
# only replay (re-execute) done tasks on the first tick
self.is_replaying = False
# save checkpoint
self._put_checkpoint({"source": "loop"})
# after execution, check if we should interrupt
@@ -567,8 +570,9 @@ class PregelLoop:
):
self.status = "interrupt_after"
raise GraphInterrupt()
# unset resuming flag
# unset resuming/replaying flags
self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
self.config[CONF].pop(CONFIG_KEY_REPLAYING, None)
def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
raise NotImplementedError
@@ -618,22 +622,28 @@ class PregelLoop:
def _first(
self, *, input_keys: str | Sequence[str], updated_channels: set[str] | None
) -> set[str] | None:
# resuming from previous checkpoint requires
# - finding a previous checkpoint
# - receiving None input (outer graph) or RESUMING flag (subgraph)
# Resuming from a previous checkpoint requires two things:
# 1. A prior checkpoint exists (channel_versions is non-empty)
# 2. The input signals continuation (not a fresh run with new input)
configurable = self.config.get(CONF, {})
is_resuming = bool(self.checkpoint["channel_versions"]) and bool(
configurable.get(
CONFIG_KEY_RESUMING,
self.input is None
or isinstance(self.input, Command)
or (
not self.is_nested
and self.config.get("metadata", {}).get("run_id")
== self.checkpoint_metadata.get("run_id", MISSING)
),
has_prior_checkpoint = bool(self.checkpoint["channel_versions"])
# For subgraphs, the parent explicitly sets CONFIG_KEY_RESUMING.
# For the outer graph, we infer from the input:
# - None input: resume after interrupt (invoke(None, config))
# - Command input: any Command operates on existing state
# - Same run_id: re-entry into an ongoing run (e.g. stream reconnect)
input_signals_resume = (
self.input is None
or isinstance(self.input, Command)
or (
not self.is_nested
and self.config.get("metadata", {}).get("run_id")
== self.checkpoint_metadata.get("run_id", MISSING)
)
)
is_resuming = has_prior_checkpoint and bool(
configurable.get(CONFIG_KEY_RESUMING, input_signals_resume)
)
# map command to writes
if isinstance(self.input, Command):
@@ -723,10 +733,14 @@ class PregelLoop:
self._put_checkpoint({"source": "input"})
elif CONFIG_KEY_RESUMING not in configurable:
raise EmptyInputError(f"Received no input for {input_keys}")
# update config
# Propagate resuming and replaying flags to subgraphs.
if not self.is_nested:
self.config = patch_configurable(
self.config, {CONFIG_KEY_RESUMING: is_resuming}
self.config,
{
CONFIG_KEY_RESUMING: is_resuming,
CONFIG_KEY_REPLAYING: self.is_replaying,
},
)
# set flag
self.status = "pending"
@@ -1085,6 +1099,14 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
saved = self.checkpointer.get_tuple(self.checkpoint_config)
else:
saved = None
# When replaying a subgraph that wasn't in the checkpoint map
# (parent checkpoint predates this subgraph), start fresh.
if (
saved is not None
and self.config[CONF].get(CONFIG_KEY_REPLAYING)
and not self.checkpoint_config.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_ID)
):
saved = None
if saved is None:
saved = CheckpointTuple(
self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
@@ -1109,6 +1131,19 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
if saved.pending_writes is not None
else []
)
# When replaying from a specific checkpoint, drop cached RESUME
# writes so that interrupt() calls re-fire instead of returning
# stale values. But if a resume value is being provided (e.g.
# Command(resume=...) on a specific checkpoint), keep them —
# multi-interrupt scenarios need previously resolved values preserved.
if self.is_replaying:
has_resume_value = (
isinstance(self.input, Command) and self.input.resume is not None
) or self.config.get(CONF, {}).get(CONFIG_KEY_RESUMING, False)
if not has_resume_value:
self.checkpoint_pending_writes = [
w for w in self.checkpoint_pending_writes if w[1] != RESUME
]
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
self.channels, self.managed = channels_from_checkpoint(
@@ -1264,6 +1299,14 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
else:
saved = None
# When replaying a subgraph that wasn't in the checkpoint map
# (parent checkpoint predates this subgraph), start fresh.
if (
saved is not None
and self.config[CONF].get(CONFIG_KEY_REPLAYING)
and not self.checkpoint_config.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_ID)
):
saved = None
if saved is None:
saved = CheckpointTuple(
self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
@@ -1288,6 +1331,19 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
if saved.pending_writes is not None
else []
)
# When replaying from a specific checkpoint, drop cached RESUME
# writes so that interrupt() calls re-fire instead of returning
# stale values. But if a resume value is being provided (e.g.
# Command(resume=...) on a specific checkpoint), keep them —
# multi-interrupt scenarios need previously resolved values preserved.
if self.is_replaying:
has_resume_value = (
isinstance(self.input, Command) and self.input.resume is not None
) or self.config.get(CONF, {}).get(CONFIG_KEY_RESUMING, False)
if not has_resume_value:
self.checkpoint_pending_writes = [
w for w in self.checkpoint_pending_writes if w[1] != RESUME
]
self.submit = await self.stack.enter_async_context(
AsyncBackgroundExecutor(self.config)
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff