← back to Exo
Support image generation cancellation (#1774)
6de14cfedb806b161aff793ac51b3750c191035b · 2026-03-25 16:56:04 +0000 · ciaranbor
## 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
Files touched
M dashboard/src/lib/stores/app.svelte.tsM src/exo/worker/engines/image/distributed_model.pyM src/exo/worker/engines/image/generate.pyM src/exo/worker/engines/image/pipeline/runner.pyM src/exo/worker/runner/image_models/runner.pyA src/exo/worker/tests/unittests/test_image/__init__.pyA src/exo/worker/tests/unittests/test_image/test_cancellation.py
Diff
commit 6de14cfedb806b161aff793ac51b3750c191035b
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date: Wed Mar 25 16:56:04 2026 +0000
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
---
dashboard/src/lib/stores/app.svelte.ts | 62 ++-
src/exo/worker/engines/image/distributed_model.py | 4 +-
src/exo/worker/engines/image/generate.py | 3 +
src/exo/worker/engines/image/pipeline/runner.py | 157 ++++----
src/exo/worker/runner/image_models/runner.py | 200 ++++------
.../worker/tests/unittests/test_image/__init__.py | 0
.../unittests/test_image/test_cancellation.py | 433 +++++++++++++++++++++
7 files changed, 660 insertions(+), 199 deletions(-)
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index 12ca231f..65c7e20c 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -2650,6 +2650,9 @@ class AppStore {
this.syncActiveMessagesIfNeeded(targetConversationId);
this.saveConversationsToStorage();
+ const abortController = new AbortController();
+ this.currentAbortController = abortController;
+
try {
// Determine the model to use
const model = this.getModelForRequest(modelId);
@@ -2704,6 +2707,7 @@ class AppStore {
"Content-Type": "application/json",
},
body: JSON.stringify(requestBody),
+ signal: abortController.signal,
});
if (!response.ok) {
@@ -2843,14 +2847,27 @@ class AppStore {
);
}
} catch (error) {
- console.error("Error generating image:", error);
- this.handleStreamingError(
- error,
- targetConversationId,
- assistantMessage.id,
- "Failed to generate image",
- );
+ if (abortController.signal.aborted) {
+ this.updateConversationMessage(
+ targetConversationId,
+ assistantMessage.id,
+ (msg) => {
+ msg.content = "Cancelled";
+ msg.attachments = [];
+ },
+ );
+ this.syncActiveMessagesIfNeeded(targetConversationId);
+ } else {
+ console.error("Error generating image:", error);
+ this.handleStreamingError(
+ error,
+ targetConversationId,
+ assistantMessage.id,
+ "Failed to generate image",
+ );
+ }
} finally {
+ this.currentAbortController = null;
this.isLoading = false;
this.saveConversationsToStorage();
}
@@ -2914,6 +2931,9 @@ class AppStore {
// Clear editing state
this.editingImage = null;
+ const abortController = new AbortController();
+ this.currentAbortController = abortController;
+
try {
// Determine the model to use
const model = this.getModelForRequest(modelId);
@@ -2975,6 +2995,7 @@ class AppStore {
const apiResponse = await fetch("/v1/images/edits", {
method: "POST",
body: formData,
+ signal: abortController.signal,
});
if (!apiResponse.ok) {
@@ -3075,14 +3096,27 @@ class AppStore {
);
}
} catch (error) {
- console.error("Error editing image:", error);
- this.handleStreamingError(
- error,
- targetConversationId,
- assistantMessage.id,
- "Failed to edit image",
- );
+ if (abortController.signal.aborted) {
+ this.updateConversationMessage(
+ targetConversationId,
+ assistantMessage.id,
+ (msg) => {
+ msg.content = "Cancelled";
+ msg.attachments = [];
+ },
+ );
+ this.syncActiveMessagesIfNeeded(targetConversationId);
+ } else {
+ console.error("Error editing image:", error);
+ this.handleStreamingError(
+ error,
+ targetConversationId,
+ assistantMessage.id,
+ "Failed to edit image",
+ );
+ }
} finally {
+ this.currentAbortController = null;
this.isLoading = false;
this.saveConversationsToStorage();
}
diff --git a/src/exo/worker/engines/image/distributed_model.py b/src/exo/worker/engines/image/distributed_model.py
index 40a4fd38..7d866a3a 100644
--- a/src/exo/worker/engines/image/distributed_model.py
+++ b/src/exo/worker/engines/image/distributed_model.py
@@ -1,4 +1,4 @@
-from collections.abc import Generator
+from collections.abc import Callable, Generator
from pathlib import Path
from typing import Any, Literal, Optional
@@ -116,6 +116,7 @@ class DistributedImageModel:
image_path: Path | None = None,
partial_images: int = 0,
advanced_params: AdvancedImageParams | None = None,
+ cancel_checker: Callable[[], bool] | None = None,
) -> Generator[Image.Image | tuple[Image.Image, int, int], None, None]:
if (
advanced_params is not None
@@ -163,6 +164,7 @@ class DistributedImageModel:
guidance_override=guidance_override,
negative_prompt=negative_prompt,
num_sync_steps=num_sync_steps,
+ cancel_checker=cancel_checker,
):
if isinstance(result, tuple):
# Partial image: (GeneratedImage, partial_index, total_partials)
diff --git a/src/exo/worker/engines/image/generate.py b/src/exo/worker/engines/image/generate.py
index 15523e51..fbc23815 100644
--- a/src/exo/worker/engines/image/generate.py
+++ b/src/exo/worker/engines/image/generate.py
@@ -3,6 +3,7 @@ import io
import random
import tempfile
import time
+from collections.abc import Callable
from pathlib import Path
from typing import Generator, Literal
@@ -69,6 +70,7 @@ def warmup_image_generator(model: DistributedImageModel) -> Image.Image | None:
def generate_image(
model: DistributedImageModel,
task: ImageGenerationTaskParams | ImageEditsTaskParams,
+ cancel_checker: Callable[[], bool] | None = None,
) -> Generator[ImageGenerationResponse | PartialImageResponse, None, None]:
"""Generate image(s), optionally yielding partial results.
@@ -127,6 +129,7 @@ def generate_image(
image_path=image_path,
partial_images=partial_images,
advanced_params=advanced_params,
+ cancel_checker=cancel_checker,
):
if isinstance(result, tuple):
# Partial image: (Image, partial_index, total_partials)
diff --git a/src/exo/worker/engines/image/pipeline/runner.py b/src/exo/worker/engines/image/pipeline/runner.py
index 01ad7b05..2390d991 100644
--- a/src/exo/worker/engines/image/pipeline/runner.py
+++ b/src/exo/worker/engines/image/pipeline/runner.py
@@ -1,4 +1,4 @@
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
from dataclasses import dataclass
from math import ceil
from typing import Any, Optional, final
@@ -100,6 +100,8 @@ class DiffusionRunner:
self.total_layers = config.total_blocks
self._guidance_override: float | None = None
+ self._cancel_checker: Callable[[], bool] | None = None
+ self._cancelling: bool = False
self._compute_assigned_blocks()
@@ -240,6 +242,43 @@ class DiffusionRunner:
def is_distributed(self) -> bool:
return self.group is not None
+ def _is_sentinel(self, tensor: mx.array) -> bool:
+ return bool(mx.all(mx.isnan(tensor)).item())
+
+ def _check_cancellation(self) -> None:
+ if self._cancelling:
+ return
+ if (
+ self.is_first_stage
+ and self._cancel_checker is not None
+ and self._cancel_checker()
+ ):
+ self._cancelling = True
+
+ def _send(self, data: mx.array, dst: int) -> mx.array:
+ assert self.group is not None
+ if self._cancelling:
+ data = mx.full(data.shape, float("nan"), dtype=data.dtype)
+ return mx.distributed.send(data, dst, group=self.group)
+
+ def _recv_and_check(self, result: mx.array) -> mx.array:
+ mx.eval(result)
+ if self._is_sentinel(result):
+ self._cancelling = True
+ return result
+
+ def _recv(self, shape: tuple[int, ...], dtype: mx.Dtype, src: int) -> mx.array:
+ assert self.group is not None
+ return self._recv_and_check(
+ mx.distributed.recv(shape, dtype, src, group=self.group)
+ )
+
+ def _recv_like(self, template: mx.array, src: int) -> mx.array:
+ assert self.group is not None
+ return self._recv_and_check(
+ mx.distributed.recv_like(template, src, group=self.group)
+ )
+
def _get_effective_guidance_scale(self) -> float | None:
if self._guidance_override is not None:
return self._guidance_override
@@ -313,19 +352,13 @@ class DiffusionRunner:
assert self.cfg_peer_rank is not None
if is_positive:
- noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group)
+ noise = self._send(noise, self.cfg_peer_rank)
mx.async_eval(noise)
- noise_neg = mx.distributed.recv_like(
- noise, self.cfg_peer_rank, group=self.group
- )
- mx.eval(noise_neg)
+ noise_neg = self._recv_like(noise, src=self.cfg_peer_rank)
noise_pos = noise
else:
- noise_pos = mx.distributed.recv_like(
- noise, self.cfg_peer_rank, group=self.group
- )
- mx.eval(noise_pos)
- noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group)
+ noise_pos = self._recv_like(noise, src=self.cfg_peer_rank)
+ noise = self._send(noise, self.cfg_peer_rank)
mx.async_eval(noise)
noise_neg = noise
@@ -432,6 +465,7 @@ class DiffusionRunner:
guidance_override: float | None = None,
negative_prompt: str | None = None,
num_sync_steps: int = 1,
+ cancel_checker: Callable[[], bool] | None = None,
):
"""Primary entry point for image generation.
@@ -454,6 +488,8 @@ class DiffusionRunner:
Final GeneratedImage
"""
self._guidance_override = guidance_override
+ self._cancel_checker = cancel_checker
+ self._cancelling = False
latents = self.adapter.create_latents(seed, runtime_config)
prompt_data = self.adapter.encode_prompt(prompt, negative_prompt)
@@ -495,7 +531,7 @@ class DiffusionRunner:
except StopIteration as e:
latents = e.value # pyright: ignore[reportAny]
- if self.is_last_stage:
+ if self.is_last_stage and not self._cancelling:
yield self.adapter.decode_latents(latents, runtime_config, seed, prompt) # pyright: ignore[reportAny]
def _run_diffusion_loop(
@@ -524,7 +560,12 @@ class DiffusionRunner:
latents=latents,
)
+ t = -1 # default if time_steps is empty; drain condition uses t
for t in time_steps:
+ self._check_cancellation()
+ if self._cancelling and self.group is None:
+ break
+
try:
latents = self._diffusion_step(
t=t,
@@ -542,7 +583,7 @@ class DiffusionRunner:
mx.eval(latents)
- if t in capture_steps and self.is_last_stage:
+ if t in capture_steps and self.is_last_stage and not self._cancelling:
yield (latents, t)
except KeyboardInterrupt: # noqa: PERF203
@@ -551,6 +592,24 @@ class DiffusionRunner:
f"Stopping image generation at step {t + 1}/{len(time_steps)}"
) from None
+ if self._cancelling:
+ break
+
+ # Drain pending ring recvs after cancellation during async steps.
+ # The last stage sent patches during the final completed step, but
+ # the first stage will never enter the next step to recv them.
+ if (
+ self._cancelling
+ and self.is_first_stage
+ and not self.is_last_stage
+ and self.group is not None
+ and t >= runtime_config.init_time_step + num_sync_steps
+ and t != runtime_config.num_inference_steps - 1
+ ):
+ patch_latents_drain, _ = self._create_patches(latents, runtime_config)
+ for patch in patch_latents_drain:
+ self._recv_like(patch, src=self.last_pipeline_rank)
+
ctx.after_loop(latents=latents) # pyright: ignore[reportAny]
return latents
@@ -777,19 +836,16 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- hidden_states = mx.distributed.recv(
+ hidden_states = self._recv(
(batch_size, num_img_tokens, hidden_dim),
dtype,
self.prev_pipeline_rank,
- group=self.group,
)
- encoder_hidden_states = mx.distributed.recv(
+ encoder_hidden_states = self._recv(
(batch_size, text_seq_len, hidden_dim),
dtype,
self.prev_pipeline_rank,
- group=self.group,
)
- mx.eval(hidden_states, encoder_hidden_states)
assert self.joint_block_wrappers is not None
assert encoder_hidden_states is not None
@@ -825,9 +881,7 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- concatenated = mx.distributed.send(
- concatenated, self.next_pipeline_rank, group=self.group
- )
+ concatenated = self._send(concatenated, self.next_pipeline_rank)
mx.async_eval(concatenated)
elif self.has_joint_blocks and not self.is_last_stage:
@@ -838,11 +892,9 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- hidden_states = mx.distributed.send(
- hidden_states, self.next_pipeline_rank, group=self.group
- )
- encoder_hidden_states = mx.distributed.send(
- encoder_hidden_states, self.next_pipeline_rank, group=self.group
+ hidden_states = self._send(hidden_states, self.next_pipeline_rank)
+ encoder_hidden_states = self._send(
+ encoder_hidden_states, self.next_pipeline_rank
)
mx.async_eval(hidden_states, encoder_hidden_states)
@@ -854,13 +906,11 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- hidden_states = mx.distributed.recv(
+ hidden_states = self._recv(
(batch_size, text_seq_len + num_img_tokens, hidden_dim),
dtype,
self.prev_pipeline_rank,
- group=self.group,
)
- mx.eval(hidden_states)
assert self.single_block_wrappers is not None
with trace(
@@ -886,9 +936,7 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- hidden_states = mx.distributed.send(
- hidden_states, self.next_pipeline_rank, group=self.group
- )
+ hidden_states = self._send(hidden_states, self.next_pipeline_rank)
mx.async_eval(hidden_states)
hidden_states = hidden_states[:, text_seq_len:, ...]
@@ -961,16 +1009,11 @@ class DiffusionRunner:
)
if not self.is_first_stage:
- hidden_states = mx.distributed.send(
- hidden_states, self.first_pipeline_rank, group=self.group
- )
+ hidden_states = self._send(hidden_states, self.first_pipeline_rank)
mx.async_eval(hidden_states)
elif self.is_first_stage:
- hidden_states = mx.distributed.recv_like(
- prev_latents, src=self.last_pipeline_rank, group=self.group
- )
- mx.eval(hidden_states)
+ hidden_states = self._recv_like(prev_latents, src=self.last_pipeline_rank)
else:
hidden_states = prev_latents
@@ -1006,10 +1049,7 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- patch = mx.distributed.recv_like(
- patch, src=self.last_pipeline_rank, group=self.group
- )
- mx.eval(patch)
+ patch = self._recv_like(patch, src=self.last_pipeline_rank)
results: list[tuple[bool, mx.array]] = []
@@ -1066,10 +1106,9 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- patch_latents[patch_idx] = mx.distributed.send(
+ patch_latents[patch_idx] = self._send(
patch_latents[patch_idx],
self.first_pipeline_rank,
- group=self.group,
)
mx.async_eval(patch_latents[patch_idx])
@@ -1116,13 +1155,11 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- patch = mx.distributed.recv(
+ patch = self._recv(
(batch_size, patch_len, hidden_dim),
patch.dtype,
self.prev_pipeline_rank,
- group=self.group,
)
- mx.eval(patch)
if patch_idx == 0:
with trace(
@@ -1130,13 +1167,11 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- encoder_hidden_states = mx.distributed.recv(
+ encoder_hidden_states = self._recv(
(batch_size, text_seq_len, hidden_dim),
patch.dtype,
self.prev_pipeline_rank,
- group=self.group,
)
- mx.eval(encoder_hidden_states)
if self.is_first_stage:
patch, encoder_hidden_states = self.adapter.compute_embeddings(
@@ -1175,9 +1210,7 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- patch_concat = mx.distributed.send(
- patch_concat, self.next_pipeline_rank, group=self.group
- )
+ patch_concat = self._send(patch_concat, self.next_pipeline_rank)
mx.async_eval(patch_concat)
elif self.has_joint_blocks and not self.is_last_stage:
@@ -1187,9 +1220,7 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- patch = mx.distributed.send(
- patch, self.next_pipeline_rank, group=self.group
- )
+ patch = self._send(patch, self.next_pipeline_rank)
mx.async_eval(patch)
if patch_idx == 0:
@@ -1199,8 +1230,8 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- encoder_hidden_states = mx.distributed.send(
- encoder_hidden_states, self.next_pipeline_rank, group=self.group
+ encoder_hidden_states = self._send(
+ encoder_hidden_states, self.next_pipeline_rank
)
mx.async_eval(encoder_hidden_states)
@@ -1213,13 +1244,11 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- patch = mx.distributed.recv(
+ patch = self._recv(
(batch_size, text_seq_len + patch_len, hidden_dim),
patch.dtype,
self.prev_pipeline_rank,
- group=self.group,
)
- mx.eval(patch)
assert self.single_block_wrappers is not None
with trace(
@@ -1245,9 +1274,7 @@ class DiffusionRunner:
rank=self.rank,
category="comms",
):
- patch = mx.distributed.send(
- patch, self.next_pipeline_rank, group=self.group
- )
+ patch = self._send(patch, self.next_pipeline_rank)
mx.async_eval(patch)
noise: mx.array | None = None
diff --git a/src/exo/worker/runner/image_models/runner.py b/src/exo/worker/runner/image_models/runner.py
index 79eaf04f..4fdecec2 100644
--- a/src/exo/worker/runner/image_models/runner.py
+++ b/src/exo/worker/runner/image_models/runner.py
@@ -4,7 +4,11 @@ from typing import TYPE_CHECKING, Literal
import mlx.core as mx
-from exo.api.types import ImageGenerationStats
+from exo.api.types import (
+ ImageEditsTaskParams,
+ ImageGenerationStats,
+ ImageGenerationTaskParams,
+)
from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED
from exo.shared.models.model_cards import ModelTask
from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
@@ -235,6 +239,77 @@ class Runner:
def acknowledge_task(self, task: Task):
self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
+ def _check_cancelled(self, task_id: TaskId) -> bool:
+ for cancel_id in self.cancel_receiver.collect():
+ self.cancelled_tasks.add(cancel_id)
+ return (
+ task_id in self.cancelled_tasks or CANCEL_ALL_TASKS in self.cancelled_tasks
+ )
+
+ def _run_image_task(
+ self,
+ task: Task,
+ task_params: ImageGenerationTaskParams | ImageEditsTaskParams,
+ command_id: CommandId,
+ ) -> None:
+ assert self.image_model
+ logger.info(f"received image task: {str(task)[:500]}")
+ logger.info("runner running")
+ self.update_status(RunnerRunning())
+ self.acknowledge_task(task)
+
+ def cancel_checker() -> bool:
+ return self._check_cancelled(task.task_id)
+
+ try:
+ image_index = 0
+ for response in generate_image(
+ model=self.image_model,
+ task=task_params,
+ cancel_checker=cancel_checker,
+ ):
+ if _is_primary_output_node(self.shard_metadata):
+ match response:
+ case PartialImageResponse():
+ logger.info(
+ f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
+ )
+ _process_image_response(
+ response,
+ command_id,
+ self.shard_metadata,
+ self.event_sender,
+ image_index,
+ )
+ case ImageGenerationResponse():
+ logger.info("sending final ImageChunk")
+ _process_image_response(
+ response,
+ command_id,
+ self.shard_metadata,
+ self.event_sender,
+ image_index,
+ )
+ image_index += 1
+ except Exception as e:
+ if _is_primary_output_node(self.shard_metadata):
+ self.event_sender.send(
+ ChunkGenerated(
+ command_id=command_id,
+ chunk=ErrorChunk(
+ model=self.shard_metadata.model_card.model_id,
+ finish_reason="error",
+ error_message=str(e),
+ ),
+ )
+ )
+ raise
+ finally:
+ _send_traces_if_enabled(self.event_sender, task.task_id, self.device_rank)
+
+ self.current_status = RunnerReady()
+ logger.info("runner ready")
+
def main(self):
with self.task_receiver as tasks:
for task in tasks:
@@ -306,124 +381,11 @@ class Runner:
self.current_status = RunnerReady()
logger.info("runner ready")
- case ImageGeneration(task_params=task_params, command_id=command_id) if (
- isinstance(self.current_status, RunnerReady)
- ):
- assert self.image_model
- logger.info(f"received image generation request: {str(task)[:500]}")
- logger.info("runner running")
- self.update_status(RunnerRunning())
- self.acknowledge_task(task)
-
- try:
- image_index = 0
- for response in generate_image(
- model=self.image_model, task=task_params
- ):
- is_primary_output = _is_primary_output_node(self.shard_metadata)
-
- if is_primary_output:
- match response:
- case PartialImageResponse():
- logger.info(
- f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
- )
- _process_image_response(
- response,
- command_id,
- self.shard_metadata,
- self.event_sender,
- image_index,
- )
- case ImageGenerationResponse():
- logger.info("sending final ImageChunk")
- _process_image_response(
- response,
- command_id,
- self.shard_metadata,
- self.event_sender,
- image_index,
- )
- image_index += 1
- # can we make this more explicit?
- except Exception as e:
- if _is_primary_output_node(self.shard_metadata):
- self.event_sender.send(
- ChunkGenerated(
- command_id=command_id,
- chunk=ErrorChunk(
- model=self.shard_metadata.model_card.model_id,
- finish_reason="error",
- error_message=str(e),
- ),
- )
- )
- raise
- finally:
- _send_traces_if_enabled(
- self.event_sender, task.task_id, self.device_rank
- )
-
- self.current_status = RunnerReady()
- logger.info("runner ready")
-
- case ImageEdits(task_params=task_params, command_id=command_id) if (
- isinstance(self.current_status, RunnerReady)
- ):
- assert self.image_model
- logger.info(f"received image edits request: {str(task)[:500]}")
- logger.info("runner running")
- self.update_status(RunnerRunning())
- self.acknowledge_task(task)
-
- try:
- image_index = 0
- for response in generate_image(
- model=self.image_model, task=task_params
- ):
- if _is_primary_output_node(self.shard_metadata):
- match response:
- case PartialImageResponse():
- logger.info(
- f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
- )
- _process_image_response(
- response,
- command_id,
- self.shard_metadata,
- self.event_sender,
- image_index,
- )
- case ImageGenerationResponse():
- logger.info("sending final ImageChunk")
- _process_image_response(
- response,
- command_id,
- self.shard_metadata,
- self.event_sender,
- image_index,
- )
- image_index += 1
- except Exception as e:
- if _is_primary_output_node(self.shard_metadata):
- self.event_sender.send(
- ChunkGenerated(
- command_id=command_id,
- chunk=ErrorChunk(
- model=self.shard_metadata.model_card.model_id,
- finish_reason="error",
- error_message=str(e),
- ),
- )
- )
- raise
- finally:
- _send_traces_if_enabled(
- self.event_sender, task.task_id, self.device_rank
- )
-
- self.current_status = RunnerReady()
- logger.info("runner ready")
+ case (
+ ImageGeneration(task_params=task_params, command_id=command_id)
+ | ImageEdits(task_params=task_params, command_id=command_id)
+ ) if isinstance(self.current_status, RunnerReady):
+ self._run_image_task(task, task_params, command_id)
case Shutdown():
logger.info("runner shutting down")
diff --git a/src/exo/worker/tests/unittests/test_image/__init__.py b/src/exo/worker/tests/unittests/test_image/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/exo/worker/tests/unittests/test_image/test_cancellation.py b/src/exo/worker/tests/unittests/test_image/test_cancellation.py
new file mode 100644
index 00000000..61d8d0a8
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_image/test_cancellation.py
@@ -0,0 +1,433 @@
+# pyright: reportPrivateUsage=false
+"""Tests for image generation cancellation logic.
+
+Tests the NaN sentinel protocol, cancellation checking, and the
+image runner's cancel_checker integration.
+"""
+
+from collections.abc import Callable
+from unittest.mock import MagicMock
+
+import mlx.core as mx
+
+from exo.shared.types.tasks import CANCEL_ALL_TASKS, TaskId
+from exo.worker.engines.image.pipeline.runner import DiffusionRunner
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_runner() -> DiffusionRunner:
+ """Create a DiffusionRunner with minimal config for unit testing.
+
+ Uses a mock adapter and no distributed group (single-node).
+ """
+ mock_config = MagicMock()
+ mock_config.joint_block_count = 10
+ mock_config.single_block_count = 10
+ mock_config.total_blocks = 20
+ mock_config.guidance_scale = None
+
+ mock_adapter = MagicMock()
+
+ mock_shard = MagicMock()
+ mock_shard.device_rank = 0
+ mock_shard.world_size = 1
+ mock_shard.start_layer = 0
+ mock_shard.end_layer = 20
+
+ runner = DiffusionRunner(
+ config=mock_config,
+ adapter=mock_adapter,
+ group=None,
+ shard_metadata=mock_shard,
+ )
+ return runner
+
+
+class FakeCancelReceiver:
+ """Fake MpReceiver that returns pre-loaded items from collect()."""
+
+ def __init__(self, items: list[TaskId] | None = None):
+ self._items = list(items) if items else []
+
+ def collect(self) -> list[TaskId]:
+ result = self._items
+ self._items = []
+ return result
+
+
+class FakeImageRunner:
+ """Fake image runner for testing _check_cancelled logic."""
+
+ def __init__(self, cancel_items: list[TaskId] | None = None) -> None:
+ self.cancel_receiver = FakeCancelReceiver(cancel_items)
+ self.cancelled_tasks = set[TaskId]()
+
+ def _check_cancelled(self, task_id: TaskId) -> bool:
+ for cancel_id in self.cancel_receiver.collect():
+ self.cancelled_tasks.add(cancel_id)
+ return (
+ task_id in self.cancelled_tasks or CANCEL_ALL_TASKS in self.cancelled_tasks
+ )
+
+
+# ---------------------------------------------------------------------------
+# _is_sentinel
+# ---------------------------------------------------------------------------
+
+
+class TestIsSentinel:
+ def test_all_nan_is_sentinel(self) -> None:
+ runner = _make_runner()
+ tensor = mx.full((2, 3), float("nan"))
+ mx.eval(tensor)
+ assert runner._is_sentinel(tensor) is True
+
+ def test_all_zeros_is_not_sentinel(self) -> None:
+ runner = _make_runner()
+ tensor = mx.zeros((2, 3))
+ mx.eval(tensor)
+ assert runner._is_sentinel(tensor) is False
+
+ def test_mixed_nan_and_real_is_not_sentinel(self) -> None:
+ """A tensor with some NaN and some real values must NOT be a sentinel.
+ Using mx.any(isnan) would incorrectly flag this as a sentinel.
+ """
+ runner = _make_runner()
+ tensor = mx.array([float("nan"), 1.0, 2.0])
+ mx.eval(tensor)
+ assert runner._is_sentinel(tensor) is False
+
+ def test_single_element_nan(self) -> None:
+ runner = _make_runner()
+ tensor = mx.array([float("nan")])
+ mx.eval(tensor)
+ assert runner._is_sentinel(tensor) is True
+
+ def test_large_tensor_all_nan(self) -> None:
+ runner = _make_runner()
+ tensor = mx.full((64, 128, 32), float("nan"))
+ mx.eval(tensor)
+ assert runner._is_sentinel(tensor) is True
+
+ def test_real_data_not_sentinel(self) -> None:
+ runner = _make_runner()
+ tensor = mx.random.normal((4, 8))
+ mx.eval(tensor)
+ assert runner._is_sentinel(tensor) is False
+
+
+# ---------------------------------------------------------------------------
+# _check_cancellation
+# ---------------------------------------------------------------------------
+
+
+class TestCheckCancellation:
+ def test_first_stage_polls_checker(self) -> None:
+ runner = _make_runner()
+ assert runner.is_first_stage # single-node is always first stage
+
+ checker: Callable[[], bool] = MagicMock(return_value=True)
+ runner._cancel_checker = checker
+
+ runner._check_cancellation()
+
+ checker.assert_called_once()
+ assert runner._cancelling is True
+
+ def test_checker_returning_false_does_not_cancel(self) -> None:
+ runner = _make_runner()
+ checker: Callable[[], bool] = MagicMock(return_value=False)
+ runner._cancel_checker = checker
+
+ runner._check_cancellation()
+
+ assert runner._cancelling is False
+
+ def test_no_checker_does_not_cancel(self) -> None:
+ runner = _make_runner()
+ runner._cancel_checker = None
+
+ runner._check_cancellation()
+
+ assert runner._cancelling is False
+
+ def test_already_cancelling_skips_checker(self) -> None:
+ runner = _make_runner()
+ runner._cancelling = True
+ checker: Callable[[], bool] = MagicMock(return_value=False)
+ runner._cancel_checker = checker
+
+ runner._check_cancellation()
+
+ checker.assert_not_called()
+ assert runner._cancelling is True # stays True
+
+ def test_cancelling_flag_is_false_on_init(self) -> None:
+ """_cancelling defaults to False on a fresh runner."""
+ runner = _make_runner()
+ assert runner._cancelling is False
+
+
+# ---------------------------------------------------------------------------
+# _send wrapper
+# ---------------------------------------------------------------------------
+
+
+class TestSendWrapper:
+ def test_send_replaces_data_with_nan_when_cancelling(self) -> None:
+ """When _cancelling is True, _send should replace data with NaN."""
+ runner = _make_runner()
+ runner._cancelling = True
+ # _send asserts group is not None, so we need a mock group
+ runner.group = MagicMock()
+
+ data = mx.ones((2, 3))
+ mx.eval(data)
+
+ # Mock mx.distributed.send to capture what's sent
+ original_send = mx.distributed.send
+ sent_data: list[mx.array] = []
+
+ def mock_send(d: mx.array, dst: int, group: mx.distributed.Group) -> mx.array:
+ mx.eval(d)
+ sent_data.append(d)
+ return d
+
+ mx.distributed.send = mock_send
+ try:
+ runner._send(data, dst=1)
+ assert len(sent_data) == 1
+ mx.eval(sent_data[0])
+ assert mx.all(mx.isnan(sent_data[0])).item()
+ assert sent_data[0].shape == (2, 3)
+ finally:
+ mx.distributed.send = original_send
+
+ def test_send_passes_real_data_when_not_cancelling(self) -> None:
+ runner = _make_runner()
+ runner._cancelling = False
+ runner.group = MagicMock()
+
+ data = mx.ones((2, 3))
+ mx.eval(data)
+
+ sent_data: list[mx.array] = []
+
+ def mock_send(d: mx.array, dst: int, group: mx.distributed.Group) -> mx.array:
+ mx.eval(d)
+ sent_data.append(d)
+ return d
+
+ original_send = mx.distributed.send
+ mx.distributed.send = mock_send
+ try:
+ runner._send(data, dst=1)
+ assert len(sent_data) == 1
+ mx.eval(sent_data[0])
+ assert not mx.any(mx.isnan(sent_data[0])).item()
+ finally:
+ mx.distributed.send = original_send
+
+
+# ---------------------------------------------------------------------------
+# Image runner _check_cancelled
+# ---------------------------------------------------------------------------
+
+
+class TestImageRunnerCheckCancelled:
+ """Tests for the image runner's _check_cancelled method."""
+
+ def test_no_cancellation(self) -> None:
+ runner = FakeImageRunner()
+ assert runner._check_cancelled(TaskId("task-1")) is False
+
+ def test_specific_task_cancelled(self) -> None:
+ task_id = TaskId("task-1")
+ runner = FakeImageRunner([task_id])
+ assert runner._check_cancelled(task_id) is True
+
+ def test_different_task_not_cancelled(self) -> None:
+ runner = FakeImageRunner([TaskId("task-2")])
+ assert runner._check_cancelled(TaskId("task-1")) is False
+
+ def test_cancel_all_tasks(self) -> None:
+ runner = FakeImageRunner([CANCEL_ALL_TASKS])
+ assert runner._check_cancelled(TaskId("any-task")) is True
+
+ def test_collect_accumulates(self) -> None:
+ """Multiple collect() calls accumulate cancelled task IDs."""
+ runner = FakeImageRunner([TaskId("task-1")])
+ runner._check_cancelled(TaskId("task-1"))
+
+ # First collect drained the receiver, but task-1 is in cancelled_tasks
+ assert runner._check_cancelled(TaskId("task-1")) is True
+
+ def test_collect_empty_after_drain(self) -> None:
+ """After draining, collect returns empty and previous cancellations persist."""
+ runner = FakeImageRunner([TaskId("task-1")])
+
+ # First call drains
+ runner._check_cancelled(TaskId("other"))
+ # task-1 is now in cancelled_tasks but "other" was never cancelled
+ assert runner._check_cancelled(TaskId("other")) is False
+ assert runner._check_cancelled(TaskId("task-1")) is True
+
+
+# ---------------------------------------------------------------------------
+# Drain condition logic
+# ---------------------------------------------------------------------------
+
+
+class TestDrainCondition:
+ """Verify the drain condition evaluates correctly for various scenarios."""
+
+ def _should_drain(
+ self,
+ *,
+ cancelling: bool,
+ is_first_stage: bool,
+ is_last_stage: bool,
+ is_distributed: bool,
+ t: int,
+ init_time_step: int,
+ num_sync_steps: int,
+ num_inference_steps: int,
+ ) -> bool:
+ """Replicate the drain condition from _run_diffusion_loop."""
+ return (
+ cancelling
+ and is_first_stage
+ and not is_last_stage
+ and is_distributed
+ and t >= init_time_step + num_sync_steps
+ and t != num_inference_steps - 1
+ )
+
+ def test_no_drain_during_sync_step(self) -> None:
+ """Sync steps have no cross-timestep ring state."""
+ assert not self._should_drain(
+ cancelling=True,
+ is_first_stage=True,
+ is_last_stage=False,
+ is_distributed=True,
+ t=0, # sync step
+ init_time_step=0,
+ num_sync_steps=2,
+ num_inference_steps=10,
+ )
+
+ def test_drain_during_async_step(self) -> None:
+ assert self._should_drain(
+ cancelling=True,
+ is_first_stage=True,
+ is_last_stage=False,
+ is_distributed=True,
+ t=3, # async step
+ init_time_step=0,
+ num_sync_steps=2,
+ num_inference_steps=10,
+ )
+
+ def test_no_drain_on_last_step(self) -> None:
+ """Last step doesn't send, so nothing to drain."""
+ assert not self._should_drain(
+ cancelling=True,
+ is_first_stage=True,
+ is_last_stage=False,
+ is_distributed=True,
+ t=9, # last step
+ init_time_step=0,
+ num_sync_steps=2,
+ num_inference_steps=10,
+ )
+
+ def test_no_drain_when_not_cancelling(self) -> None:
+ assert not self._should_drain(
+ cancelling=False,
+ is_first_stage=True,
+ is_last_stage=False,
+ is_distributed=True,
+ t=5,
+ init_time_step=0,
+ num_sync_steps=2,
+ num_inference_steps=10,
+ )
+
+ def test_no_drain_on_last_stage(self) -> None:
+ """Last stage is also first stage (single pipeline) — no ring."""
+ assert not self._should_drain(
+ cancelling=True,
+ is_first_stage=True,
+ is_last_stage=True,
+ is_distributed=True,
+ t=5,
+ init_time_step=0,
+ num_sync_steps=2,
+ num_inference_steps=10,
+ )
+
+ def test_no_drain_single_node(self) -> None:
+ assert not self._should_drain(
+ cancelling=True,
+ is_first_stage=True,
+ is_last_stage=False,
+ is_distributed=False,
+ t=5,
+ init_time_step=0,
+ num_sync_steps=2,
+ num_inference_steps=10,
+ )
+
+ def test_no_drain_not_first_stage(self) -> None:
+ """Only first stage needs to drain (it's the one receiving)."""
+ assert not self._should_drain(
+ cancelling=True,
+ is_first_stage=False,
+ is_last_stage=False,
+ is_distributed=True,
+ t=5,
+ init_time_step=0,
+ num_sync_steps=2,
+ num_inference_steps=10,
+ )
+
+ def test_drain_first_async_step(self) -> None:
+ """First async step: last stage sends, so drain is needed."""
+ assert self._should_drain(
+ cancelling=True,
+ is_first_stage=True,
+ is_last_stage=False,
+ is_distributed=True,
+ t=2, # first async step (init=0, sync=2)
+ init_time_step=0,
+ num_sync_steps=2,
+ num_inference_steps=10,
+ )
+
+ def test_drain_with_nonzero_init_time_step(self) -> None:
+ """img2img can have init_time_step > 0."""
+ assert self._should_drain(
+ cancelling=True,
+ is_first_stage=True,
+ is_last_stage=False,
+ is_distributed=True,
+ t=5,
+ init_time_step=3,
+ num_sync_steps=1,
+ num_inference_steps=10,
+ )
+
+ def test_no_drain_sync_with_nonzero_init(self) -> None:
+ assert not self._should_drain(
+ cancelling=True,
+ is_first_stage=True,
+ is_last_stage=False,
+ is_distributed=True,
+ t=3,
+ init_time_step=3,
+ num_sync_steps=1,
+ num_inference_steps=10,
+ )
← fc1ae901 fix: DeepSeek V3.2 warmup crash and tool calling + add catal
·
back to Exo
·
Prefer higher % downloaded nodes for API placement previews 30fd5aa1 →