[Ayush Baluni (aayushbaluni)]

Fixes #5672

Cancelling a run mid-superstep does not persist the cancellation when durability="exit", so get_state().tasks[...].error is None and callers cannot tell a task that was cancelled in flight from one that never started. "async" and "sync" both record it correctly.

The inconsistency

Two nodes in one superstep, p1 fast and p2 slow; cancel while p2 is still running:

durability=async   tasks=[('p1', None), ('p2', 'CancelledError()')]
durability=sync    tasks=[('p1', None), ('p2', 'CancelledError()')]
durability=exit    tasks=[('p1', None), ('p2', None)]     <-- cancellation lost

Deterministic across repeated trials.

Root cause

PregelRunner.commit() appends (ERROR, exception) for a cancelled task (_runner.py:579-583), but it runs from a done-callback that fires after AsyncPregelLoop.__aexit__ unwinds the stack. Instrumenting the call order:

durability=exit
  commit(task=p1, exc=None)
  _suppress_interrupt(enter, exc=CancelledError)
  _put_pending_writes(tasks=1, channels=['steps'])   <-- the only flush
  _suppress_interrupt(exit)
  commit(task=p2, exc=CancelledError)                <-- ERROR write arrives after

With "async"/"sync", put_writes() forwards to the checkpointer immediately (_loop.py:466), so the record survives. With "exit", writes are only flushed by _put_pending_writes() inside _suppress_interrupt (_loop.py:1334), which has already run, so the write lands in checkpoint_pending_writes too late and is never persisted.

Fix

Stage the cancellation for still-running tasks before the exit flush. Tasks that already produced writes are skipped, so a completed node's work is untouched and is never misreported as cancelled.

How I verified this

• New parametrized regression test over all three durability modes.
• Mutation-checked: removing only the new call makes the exit case fail while async/sync still pass, so the test measures the fix rather than decorating it.
• tests/test_pregel_async.py: 313 passed (310 on unmodified main, +3 new).
• Full suite: 2012 passed, 6 skipped, 0 failed (2009 on unmodified main). The 208 errors are byte-identical before and after and are postgres/redis fixtures I have no local service for.
• ruff check and ruff format clean.

Scope

Deliberately narrow. It does not attempt the terminal-receipt / lineage contract discussed in the issue thread, and does not touch the Platform-side question of whether Agent Server awaits the loop's cancellation cleanup before publishing interrupted, which still looks Platform-owned.

Note on process: the issue is still unassigned despite several contributors asking. Happy to close this if you would rather assign it to someone already waiting, or to re-scope it, since the defect fixed here is narrower than the originally reported symptom (which I could not reproduce on current main; details in my comment).

[github-actions]

This PR has been automatically closed because you are not assigned to the linked issue.

External contributors must be assigned to an issue before opening a PR for it. Please:

1. Comment on the linked issue to request assignment from a maintainer
2. Once assigned, your PR will be reopened automatically

Maintainers: reopen this PR or remove the missing-issue-link label to bypass this check.