maybe a fix

This commit is contained in:
Sydney Runkle
2026-03-04 18:18:06 -08:00
parent 63528f25af
commit 61fb3563b4
3 changed files with 48 additions and 40 deletions
+16 -7
View File
@@ -725,8 +725,17 @@ class PregelLoop:
raise EmptyInputError(f"Received no input for {input_keys}")
# update config
if not self.is_nested:
has_resume = (
isinstance(self.input, Command) and self.input.resume is not None
)
# Propagate CONFIG_KEY_RESUMING to subgraphs:
# - True: genuine resume → subgraphs resume and preserve RESUME writes
# - False: fork from old checkpoint → subgraphs start fresh
# skip_done_tasks is False when a specific checkpoint_id was provided
# (fork scenario), True when resuming from latest checkpoint.
self.config = patch_configurable(
self.config, {CONFIG_KEY_RESUMING: is_resuming}
self.config,
{CONFIG_KEY_RESUMING: has_resume or (is_resuming and self.skip_done_tasks)},
)
# set flag
self.status = "pending"
@@ -1114,9 +1123,9 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
# We must NOT drop them when resuming (subgraph via CONFIG_KEY_RESUMING,
# or top graph via Command(resume=...)) because with multiple interrupts
# previously resolved RESUME values need to be preserved.
is_resuming = CONFIG_KEY_RESUMING in self.config.get(CONF, {}) or (
isinstance(self.input, Command) and self.input.resume is not None
)
is_resuming = (
self.config.get(CONF, {}).get(CONFIG_KEY_RESUMING) is True
) or (isinstance(self.input, Command) and self.input.resume is not None)
if not self.skip_done_tasks and not is_resuming:
self.checkpoint_pending_writes = [
w for w in self.checkpoint_pending_writes if w[1] != RESUME
@@ -1305,9 +1314,9 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
# We must NOT drop them when resuming (subgraph via CONFIG_KEY_RESUMING,
# or top graph via Command(resume=...)) because with multiple interrupts
# previously resolved RESUME values need to be preserved.
is_resuming = CONFIG_KEY_RESUMING in self.config.get(CONF, {}) or (
isinstance(self.input, Command) and self.input.resume is not None
)
is_resuming = (
self.config.get(CONF, {}).get(CONFIG_KEY_RESUMING) is True
) or (isinstance(self.input, Command) and self.input.resume is not None)
if not self.skip_done_tasks and not is_resuming:
self.checkpoint_pending_writes = [
w for w in self.checkpoint_pending_writes if w[1] != RESUME
+17 -19
View File
@@ -5576,10 +5576,9 @@ def test_fork_subgraph_interrupt_no_checkpointer(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Fork/replay with a subgraph that has no checkpointer (checkpointer=False/None).
The subgraph inherits the parent's checkpointer via config and its saved
checkpoint retains RESUME writes (CONFIG_KEY_RESUMING is propagated from
the parent). So the subgraph does NOT re-fire the interrupt on replay
it uses the cached resume value and completes."""
On fork (input=None), the parent propagates CONFIG_KEY_RESUMING=False so the
subgraph strips RESUME writes and re-fires the interrupt consistent with
top-level interrupt behavior on fork/replay."""
called: list[str] = []
@@ -5630,21 +5629,20 @@ def test_fork_subgraph_interrupt_no_checkpointer(
history = list(graph.get_state_history(config))
before_sub = [s for s in history if s.next == ("call_subgraph",)][-1]
# 4. Replay — subgraph uses cached resume value, does NOT re-fire interrupt
# 4. Replay — subgraph re-fires interrupt (consistent with top-level behavior)
called.clear()
replay_result = graph.invoke(None, before_sub.config)
assert "__interrupt__" not in replay_result
assert replay_result == {"value": ["sub:answer", "after"]}
assert "__interrupt__" in replay_result
assert replay_result["__interrupt__"][0].value == "Sub question?"
assert "call_subgraph" in called
assert "after" in called
assert "after" not in called
def test_fork_subgraph_interrupt_checkpointer_true(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Fork/replay with a subgraph that has checkpointer=True.
Same behavior as no checkpointer the subgraph's checkpoint retains
RESUME writes and the interrupt does NOT re-fire on replay."""
On fork, subgraph re-fires interrupts consistent with top-level behavior."""
called: list[str] = []
@@ -5701,21 +5699,21 @@ def test_fork_subgraph_interrupt_checkpointer_true(
history = list(graph.get_state_history(config))
before_sub = [s for s in history if s.next == ("call_subgraph",)][-1]
# 4. Replay — subgraph uses cached resume value, does NOT re-fire interrupt
# 4. Replay — subgraph re-fires interrupt (consistent with top-level behavior)
called.clear()
replay_result = graph.invoke(None, before_sub.config)
assert "__interrupt__" not in replay_result
assert replay_result == {"value": ["sub_node", "sub:answer", "after"]}
assert "__interrupt__" in replay_result
assert replay_result["__interrupt__"][0].value == "Sub question?"
assert "call_subgraph" in called
assert "after" in called
assert "after" not in called
def test_fork_subgraph_two_interrupts_no_checkpointer(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Fork/replay with a subgraph (no checkpointer) containing two interrupt
nodes. Same as single interrupt the subgraph uses cached resume values
and completes without re-firing interrupts."""
nodes. On fork, subgraph re-fires the first interrupt consistent with
top-level behavior."""
called: list[str] = []
@@ -5768,14 +5766,14 @@ def test_fork_subgraph_two_interrupts_no_checkpointer(
result = graph.invoke(Command(resume="a2"), config)
assert result == {"value": ["s1:a1", "s2:a2"]}
# 4. Replay from before subgraph — uses cached resume values
# 4. Replay from before subgraph — re-fires first interrupt
history = list(graph.get_state_history(config))
before_sub = [s for s in history if s.next == ("call_subgraph",)][-1]
called.clear()
replay_result = graph.invoke(None, before_sub.config)
assert "__interrupt__" not in replay_result
assert replay_result == {"value": ["s1:a1", "s2:a2"]}
assert "__interrupt__" in replay_result
assert replay_result["__interrupt__"][0].value == "Sub Q1?"
def test_concurrent_execution_thread_safety():
+15 -14
View File
@@ -6779,9 +6779,9 @@ async def test_fork_after_all_interrupts(
async def test_fork_subgraph_interrupt_no_checkpointer(
async_checkpointer: BaseCheckpointSaver,
) -> None:
"""Fork/replay with subgraph (no checkpointer). The subgraph inherits the
parent's checkpointer and CONFIG_KEY_RESUMING, so it uses cached resume
values and does NOT re-fire the interrupt."""
"""Fork/replay with subgraph (no checkpointer). On fork, the parent
propagates CONFIG_KEY_RESUMING=False so the subgraph strips RESUME writes
and re-fires the interrupt consistent with top-level behavior."""
called: list[str] = []
@@ -6831,17 +6831,17 @@ async def test_fork_subgraph_interrupt_no_checkpointer(
called.clear()
replay_result = await graph.ainvoke(None, before_sub.config)
assert "__interrupt__" not in replay_result
assert replay_result == {"value": ["sub:answer", "after"]}
assert "__interrupt__" in replay_result
assert replay_result["__interrupt__"][0].value == "Sub question?"
assert "call_subgraph" in called
assert "after" in called
assert "after" not in called
async def test_fork_subgraph_interrupt_checkpointer_true(
async_checkpointer: BaseCheckpointSaver,
) -> None:
"""Fork/replay with subgraph (checkpointer=True). Same as no checkpointer —
uses cached resume values, does NOT re-fire the interrupt."""
"""Fork/replay with subgraph (checkpointer=True). On fork, subgraph
re-fires interrupts consistent with top-level behavior."""
called: list[str] = []
@@ -6897,17 +6897,18 @@ async def test_fork_subgraph_interrupt_checkpointer_true(
called.clear()
replay_result = await graph.ainvoke(None, before_sub.config)
assert "__interrupt__" not in replay_result
assert replay_result == {"value": ["sub_node", "sub:answer", "after"]}
assert "__interrupt__" in replay_result
assert replay_result["__interrupt__"][0].value == "Sub question?"
assert "call_subgraph" in called
assert "after" in called
assert "after" not in called
async def test_fork_subgraph_two_interrupts_no_checkpointer(
async_checkpointer: BaseCheckpointSaver,
) -> None:
"""Fork/replay with subgraph (no checkpointer) with two interrupt nodes.
Same behavior uses cached resume values, does NOT re-fire interrupts."""
On fork, subgraph re-fires the first interrupt consistent with
top-level behavior."""
called: list[str] = []
@@ -6962,8 +6963,8 @@ async def test_fork_subgraph_two_interrupts_no_checkpointer(
called.clear()
replay_result = await graph.ainvoke(None, before_sub.config)
assert "__interrupt__" not in replay_result
assert replay_result == {"value": ["s1:a1", "s2:a2"]}
assert "__interrupt__" in replay_result
assert replay_result["__interrupt__"][0].value == "Sub Q1?"
async def test_concurrent_execution():