Files
exo/dashboard
ciaranbor 6de14cfedb Support image generation cancellation (#1774)
## Motivation

Support cancelling image generation, similar to existing support for
cancelling text generation

## Changes

- Dashboard (app.svelte.ts): Wire up AbortController for both
generateImage and editImage API calls. On abort, show "Cancelled"
instead of an error. Clean up the controller in finally.
- Pipeline runner (pipeline/runner.py): Introduce a cancel_checker
callback and NaN-sentinel cancellation protocol for distributed
diffusion:
  - _check_cancellation() - only rank 0 polls the cancel callback
- _send() - replaces data with NaN sentinels when cancelling, so
downstream ranks detect cancellation via _recv_and_check()
  - _recv() / _recv_like() wrappers that eval and check for NaN sentinel
  - After cancellation, drains any pending ring recv to prevent deadlock
  - Skips partial image yields and final decode when cancelled
- Image runner (runner/image_models/runner.py): Deduplicate the
ImageGeneration and ImageEdits match arms into a shared
_run_image_task() method. Thread a cancel_checker closure (backed by the
existing cancel_receiver + cancelled_tasks set) into generate_image().
- Plumbing (distributed_model.py, generate.py): Pass cancel_checker
through the call chain.

## Why It Works

- Rank 0 is the only node that knows about task-level cancellation. When
it detects cancellation, it sends NaN tensors instead of real data.
Higher-order ranks detect the NaN sentinel on recv, set their own
_cancelling flag, and propagate NaN forward
- A drain step after the loop prevents the deadlock case where the last
rank already sent patches that the first would never consume.
- For single-node mode, the loop simply breaks immediately on
cancellation.

## Test Plan

### Automated Testing

New tests in src/exo/worker/tests/unittests/test_image
2026-03-25 16:56:04 +00:00
..
2025-12-17 12:22:22 +00:00
2025-10-22 11:56:52 +01:00
2025-10-22 11:56:52 +01:00
2025-08-25 16:41:13 +01:00
2025-10-22 11:56:52 +01:00