electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Support image generation cancellation (#1774)

Browse Source 
## 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
This commit is contained in:
ciaranbor
2026-03-25 16:56:04 +00:00
committed by GitHub
parent fc1ae90111
commit 6de14cfedb
7 changed files with 660 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dashboard/src/lib/stores/app.svelte.ts
+48 -14
View File
@@ -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();
}
src/exo/worker/engines/image/distributed_model.py
+3 -1
View File
@@ -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)
src/exo/worker/engines/image/generate.py
+3
View File
@@ -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)
src/exo/worker/engines/image/pipeline/runner.py
+92 -65
View File
@@ -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
src/exo/worker/runner/image_models/runner.py
+81 -119
View File
@@ -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")
src/exo/worker/tests/unittests/test_image/__init__.py
View File
src/exo/worker/tests/unittests/test_image/test_cancellation.py
+433
View File
@@ -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,
)