Fix missing interrupts in stream (#3886)

- When multiple parallel tasks and/or subgraphs emit interrupts some
were missing from stream output
This commit is contained in:
Nuno Campos
2025-03-17 18:51:42 -07:00
committed by GitHub
5 changed files with 52 additions and 16 deletions
+8 -2
View File
@@ -137,7 +137,12 @@ def map_debug_task_results(
"result": [
w for w in writes if w[0] in stream_channels_list or w[0] == RETURN
],
"interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT],
"interrupts": [
asdict(v)
for w in writes
if w[0] == INTERRUPT
for v in (w[1] if isinstance(w[1], Sequence) else [w[1]])
],
},
}
@@ -293,8 +298,9 @@ def tasks_w_writes(
),
tuple(
v
for tid, n, v in pending_writes
for tid, n, vv in pending_writes
if tid == task.id and n == INTERRUPT
for v in (vv if isinstance(vv, Sequence) else [vv])
),
states.get(task.id) if states else None,
(
+1 -1
View File
@@ -28,7 +28,7 @@ def is_task_id(task_id: str) -> bool:
"""Check if a string is a valid task id."""
try:
UUID(task_id)
except ValueError:
except Exception:
return False
return True
+27 -6
View File
@@ -794,11 +794,14 @@ class PregelLoop(LoopProtocol):
[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]}]),
)
# emit INTERRUPT if exception is empty (otherwise emitted by put_writes)
if exc_value is not None and (not exc_value.args or not exc_value.args[0]):
self._emit(
"updates",
lambda: iter(
[{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]
),
)
# save final output
self.output = read_channels(self.channels, self.output_keys)
# suppress interrupt
@@ -829,7 +832,25 @@ class PregelLoop(LoopProtocol):
"tags", EMPTY_SEQ
):
return
if writes[0][0] != ERROR and writes[0][0] != INTERRUPT:
if writes[0][0] == INTERRUPT:
self._emit(
"updates",
lambda: iter(
[
{
INTERRUPT: tuple(
v
for w in writes
if w[0] == INTERRUPT
for v in (
w[1] if isinstance(w[1], Sequence) else (w[1],)
)
)
}
]
),
)
elif writes[0][0] != ERROR:
self._emit(
"updates",
map_output_updates,
+13 -4
View File
@@ -543,10 +543,11 @@ class PregelRunner:
elif exception:
if isinstance(exception, GraphInterrupt):
# save interrupt to checkpointer
if interrupts := [(INTERRUPT, i) for i in exception.args[0]]:
if exception.args[0]:
writes = [(INTERRUPT, exception.args[0])]
if resumes := [w for w in task.writes if w[0] == RESUME]:
interrupts.extend(resumes)
self.put_writes(task.id, interrupts)
writes.extend(resumes)
self.put_writes(task.id, writes)
elif isinstance(exception, GraphBubbleUp):
raise exception
else:
@@ -608,6 +609,7 @@ def _panic_or_proceed(
done.add(fut)
else:
inflight.add(fut)
interrupts: list[GraphInterrupt] = []
while done:
# if any task failed
if exc := _exception(done.pop()):
@@ -616,7 +618,14 @@ def _panic_or_proceed(
inflight.pop().cancel()
# raise the exception
if panic:
raise exc
if isinstance(exc, GraphInterrupt):
# collect interrupts
interrupts.append(exc)
else:
raise exc
# raise combined interrupts
if interrupts:
raise GraphInterrupt(tuple(i for exc in interrupts for i in exc.args[0]))
if inflight:
# if we got here means we timed out
while inflight:
+3 -3
View File
@@ -939,9 +939,6 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
{"my_key": "value ⛰️", "market": "DE"}, thread2
)
] == [
{
"tool_one": {"my_key": " one"},
},
{
"__interrupt__": (
Interrupt(
@@ -951,6 +948,9 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
),
)
},
{
"tool_one": {"my_key": " one"},
},
]
# resume with answer
assert [