Compare commits

...
Author SHA1 Message Date
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
3 changed files with 3194 additions and 15 deletions
+55 -15
View File
@@ -618,22 +618,28 @@ class PregelLoop:
def _first( def _first(
self, *, input_keys: str | Sequence[str], updated_channels: set[str] | None self, *, input_keys: str | Sequence[str], updated_channels: set[str] | None
) -> set[str] | None: ) -> set[str] | None:
# resuming from previous checkpoint requires # Resuming from a previous checkpoint requires two things:
# - finding a previous checkpoint # 1. A prior checkpoint exists (channel_versions is non-empty)
# - receiving None input (outer graph) or RESUMING flag (subgraph) # 2. The input signals continuation (not a fresh run with new input)
configurable = self.config.get(CONF, {}) configurable = self.config.get(CONF, {})
is_resuming = bool(self.checkpoint["channel_versions"]) and bool( has_prior_checkpoint = bool(self.checkpoint["channel_versions"])
configurable.get( # For subgraphs, the parent explicitly sets CONFIG_KEY_RESUMING.
CONFIG_KEY_RESUMING, # For the outer graph, we infer from the input:
self.input is None # - None input: resume after interrupt (invoke(None, config))
or isinstance(self.input, Command) # - Command input: any Command operates on existing state
or ( # - Same run_id: re-entry into an ongoing run (e.g. stream reconnect)
not self.is_nested input_signals_resume = (
and self.config.get("metadata", {}).get("run_id") self.input is None
== self.checkpoint_metadata.get("run_id", MISSING) 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 # map command to writes
if isinstance(self.input, Command): if isinstance(self.input, Command):
@@ -723,10 +729,18 @@ class PregelLoop:
self._put_checkpoint({"source": "input"}) self._put_checkpoint({"source": "input"})
elif CONFIG_KEY_RESUMING not in configurable: elif CONFIG_KEY_RESUMING not in configurable:
raise EmptyInputError(f"Received no input for {input_keys}") raise EmptyInputError(f"Received no input for {input_keys}")
# update config # Propagate resuming flag to subgraphs (only the outer graph does this).
if not self.is_nested: if not self.is_nested:
has_resume_value = (
isinstance(self.input, Command) and self.input.resume is not None
)
# When forking (skip_done_tasks=False, i.e. specific checkpoint_id),
# subgraphs should NOT resume — they start fresh.
# When genuinely resuming from latest, subgraphs should also resume.
is_fork = not self.skip_done_tasks
subgraph_should_resume = has_resume_value or (is_resuming and not is_fork)
self.config = patch_configurable( self.config = patch_configurable(
self.config, {CONFIG_KEY_RESUMING: is_resuming} self.config, {CONFIG_KEY_RESUMING: subgraph_should_resume}
) )
# set flag # set flag
self.status = "pending" self.status = "pending"
@@ -1109,6 +1123,19 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
if saved.pending_writes is not None if saved.pending_writes is not None
else [] else []
) )
# When replaying from a specific checkpoint (fork), drop cached
# RESUME writes so that interrupt() calls re-fire instead of
# returning stale values. But if the input directly carries a
# resume value, keep them — multi-interrupt scenarios need
# previously resolved RESUME values preserved.
is_replaying = not self.skip_done_tasks
has_resume_value = (
self.config.get(CONF, {}).get(CONFIG_KEY_RESUMING) is True
) or (isinstance(self.input, Command) and self.input.resume is not None)
if is_replaying and 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.submit = self.stack.enter_context(BackgroundExecutor(self.config))
self.channels, self.managed = channels_from_checkpoint( self.channels, self.managed = channels_from_checkpoint(
@@ -1288,6 +1315,19 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
if saved.pending_writes is not None if saved.pending_writes is not None
else [] else []
) )
# When replaying from a specific checkpoint (fork), drop cached
# RESUME writes so that interrupt() calls re-fire instead of
# returning stale values. But if the input directly carries a
# resume value, keep them — multi-interrupt scenarios need
# previously resolved RESUME values preserved.
is_replaying = not self.skip_done_tasks
has_resume_value = (
self.config.get(CONF, {}).get(CONFIG_KEY_RESUMING) is True
) or (isinstance(self.input, Command) and self.input.resume is not None)
if is_replaying and 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( self.submit = await self.stack.enter_async_context(
AsyncBackgroundExecutor(self.config) AsyncBackgroundExecutor(self.config)
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff