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, + )