api cancellation

closing the http request to the api now - sends a cancellation from the api - writes that canellation in the master - worker plans off the cancellation - runner observes that cancellation after every generation step (+1 communication per token) - cancellation happens synchronously to prevent gpu locks
2026-01-31 01:01:11 -05:00 · 2026-01-30 13:56:49 +00:00
25 changed files with 337 additions and 495 deletions
--- a/MISSED_THINGS.md
+++ b/MISSED_THINGS.md
@@ -5,18 +5,18 @@
 [X] Fetching download status of all models on start
 [X] Deduplication of tasks in plan_step.
 [X] resolve_allow_patterns should just be wildcard now.
-[] no mx_barrier in genreate.py mlx_generate at the end.
+[X] no mx_barrier in genreate.py mlx_generate at the end.
 [] cache assertion not needed in auto_parallel.py PipelineLastLayer.
-[] GPTOSS support dropped in auto_parallel.py.
-[] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
-[] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
-[] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
+[X] GPTOSS support dropped in auto_parallel.py.
+[X] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
+[X] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
+[X] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
 [] Dropped prefill/decode code in auto_parallel.py and utils_mlx.py.
 [X] KV_CACHE_BITS should be None to disable quantized KV cache.
-[] Dropped _set_nofile_limit in utils_mlx.py.
-[] We have group optional in load_mlx_items in utils_mlx.py.
+[X] Dropped _set_nofile_limit in utils_mlx.py.
+[X] We have group optional in load_mlx_items in utils_mlx.py.
 [] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
-[] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
+[X] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
 [X] We put cache limit back in utils_mlx.py.
 [] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
 [] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
--- a/src/exo/main.py
+++ b/src/exo/main.py
@@ -40,7 +40,6 @@ class Node:

    node_id: NodeId
    event_index_counter: Iterator[int]
-    _profiling_in_progress: set[NodeId] = field(default_factory=set)
    _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)

    @classmethod
@@ -87,7 +86,6 @@ class Node:
        else:
            api = None

-        profiling_in_progress: set[NodeId] = set()
        if not args.no_worker:
            worker = Worker(
                node_id,
@@ -98,7 +96,6 @@ class Node:
                command_sender=router.sender(topics.COMMANDS),
                download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                event_index_counter=event_index_counter,
-                profiling_in_progress=profiling_in_progress,
            )
        else:
            worker = None
@@ -136,7 +133,6 @@ class Node:
            api,
            node_id,
            event_index_counter,
-            profiling_in_progress,
        )

    async def run(self):
@@ -243,7 +239,6 @@ class Node:
                                topics.DOWNLOAD_COMMANDS
                            ),
                            event_index_counter=self.event_index_counter,
-                            profiling_in_progress=self._profiling_in_progress,
                        )
                        self._tg.start_soon(self.worker.run)
                    if self.api:
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -8,8 +8,8 @@ from typing import Annotated, Literal, cast
 from uuid import uuid4

 import anyio
-from anyio import BrokenResourceError, ClosedResourceError, EndOfStream, create_task_group
-from anyio.abc import SocketStream, TaskGroup
+from anyio import BrokenResourceError, create_task_group
+from anyio.abc import TaskGroup
 from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -90,6 +90,7 @@ from exo.shared.types.commands import (
    PlaceInstance,
    SendInputChunk,
    StartDownload,
+    TaskCancelled,
    TaskFinished,
 )
 from exo.shared.types.common import CommandId, Id, NodeId, SessionId
@@ -501,16 +502,14 @@ class API:
                        break

        except anyio.get_cancelled_exc_class():
-            # TODO: TaskCancelled
-            """
-            self.command_sender.send_nowait(
-                ForwarderCommand(origin=self.node_id, command=command)
-            )
-            """
+            command = TaskCancelled(cancelled_command_id=command_id)
+            with anyio.CancelScope(shield=True):
+                await self.command_sender.send(
+                    ForwarderCommand(origin=self.node_id, command=command)
+                )
            raise
        finally:
-            command = TaskFinished(finished_command_id=command_id)
-            await self._send(command)
+            await self._send(TaskFinished(finished_command_id=command_id))
            if command_id in self._chat_completion_queues:
                del self._chat_completion_queues[command_id]

@@ -548,7 +547,7 @@ class API:

    async def _collect_chat_completion(
        self, command_id: CommandId
-    ) -> ChatCompletionResponse:
+    ) -> AsyncGenerator[ChatCompletionResponse]:
        """Collect all token chunks for a chat completion and return a single response."""

        text_parts: list[str] = []
@@ -589,7 +588,7 @@ class API:
        combined_text = "".join(text_parts)
        assert model is not None

-        return ChatCompletionResponse(
+        yield ChatCompletionResponse(
            id=command_id,
            created=int(time.time()),
            model=model,
@@ -606,6 +605,7 @@ class API:
            ],
            usage=usage,
        )
+        return

    async def _collect_chat_completion_with_stats(
        self, command_id: CommandId
@@ -696,7 +696,10 @@ class API:
                media_type="text/event-stream",
            )

-        return await self._collect_chat_completion(command.command_id)
+        return StreamingResponse(
+            self._collect_chat_completion(command.command_id),
+            media_type="application/json",
+        )

    async def bench_chat_completions(
        self, payload: BenchChatCompletionTaskParams
@@ -902,6 +905,11 @@ class API:
                        del image_metadata[key]

        except anyio.get_cancelled_exc_class():
+            command = TaskCancelled(cancelled_command_id=command_id)
+            with anyio.CancelScope(shield=True):
+                await self.command_sender.send(
+                    ForwarderCommand(origin=self.node_id, command=command)
+                )
            raise
        finally:
            await self._send(TaskFinished(finished_command_id=command_id))
@@ -983,6 +991,11 @@ class API:

            return (images, stats if capture_stats else None)
        except anyio.get_cancelled_exc_class():
+            command = TaskCancelled(cancelled_command_id=command_id)
+            with anyio.CancelScope(shield=True):
+                await self.command_sender.send(
+                    ForwarderCommand(origin=self.node_id, command=command)
+                )
            raise
        finally:
            await self._send(TaskFinished(finished_command_id=command_id))
@@ -1250,7 +1263,6 @@ class API:
            tg.start_soon(self._apply_state)
            tg.start_soon(self._pause_on_new_election)
            tg.start_soon(self._cleanup_expired_images)
-            tg.start_soon(self._run_bandwidth_server)
            print_startup_banner(self.port)
            await serve(
                cast(ASGIFramework, self.app),
@@ -1258,43 +1270,6 @@ class API:
                shutdown_trigger=lambda: anyio.sleep_forever(),
            )

-    async def _run_bandwidth_server(self):
-        """TCP server for iperf-like bandwidth testing."""
-        bandwidth_port = self.port + 1
-        listener = await anyio.create_tcp_listener(local_port=bandwidth_port)
-        logger.info(f"Bandwidth test server listening on port {bandwidth_port}")
-        await listener.serve(self._handle_bandwidth_connection)
-
-    async def _handle_bandwidth_connection(self, stream: SocketStream) -> None:
-        """Handle a single bandwidth test connection."""
-        try:
-            mode = await stream.receive(1)
-            if mode == b"U":
-                # Upload test: client sends, we receive
-                bytes_received = 0
-                start = time.perf_counter()
-                while True:
-                    try:
-                        data = await stream.receive(1024 * 1024)
-                        if not data or data == b"DONE" or b"DONE" in data:
-                            break
-                        bytes_received += len(data)
-                    except EndOfStream:
-                        break
-                elapsed = time.perf_counter() - start
-                logger.debug(f"Bandwidth upload: {bytes_received} bytes in {elapsed:.3f}s")
-            elif mode == b"D":
-                # Download test: we send, client receives
-                chunk = b"X" * (1024 * 1024)
-                start = time.perf_counter()
-                while time.perf_counter() - start < 1.0:  # Send for 1s
-                    try:
-                        await stream.send(chunk)
-                    except (BrokenResourceError, ClosedResourceError):
-                        break
-        except Exception as e:
-            logger.debug(f"Bandwidth connection error: {e}")
-
        self.command_sender.close()
        self.global_event_receiver.close()

--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -21,6 +21,7 @@ from exo.shared.types.commands import (
    PlaceInstance,
    RequestEventLog,
    SendInputChunk,
+    TaskCancelled,
    TaskFinished,
    TestCommand,
 )
@@ -35,6 +36,7 @@ from exo.shared.types.events import (
    NodeTimedOut,
    TaskCreated,
    TaskDeleted,
+    TaskStatusUpdated,
 )
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
@@ -246,7 +248,7 @@ class Master:
                        case DeleteInstance():
                            placement = delete_instance(command, self.state.instances)
                            transition_events = get_transition_events(
-                                self.state.instances, placement
+                                self.state.instances, placement, self.state.tasks
                            )
                            generated_events.extend(transition_events)
                        case PlaceInstance():
@@ -258,7 +260,7 @@ class Master:
                                self.state.node_network,
                            )
                            transition_events = get_transition_events(
-                                self.state.instances, placement
+                                self.state.instances, placement, self.state.tasks
                            )
                            generated_events.extend(transition_events)
                        case CreateInstance():
@@ -268,7 +270,7 @@ class Master:
                                self.state.instances,
                            )
                            transition_events = get_transition_events(
-                                self.state.instances, placement
+                                self.state.instances, placement, self.state.tasks
                            )
                            generated_events.extend(transition_events)
                        case SendInputChunk(chunk=chunk):
@@ -278,6 +280,18 @@ class Master:
                                    chunk=chunk,
                                )
                            )
+                        case TaskCancelled():
+                            if (
+                                task_id := self.command_task_mapping.get(
+                                    command.cancelled_command_id
+                                )
+                            ) is not None:
+                                generated_events.append(
+                                    TaskStatusUpdated(
+                                        task_status=TaskStatus.Cancelled,
+                                        task_id=task_id,
+                                    )
+                                )
                        case TaskFinished():
                            generated_events.append(
                                TaskDeleted(
@@ -286,10 +300,9 @@ class Master:
                                    ]
                                )
                            )
-                            if command.finished_command_id in self.command_task_mapping:
-                                del self.command_task_mapping[
-                                    command.finished_command_id
-                                ]
+                            self.command_task_mapping.pop(
+                                command.finished_command_id, None
+                            )
                        case RequestEventLog():
                            # We should just be able to send everything, since other buffers will ignore old messages
                            for i in range(command.since_idx, len(self._event_log)):
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -20,9 +20,15 @@ from exo.shared.types.commands import (
    PlaceInstance,
 )
 from exo.shared.types.common import NodeId
-from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
+from exo.shared.types.events import (
+    Event,
+    InstanceCreated,
+    InstanceDeleted,
+    TaskStatusUpdated,
+)
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
+from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.instances import (
    Instance,
    InstanceId,
@@ -180,6 +186,7 @@ def delete_instance(
 def get_transition_events(
    current_instances: Mapping[InstanceId, Instance],
    target_instances: Mapping[InstanceId, Instance],
+    tasks: Mapping[TaskId, Task],
 ) -> Sequence[Event]:
    events: list[Event] = []

@@ -195,6 +202,18 @@ def get_transition_events(
    # find instances to delete
    for instance_id in current_instances:
        if instance_id not in target_instances:
+            for task in tasks.values():
+                if task.instance_id == instance_id and task.task_status in [
+                    TaskStatus.Pending,
+                    TaskStatus.Running,
+                ]:
+                    events.append(
+                        TaskStatusUpdated(
+                            task_status=TaskStatus.Cancelled,
+                            task_id=task.task_id,
+                        )
+                    )
+
            events.append(
                InstanceDeleted(
                    instance_id=instance_id,
--- a/src/exo/master/tests/conftest.py
+++ b/src/exo/master/tests/conftest.py
@@ -28,9 +28,6 @@ def create_node_network() -> NodeNetworkInfo:
 def create_socket_connection(ip: int, sink_port: int = 1234) -> SocketConnection:
    return SocketConnection(
        sink_multiaddr=Multiaddr(address=f"/ip4/169.254.0.{ip}/tcp/{sink_port}"),
-        latency_ms=1.0,
-        other_to_sink_bandwidth_mbps=1000.0,
-        sink_to_other_bandwidth_mbps=1000.0,
    )


--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -239,7 +239,7 @@ def test_get_transition_events_no_change(instance: Instance):
    target_instances = {instance_id: instance}

    # act
-    events = get_transition_events(current_instances, target_instances)
+    events = get_transition_events(current_instances, target_instances, {})

    # assert
    assert len(events) == 0
@@ -252,7 +252,7 @@ def test_get_transition_events_create_instance(instance: Instance):
    target_instances: dict[InstanceId, Instance] = {instance_id: instance}

    # act
-    events = get_transition_events(current_instances, target_instances)
+    events = get_transition_events(current_instances, target_instances, {})

    # assert
    assert len(events) == 1
@@ -266,7 +266,7 @@ def test_get_transition_events_delete_instance(instance: Instance):
    target_instances: dict[InstanceId, Instance] = {}

    # act
-    events = get_transition_events(current_instances, target_instances)
+    events = get_transition_events(current_instances, target_instances, {})

    # assert
    assert len(events) == 1
@@ -366,10 +366,7 @@ def test_tensor_rdma_backend_connectivity_matrix(
        ip_address="10.0.0.1",
    )
    ethernet_conn = SocketConnection(
-        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000"),
-        latency_ms=1.0,
-        other_to_sink_bandwidth_mbps=1000.0,
-        sink_to_other_bandwidth_mbps=1000.0,
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
    )

    node_network = {
--- a/src/exo/master/tests/test_topology.py
+++ b/src/exo/master/tests/test_topology.py
@@ -15,9 +15,6 @@ def topology() -> Topology:
 def socket_connection() -> SocketConnection:
    return SocketConnection(
        sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1235"),
-        latency_ms=1.0,
-        sink_to_other_bandwidth_mbps=1000.0,
-        other_to_sink_bandwidth_mbps=1000.0,
    )


--- a/src/exo/routing/router.py
+++ b/src/exo/routing/router.py
@@ -216,8 +216,6 @@ def get_node_id_keypair(
    Obtains the :class:`Keypair` associated with this node-ID.
    Obtain the :class:`PeerId` by from it.
    """
-    # TODO(evan): bring back node id persistence once we figure out how to deal with duplicates
-    return Keypair.generate_ed25519()

    def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
        return Path(str(path) + ".lock")
--- a/src/exo/shared/tests/test_node_id_persistence.py
+++ b/src/exo/shared/tests/test_node_id_persistence.py
@@ -8,7 +8,7 @@ from multiprocessing.synchronize import Event as EventT
 from multiprocessing.synchronize import Semaphore as SemaphoreT

 from loguru import logger
-from pytest import LogCaptureFixture, mark
+from pytest import LogCaptureFixture

 from exo.routing.router import get_node_id_keypair
 from exo.shared.constants import EXO_NODE_ID_KEYPAIR
@@ -74,7 +74,6 @@ def _delete_if_exists(p: str | bytes | os.PathLike[str] | os.PathLike[bytes]):
        os.remove(p)


-@mark.skip(reason="this functionality is currently disabled but may return in future")
 def test_node_id_fetching(caplog: LogCaptureFixture):
    reps = 10

--- a/src/exo/shared/tests/test_state_serialization.py
+++ b/src/exo/shared/tests/test_state_serialization.py
@@ -16,9 +16,6 @@ def test_state_serialization_roundtrip() -> None:
        sink=node_b,
        edge=SocketConnection(
            sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/10001"),
-            latency_ms=1.5,
-            sink_to_other_bandwidth_mbps=1000.0,
-            other_to_sink_bandwidth_mbps=1100.0,
        ),
    )

--- a/src/exo/shared/types/commands.py
+++ b/src/exo/shared/types/commands.py
@@ -49,6 +49,10 @@ class DeleteInstance(BaseCommand):
    instance_id: InstanceId


+class TaskCancelled(BaseCommand):
+    cancelled_command_id: CommandId
+
+
 class TaskFinished(BaseCommand):
    finished_command_id: CommandId

@@ -85,6 +89,7 @@ Command = (
    | PlaceInstance
    | CreateInstance
    | DeleteInstance
+    | TaskCancelled
    | TaskFinished
    | SendInputChunk
 )
--- a/src/exo/shared/types/tasks.py
+++ b/src/exo/shared/types/tasks.py
@@ -25,6 +25,7 @@ class TaskStatus(str, Enum):
    Complete = "Complete"
    TimedOut = "TimedOut"
    Failed = "Failed"
+    Cancelled = "Cancelled"


 class BaseTask(TaggedModel):
@@ -61,6 +62,10 @@ class ChatCompletion(BaseTask):  # emitted by Master
    error_message: str | None = Field(default=None)


+class CancelTask(BaseTask):
+    cancelled_task_id: TaskId
+
+
 class ImageGeneration(BaseTask):  # emitted by Master
    command_id: CommandId
    task_params: ImageGenerationTaskParams
@@ -88,6 +93,7 @@ Task = (
    | LoadModel
    | StartWarmup
    | ChatCompletion
+    | CancelTask
    | ImageGeneration
    | ImageEdits
    | Shutdown
--- a/src/exo/shared/types/topology.py
+++ b/src/exo/shared/types/topology.py
@@ -24,22 +24,10 @@ class RDMAConnection(FrozenModel):

 class SocketConnection(FrozenModel):
    sink_multiaddr: Multiaddr
-    latency_ms: float
-    sink_to_other_bandwidth_mbps: float
-    other_to_sink_bandwidth_mbps: float
-
-    @property
-    def bandwidth_mbps(self):
-        return min(self.sink_to_other_bandwidth_mbps, self.other_to_sink_bandwidth_mbps)

    def __hash__(self):
        return hash(self.sink_multiaddr.ip_address)

-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, SocketConnection):
-            return NotImplemented
-        return self.sink_multiaddr == other.sink_multiaddr
-

 class Connection(FrozenModel):
    source: NodeId
--- a/src/exo/utils/info_gatherer/connection_profiler.py
+++ b/src/exo/utils/info_gatherer/connection_profiler.py
@@ -1,125 +0,0 @@
-import time
-
-import anyio
-import httpx
-from anyio.abc import SocketStream
-from exo.shared.logging import logger
-from pydantic.v1 import BaseModel
-
-LATENCY_PING_COUNT = 5
-BANDWIDTH_TEST_DURATION_S = 0.5
-BANDWIDTH_TEST_PORT_OFFSET = 1  # API port + 1
-
-
-class ConnectionProfile(BaseModel):
-    latency_ms: float
-    upload_mbps: float
-    download_mbps: float
-
-
-async def measure_latency(target_ip: str, port: int = 52415) -> float:
-    if ":" in target_ip:
-        url = f"http://[{target_ip}]:{port}/node_id"
-    else:
-        url = f"http://{target_ip}:{port}/node_id"
-
-    rtts: list[float] = []
-
-    async with httpx.AsyncClient(timeout=10.0) as client:
-        for _ in range(LATENCY_PING_COUNT):
-            try:
-                start = time.perf_counter()
-                response = await client.get(url)
-                end = time.perf_counter()
-
-                if response.status_code == 200:
-                    rtts.append((end - start) * 1000)
-            except (httpx.TimeoutException, httpx.NetworkError, httpx.RemoteProtocolError) as e:
-                logger.debug(f"Latency ping failed: {e}")
-
-    if not rtts:
-        raise ConnectionError(f"Failed to measure latency to {target_ip}:{port}")
-
-    return sum(rtts) / len(rtts)
-
-
-async def _measure_upload_tcp(stream: SocketStream, duration: float) -> float:
-    """Send data for duration seconds, return Mbps."""
-    chunk = b"X" * (1024 * 1024)  # 1MB
-    bytes_sent = 0
-    start = time.perf_counter()
-    deadline = start + duration
-
-    while time.perf_counter() < deadline:
-        await stream.send(chunk)
-        bytes_sent += len(chunk)
-
-    elapsed = time.perf_counter() - start
-    return (bytes_sent * 8 / elapsed) / 1_000_000 if elapsed > 0 else 0.0
-
-
-async def _measure_download_tcp(stream: SocketStream, duration: float) -> float:
-    """Receive data for duration seconds, return Mbps."""
-    bytes_received = 0
-    start = time.perf_counter()
-
-    with anyio.move_on_after(duration):
-        while True:
-            data = await stream.receive(1024 * 1024)
-            if not data:
-                break
-            bytes_received += len(data)
-
-    elapsed = time.perf_counter() - start
-    return (bytes_received * 8 / elapsed) / 1_000_000 if elapsed > 0 else 0.0
-
-
-async def measure_bandwidth_tcp(target_ip: str, port: int) -> tuple[float, float]:
-    """Measure bandwidth using raw TCP like iperf."""
-    upload_mbps = 0.0
-    download_mbps = 0.0
-
-    try:
-        async with await anyio.connect_tcp(target_ip, port) as stream:
-            # Protocol: send 'U' for upload test, 'D' for download test
-            # Upload: client sends, server receives
-            await stream.send(b"U")
-            upload_mbps = await _measure_upload_tcp(stream, BANDWIDTH_TEST_DURATION_S)
-            await stream.send(b"DONE")
-            logger.debug(f"Upload: {upload_mbps:.1f} Mbps")
-    except Exception as e:
-        logger.debug(f"Upload TCP test failed: {e}")
-
-    try:
-        async with await anyio.connect_tcp(target_ip, port) as stream:
-            # Download: client receives, server sends
-            await stream.send(b"D")
-            download_mbps = await _measure_download_tcp(stream, BANDWIDTH_TEST_DURATION_S)
-            logger.debug(f"Download: {download_mbps:.1f} Mbps")
-    except Exception as e:
-        logger.debug(f"Download TCP test failed: {e}")
-
-    return upload_mbps, download_mbps
-
-
-async def profile_connection(target_ip: str, port: int = 52415) -> ConnectionProfile:
-    logger.debug(f"Profiling connection to {target_ip}:{port}")
-
-    latency_ms = await measure_latency(target_ip, port)
-    logger.debug(f"Measured latency to {target_ip}: {latency_ms:.2f}ms")
-
-    bandwidth_port = port + BANDWIDTH_TEST_PORT_OFFSET
-    upload_mbps, download_mbps = await measure_bandwidth_tcp(target_ip, bandwidth_port)
-    logger.debug(
-        f"Measured bandwidth to {target_ip}: "
-        f"upload={upload_mbps:.1f}Mbps, download={download_mbps:.1f}Mbps"
-    )
-
-    if upload_mbps == 0.0 and download_mbps == 0.0:
-        raise ConnectionError(f"Failed to measure bandwidth to {target_ip}:{bandwidth_port}")
-
-    return ConnectionProfile(
-        latency_ms=latency_ms,
-        upload_mbps=upload_mbps,
-        download_mbps=download_mbps,
-    )
--- a/src/exo/worker/engines/mlx/cache.py
+++ b/src/exo/worker/engines/mlx/cache.py
@@ -3,7 +3,6 @@ from copy import deepcopy
 from typing import Any, cast

 import mlx.core as mx
-import psutil
 from mlx_lm.models.cache import (
    KVCache,
    QuantizedKVCache,
@@ -13,29 +12,25 @@ from mlx_lm.models.cache import (
 from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper

-from exo.shared.types.memory import Memory
 from exo.shared.types.mlx import KVCacheType
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
 from exo.worker.runner.bootstrap import logger

 # Fraction of device memory above which LRU eviction kicks in
-_DEFAULT_MEMORY_THRESHOLD = 0.9
+_DEFAULT_MEMORY_THRESHOLD = 0.85
 _MEMORY_THRESHOLD = float(
    os.environ.get("EXO_MEMORY_THRESHOLD", _DEFAULT_MEMORY_THRESHOLD)
 )


 class KVPrefixCache:
-    def __init__(
-        self, tokenizer: TokenizerWrapper, group: mx.distributed.Group | None = None
-    ):
+    def __init__(self, tokenizer: TokenizerWrapper):
        self.prompts: list[mx.array] = []  # mx array of tokens (ints)
        self.caches: list[KVCacheType] = []
        self._last_used: list[int] = []  # monotonic counter of last access per entry
        self._access_counter: int = 0
        self._tokenizer: TokenizerWrapper = tokenizer
-        self._group = group

    def clear(self):
        """Clear all cached prompts and caches."""
@@ -86,13 +81,13 @@ class KVPrefixCache:
        best_snapshot_index, best_snapshot_length = None, 0

        for i, cached_prompt in enumerate(self.prompts):
-            length = get_prefix_length(tokenized_prompt, cached_prompt)
+            length = _get_prefix_length(tokenized_prompt, cached_prompt)

            if length == max_length:
                # Exact match - cached prompt starts with our entire prompt
                # Trim cache to prompt length - 1, return last token for stream_generate
                prompt_cache = deepcopy(self.caches[i])
-                cached_length = cache_length(self.caches[i])
+                cached_length = _cache_length(self.caches[i])
                tokens_to_trim = cached_length - (max_length - 1)
                if tokens_to_trim > 0:
                    trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
@@ -114,7 +109,7 @@ class KVPrefixCache:
            prompt_cache = deepcopy(self.caches[best_snapshot_index])

            # Trim removes tokens from the end, so we trim (cached_length - prefix_length) to keep the prefix
-            cached_length = cache_length(self.caches[best_snapshot_index])
+            cached_length = _cache_length(self.caches[best_snapshot_index])
            tokens_to_trim = cached_length - best_snapshot_length
            if tokens_to_trim > 0:
                trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
@@ -136,37 +131,29 @@ class KVPrefixCache:
            return prompt_cache, tokenized_prompt, None

    def _evict_if_needed(self):
-        """Evict least recently used entries while memory usage is high."""
+        """Evict least recently used entries while memory pressure is high."""
        if len(self.caches) == 0:
            return

+        active: int = mx.metal.get_active_memory()
+        limit = int(mx.metal.device_info()["max_recommended_working_set_size"])
+        if active < limit * _MEMORY_THRESHOLD:
+            return
+
        # Evict LRU entries until below threshold or only one entry left
-        while (
-            len(self.caches) > 1
-            and self.get_memory_used_percentage() > _MEMORY_THRESHOLD
-        ):
+        while len(self.caches) > 0:
            lru_index = self._last_used.index(min(self._last_used))
            evicted_tokens = len(self.prompts[lru_index])
            self.prompts.pop(lru_index)
            self.caches.pop(lru_index)
            self._last_used.pop(lru_index)
            logger.info(
-                f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory usage"
+                f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory pressure"
            )

-    def get_memory_used_percentage(self) -> float:
-        local_pressure: float = get_memory_used_percentage()
-
-        if self._group is None:
-            return local_pressure
-
-        all_pressure = mx.distributed.all_gather(
-            mx.array([local_pressure], dtype=mx.float32),
-            group=self._group,
-        )
-        # .item() evals.
-        max_pressure = float(mx.max(all_pressure).item())
-        return max_pressure
+            active = mx.metal.get_active_memory()
+            if active < limit * _MEMORY_THRESHOLD:
+                break


 def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
@@ -181,13 +168,13 @@ def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
    return mx.array(tokenized_prompt)


-def cache_length(cache: KVCacheType) -> int:
+def _cache_length(cache: KVCacheType) -> int:
    """Get the number of tokens in a KV cache."""
    # Use .offset attribute which all cache types have (len() not implemented in older QuantizedKVCache)
    return max(c.offset for c in cache)  # type: ignore


-def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
+def _get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
    """Find the length of the common prefix between two token arrays."""
    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]))
    if n == 0:
@@ -198,17 +185,6 @@ def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
    return int(mx.sum(prefix_mask).item())


-def get_available_memory() -> Memory:
-    mem: int = psutil.virtual_memory().available
-    return Memory.from_bytes(mem)
-
-
-def get_memory_used_percentage() -> float:
-    mem = psutil.virtual_memory()
-    # percent is 0-100
-    return float(mem.percent / 100)
-
-
 def make_kv_cache(
    model: Model, max_kv_size: int | None = None, keep: int = 0
 ) -> KVCacheType:
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -27,7 +27,6 @@ from exo.worker.engines.mlx.cache import KVPrefixCache, encode_prompt, make_kv_c
 from exo.worker.engines.mlx.constants import KV_BITS, KV_GROUP_SIZE, MAX_TOKENS
 from exo.worker.engines.mlx.utils_mlx import (
    apply_chat_template,
-    mx_barrier,
 )
 from exo.worker.runner.bootstrap import logger

@@ -136,10 +135,6 @@ def warmup_inference(

    logger.info("Generated ALL warmup tokens")

-    # TODO: Do we want an mx_barrier?
-    #  At least this version is actively incorrect, as it should use mx_barrier(group)
-    mx_barrier()
-
    return tokens_generated


@@ -309,5 +304,3 @@ def mlx_generate(
                else:
                    kv_prefix_cache.add_kv_cache(full_prompt, caches)
            break
-
-        # TODO: Do we want an mx_barrier?
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -67,8 +67,6 @@ Group = mx.distributed.Group
 resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))


-# TODO: Test this
-#  ALSO https://github.com/exo-explore/exo/pull/233#discussion_r2549683673
 def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
    return Memory.from_float_kb(
        (model_shard_meta.end_layer - model_shard_meta.start_layer)
@@ -86,30 +84,6 @@ class ModelLoadingTimeoutError(Exception):
    pass


-def mx_barrier(group: Group | None = None):
-    mx.eval(
-        mx.distributed.all_sum(
-            mx.array(1.0),
-            stream=mx.default_stream(mx.Device(mx.cpu)),
-            group=group,
-        )
-    )
-
-
-def broadcast_from_zero(value: int, group: Group | None = None):
-    if group is None:
-        return value
-
-    if group.rank() == 0:
-        a = mx.array([value], dtype=mx.int32)
-    else:
-        a = mx.array([0], dtype=mx.int32)
-
-    m = mx.distributed.all_sum(a, stream=mx.Device(mx.DeviceType.cpu), group=group)
-    mx.eval(m)
-    return int(m.item())
-
-
 class HostList(RootModel[list[str]]):
    @classmethod
    def from_hosts(cls, hosts: list[Host]) -> "HostList":
@@ -538,3 +512,23 @@ def mlx_cleanup(
    import gc

    gc.collect()
+
+
+def mx_any(bool_: bool, group: Group | None) -> bool:
+    if group is None:
+        return bool_
+    num_true = mx.distributed.all_sum(
+        mx.array(bool_), group=group, stream=mx.default_stream(mx.Device(mx.cpu))
+    )
+    mx.eval(num_true)
+    return num_true.item() > 0
+
+
+def mx_barrier(group: Group | None):
+    if group is None:
+        return
+    mx.eval(
+        mx.distributed.all_sum(
+            mx.array(1.0), group=group, stream=mx.default_stream(mx.Device(mx.cpu))
+        )
+    )
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -33,6 +33,7 @@ from exo.shared.types.events import (
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
+    CancelTask,
    CreateRunner,
    DownloadModel,
    ImageEdits,
@@ -44,7 +45,6 @@ from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.runners import RunnerId
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.event_buffer import OrderedBuffer
-from exo.utils.info_gatherer.connection_profiler import profile_connection
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
 from exo.utils.info_gatherer.net_profile import check_reachable
 from exo.utils.keyed_backoff import KeyedBackoff
@@ -66,7 +66,6 @@ class Worker:
        command_sender: Sender[ForwarderCommand],
        download_command_sender: Sender[ForwarderDownloadCommand],
        event_index_counter: Iterator[int],
-        profiling_in_progress: set[NodeId] | None = None,
    ):
        self.node_id: NodeId = node_id
        self.session_id: SessionId = session_id
@@ -83,8 +82,6 @@ class Worker:
        self.state: State = State()
        self.runners: dict[RunnerId, RunnerSupervisor] = {}
        self._tg: TaskGroup = create_task_group()
-        self._profiling_in_progress: set[NodeId] = profiling_in_progress if profiling_in_progress is not None else set()
-        self._profiling_lock = anyio.Lock()

        self._nack_cancel_scope: CancelScope | None = None
        self._nack_attempts: int = 0
@@ -119,8 +116,9 @@ class Worker:
        self.local_event_sender.close()
        self.command_sender.close()
        self.download_command_sender.close()
-        for runner in self.runners.values():
-            runner.shutdown()
+        async with create_task_group() as tg:
+            for runner in self.runners.values():
+                tg.start_soon(runner.shutdown)

    async def _forward_info(self, recv: Receiver[GatheredInfo]):
        with recv as info_stream:
@@ -224,15 +222,22 @@ class Worker:
                        )
                    )
                case Shutdown(runner_id=runner_id):
+                    runner = self.runners.pop(runner_id)
                    try:
                        with fail_after(3):
-                            await self.runners.pop(runner_id).start_task(task)
+                            await runner.start_task(task)
                    except TimeoutError:
                        await self.event_sender.send(
                            TaskStatusUpdated(
                                task_id=task.task_id, task_status=TaskStatus.TimedOut
                            )
                        )
+                    finally:
+                        await runner.shutdown()
+                case CancelTask(cancelled_task_id=cancelled_task_id):
+                    await self.runners[self._task_to_runner_id(task)].cancel_task(
+                        cancelled_task_id
+                    )
                case ImageEdits() if task.task_params.total_input_chunks > 0:
                    # Assemble image from chunks and inject into task
                    cmd_id = task.command_id
@@ -286,73 +291,37 @@ class Worker:
    async def _connection_message_event_writer(self):
        with self.connection_message_receiver as connection_messages:
            async for msg in connection_messages:
-                event = await self._convert_connection_message_to_event(msg)
-                if event:
-                    await self.event_sender.send(event)
+                await self.event_sender.send(
+                    self._convert_connection_message_to_event(msg)
+                )

-    async def _convert_connection_message_to_event(self, msg: ConnectionMessage):
-        if msg.connection_type == ConnectionMessageType.Connected:
-            async with self._profiling_lock:
-                if msg.node_id in self._profiling_in_progress:
-                    return None
-                self._profiling_in_progress.add(msg.node_id)
-            return await self._profile_and_emit_connection(
-                msg.node_id,
-                msg.remote_ipv4,
-                msg.remote_tcp_port,
-            )
-        elif msg.connection_type == ConnectionMessageType.Disconnected:
-            self._profiling_in_progress.discard(msg.node_id)
-            target_ip = msg.remote_ipv4
-            for connection in self.state.topology.list_connections():
-                if (
-                    isinstance(connection.edge, SocketConnection)
-                    and connection.edge.sink_multiaddr.ip_address == target_ip
-                ):
-                    return TopologyEdgeDeleted(conn=connection)
-
-    async def _profile_and_emit_connection(
-        self,
-        sink_node_id: NodeId,
-        remote_ip: str,
-        remote_port: int,
-    ):
-        try:
-            profile = await profile_connection(remote_ip)
-        except ConnectionError as e:
-            logger.warning(f"Failed to profile connection to {sink_node_id}: {e}")
-            profile = None
-
-        if profile:
-            latency_ms = profile.latency_ms
-            other_to_sink_mbps = profile.upload_mbps
-            sink_to_other_mbps = profile.download_mbps
-        else:
-            latency_ms = 0.0
-            other_to_sink_mbps = 0.0
-            sink_to_other_mbps = 0.0
-
-        logger.info(
-            f"Connection to {sink_node_id} profiled: "
-            f"latency={latency_ms:.2f}ms, "
-            f"upload={other_to_sink_mbps:.1f}Mbps, "
-            f"download={sink_to_other_mbps:.1f}Mbps"
-        )
-
-        return TopologyEdgeCreated(
-            conn=Connection(
-                source=self.node_id,
-                sink=sink_node_id,
-                edge=SocketConnection(
-                    sink_multiaddr=Multiaddr(
-                        address=f"/ip4/{remote_ip}/tcp/{remote_port}"
+    def _convert_connection_message_to_event(self, msg: ConnectionMessage):
+        match msg.connection_type:
+            case ConnectionMessageType.Connected:
+                return TopologyEdgeCreated(
+                    conn=Connection(
+                        source=self.node_id,
+                        sink=msg.node_id,
+                        edge=SocketConnection(
+                            sink_multiaddr=Multiaddr(
+                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                            ),
+                        ),
                    ),
-                    latency_ms=latency_ms,
-                    sink_to_other_bandwidth_mbps=sink_to_other_mbps,
-                    other_to_sink_bandwidth_mbps=other_to_sink_mbps,
-                ),
-            ),
-        )
+                )
+
+            case ConnectionMessageType.Disconnected:
+                return TopologyEdgeDeleted(
+                    conn=Connection(
+                        source=self.node_id,
+                        sink=msg.node_id,
+                        edge=SocketConnection(
+                            sink_multiaddr=Multiaddr(
+                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                            ),
+                        ),
+                    ),
+                )

    async def _nack_request(self, since_idx: int) -> None:
        # We request all events after (and including) the missing index.
@@ -391,8 +360,6 @@ class Worker:
            for event in self.out_for_delivery.copy().values():
                await self.local_event_sender.send(event)

-    ## Op Executors
-
    def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor:
        """Creates and stores a new AssignedRunner with initial downloading status."""
        runner = RunnerSupervisor.create(
@@ -429,21 +396,21 @@ class Worker:
            )
            for nid in conns:
                for ip in conns[nid]:
-                    temp_edge = SocketConnection(
+                    edge = SocketConnection(
+                        # nonsense multiaddr
                        sink_multiaddr=Multiaddr(address=f"/ip4/{ip}/tcp/52415")
                        if "." in ip
+                        # nonsense multiaddr
                        else Multiaddr(address=f"/ip6/{ip}/tcp/52415"),
-                        latency_ms=0.0,
-                        sink_to_other_bandwidth_mbps=0.0,
-                        other_to_sink_bandwidth_mbps=0.0,
                    )
-                    if temp_edge not in edges:
-                        logger.debug(f"ping discovered new connection to {nid} at {ip}")
-                        self._tg.start_soon(
-                            self._profile_and_emit_connection,
-                            nid,
-                            ip,
-                            52415,
+                    if edge not in edges:
+                        logger.debug(f"ping discovered {edge=}")
+                        await self.event_sender.send(
+                            TopologyEdgeCreated(
+                                conn=Connection(
+                                    source=self.node_id, sink=nid, edge=edge
+                                )
+                            )
                        )

            for conn in self.state.topology.out_edges(self.node_id):
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -4,6 +4,7 @@ from collections.abc import Mapping, Sequence

 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.tasks import (
+    CancelTask,
    ChatCompletion,
    ConnectToGroup,
    CreateRunner,
@@ -59,7 +60,8 @@ def plan(
        or _init_distributed_backend(runners, all_runners)
        or _load_model(runners, all_runners, global_download_status)
        or _ready_to_warmup(runners, all_runners)
-        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer)
+        or _cancel_tasks(runners, tasks)
+        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer or {})
    )


@@ -270,7 +272,7 @@ def _pending_tasks(
    runners: Mapping[RunnerId, RunnerSupervisor],
    tasks: Mapping[TaskId, Task],
    all_runners: Mapping[RunnerId, RunnerStatus],
-    input_chunk_buffer: Mapping[CommandId, dict[int, str]] | None = None,
+    input_chunk_buffer: Mapping[CommandId, dict[int, str]],
 ) -> Task | None:
    for task in tasks.values():
        # for now, just forward chat completions
@@ -284,7 +286,7 @@ def _pending_tasks(
        if isinstance(task, ImageEdits) and task.task_params.total_input_chunks > 0:
            cmd_id = task.command_id
            expected = task.task_params.total_input_chunks
-            received = len((input_chunk_buffer or {}).get(cmd_id, {}))
+            received = len(input_chunk_buffer.get(cmd_id, {}))
            if received < expected:
                continue  # Wait for all chunks to arrive

@@ -292,16 +294,31 @@ def _pending_tasks(
            if task.instance_id != runner.bound_instance.instance.instance_id:
                continue

-            # I have a design point here; this is a state race in disguise as the task status doesn't get updated to completed fast enough
-            # however, realistically the task status should be set to completed by the LAST runner, so this is a true race
-            # the actual solution is somewhat deeper than this bypass - TODO!
+            # the task status _should_ be set to completed by the LAST runner
+            # it is currently set by the first
+            # this is definitely a hack
            if task.task_id in runner.completed:
                continue

-            # TODO: Check ordering aligns with MLX distributeds expectations.
-
            if isinstance(runner.status, RunnerReady) and all(
                isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
            ):
                return task
+
+
+def _cancel_tasks(
+    runners: Mapping[RunnerId, RunnerSupervisor],
+    tasks: Mapping[TaskId, Task],
+) -> Task | None:
+    for task in tasks.values():
+        if task.task_status != TaskStatus.Cancelled:
+            continue
+        for runner in runners.values():
+            if task.instance_id != runner.bound_instance.instance.instance_id:
+                continue
+            if task.task_id in runner.cancelled:
+                continue
+            return CancelTask(
+                instance_id=task.instance_id, cancelled_task_id=task.task_id
+            )
--- a/src/exo/worker/runner/bootstrap.py
+++ b/src/exo/worker/runner/bootstrap.py
@@ -3,7 +3,7 @@ import os
 import loguru

 from exo.shared.types.events import Event, RunnerStatusUpdated
-from exo.shared.types.tasks import Task
+from exo.shared.types.tasks import Task, TaskId
 from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
@@ -15,6 +15,7 @@ def entrypoint(
    bound_instance: BoundInstance,
    event_sender: MpSender[Event],
    task_receiver: MpReceiver[Task],
+    cancel_receiver: MpReceiver[TaskId],
    _logger: "loguru.Logger",
 ) -> None:
    fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
@@ -38,7 +39,7 @@ def entrypoint(
    try:
        from exo.worker.runner.runner import main

-        main(bound_instance, event_sender, task_receiver)
+        main(bound_instance, event_sender, task_receiver, cancel_receiver)
    except ClosedResourceError:
        logger.warning("Runner communication closed unexpectedly")
    except Exception as e:
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -1,5 +1,6 @@
 import base64
 import json
+import math
 import time
 from collections.abc import Generator
 from functools import cache
@@ -79,6 +80,7 @@ from exo.worker.engines.mlx.utils_mlx import (
    initialize_mlx,
    load_mlx_items,
    mlx_force_oom,
+    mx_any,
 )
 from exo.worker.runner.bootstrap import logger

@@ -87,6 +89,7 @@ def main(
    bound_instance: BoundInstance,
    event_sender: MpSender[Event],
    task_receiver: MpReceiver[Task],
+    cancel_receiver: MpReceiver[TaskId],
 ):
    instance, runner_id, shard_metadata = (
        bound_instance.instance,
@@ -101,23 +104,24 @@ def main(
        time.sleep(timeout)

    setup_start_time = time.time()
+    cancelled_tasks = set[TaskId]()

-    model: Model | DistributedImageModel | None = None
+    # type checker was unhappy with me - splitting these fixed it
+    inference_model: Model | None = None
+    image_model: DistributedImageModel | None = None
    tokenizer = None
    group = None
    kv_prefix_cache: KVPrefixCache | None = None
+    check_for_cancel_every: int | None = None

    current_status: RunnerStatus = RunnerIdle()
    logger.info("runner created")
    event_sender.send(
        RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
    )
-    seen = set[TaskId]()
    with task_receiver as tasks:
        for task in tasks:
-            if task.task_id in seen:
-                logger.warning("repeat task - potential error")
-            seen.add(task.task_id)
+            cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
            event_sender.send(
                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
            )
@@ -162,19 +166,19 @@ def main(
                        time.sleep(0.5)

                    if ModelTask.TextGeneration in shard_metadata.model_card.tasks:
-                        model, tokenizer = load_mlx_items(
+                        inference_model, tokenizer = load_mlx_items(
                            bound_instance, group, on_timeout=on_model_load_timeout
                        )
                        logger.info(
                            f"model has_tool_calling={tokenizer.has_tool_calling}"
                        )
-                        kv_prefix_cache = KVPrefixCache(tokenizer, group)
+                        kv_prefix_cache = KVPrefixCache(tokenizer)

                    elif (
                        ModelTask.TextToImage in shard_metadata.model_card.tasks
                        or ModelTask.ImageToImage in shard_metadata.model_card.tasks
                    ):
-                        model = initialize_image_model(bound_instance)
+                        image_model = initialize_image_model(bound_instance)
                    else:
                        raise ValueError(
                            f"Unknown model task(s): {shard_metadata.model_card.tasks}"
@@ -182,8 +186,6 @@ def main(
                    current_status = RunnerLoaded()
                    logger.info("runner loaded")
                case StartWarmup() if isinstance(current_status, RunnerLoaded):
-                    assert model
-
                    current_status = RunnerWarmingUp()
                    logger.info("runner warming up")
                    event_sender.send(
@@ -194,15 +196,30 @@ def main(

                    logger.info(f"warming up inference for instance: {instance}")
                    if ModelTask.TextGeneration in shard_metadata.model_card.tasks:
-                        assert not isinstance(model, DistributedImageModel)
+                        assert inference_model
                        assert tokenizer

+                        t = time.perf_counter()
                        toks = warmup_inference(
-                            model=model,
+                            model=inference_model,
                            tokenizer=tokenizer,
-                            # kv_prefix_cache=kv_prefix_cache,  # supply for warmup-time prefix caching
                        )
                        logger.info(f"warmed up by generating {toks} tokens")
+                        check_for_cancel_every = min(
+                            math.ceil(toks / (time.perf_counter() - t)), 100
+                        )
+                        if group is not None:
+                            check_for_cancel_every = int(
+                                mx.max(
+                                    mx.distributed.all_gather(
+                                        mx.array([check_for_cancel_every]), group=group
+                                    )
+                                ).item()
+                            )
+
+                        logger.info(
+                            f"runner checking for cancellation every {check_for_cancel_every} tokens"
+                        )
                        logger.info(
                            f"runner initialized in {time.time() - setup_start_time} seconds"
                        )
@@ -210,8 +227,8 @@ def main(
                        ModelTask.TextToImage in shard_metadata.model_card.tasks
                        or ModelTask.ImageToImage in shard_metadata.model_card.tasks
                    ):
-                        assert isinstance(model, DistributedImageModel)
-                        image = warmup_image_generator(model=model)
+                        assert image_model
+                        image = warmup_image_generator(model=image_model)
                        if image is not None:
                            logger.info(f"warmed up by generating {image.size} image")
                        else:
@@ -230,8 +247,9 @@ def main(
                            runner_id=runner_id, runner_status=current_status
                        )
                    )
-                    assert model and not isinstance(model, DistributedImageModel)
+                    assert inference_model
                    assert tokenizer
+                    assert check_for_cancel_every
                    assert task_params.messages[0].content is not None

                    try:
@@ -242,7 +260,7 @@ def main(

                        # Generate responses using the actual MLX generation
                        mlx_generator = mlx_generate(
-                            model=model,
+                            model=inference_model,
                            tokenizer=tokenizer,
                            task=task_params,
                            prompt=prompt,
@@ -266,11 +284,11 @@ def main(
                            patch_glm_tokenizer(tokenizer)

                        # GPT-OSS specific parsing to match other model formats.
-                        elif isinstance(model, GptOssModel):
+                        elif isinstance(inference_model, GptOssModel):
                            mlx_generator = parse_gpt_oss(mlx_generator)

                        if tokenizer.has_tool_calling and not isinstance(
-                            model, GptOssModel
+                            inference_model, GptOssModel
                        ):
                            assert tokenizer.tool_call_start
                            assert tokenizer.tool_call_end
@@ -283,7 +301,18 @@ def main(
                            )

                        completion_tokens = 0
+                        tokens_since_last_cancel_check = 0
                        for response in mlx_generator:
+                            tokens_since_last_cancel_check += 1
+                            if tokens_since_last_cancel_check >= check_for_cancel_every:
+                                tokens_since_last_cancel_check = 0
+                                cancelled_tasks.update(cancel_receiver.collect())
+                                want_to_cancel = (task.task_id in cancelled_tasks) or (
+                                    TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks
+                                )
+                                if mx_any(want_to_cancel, group):
+                                    break
+
                            match response:
                                case GenerationResponse():
                                    completion_tokens += 1
@@ -353,7 +382,7 @@ def main(
                case ImageGeneration(
                    task_params=task_params, command_id=command_id
                ) if isinstance(current_status, RunnerReady):
-                    assert isinstance(model, DistributedImageModel)
+                    assert image_model
                    logger.info(f"received image generation request: {str(task)[:500]}")
                    current_status = RunnerRunning()
                    logger.info("runner running")
@@ -367,7 +396,9 @@ def main(
                        # Generate images using the image generation backend
                        # Track image_index for final images only
                        image_index = 0
-                        for response in generate_image(model=model, task=task_params):
+                        for response in generate_image(
+                            model=image_model, task=task_params
+                        ):
                            if (
                                shard_metadata.device_rank
                                == shard_metadata.world_size - 1
@@ -414,7 +445,7 @@ def main(
                case ImageEdits(task_params=task_params, command_id=command_id) if (
                    isinstance(current_status, RunnerReady)
                ):
-                    assert isinstance(model, DistributedImageModel)
+                    assert image_model
                    logger.info(f"received image edits request: {str(task)[:500]}")
                    current_status = RunnerRunning()
                    logger.info("runner running")
@@ -426,7 +457,9 @@ def main(

                    try:
                        image_index = 0
-                        for response in generate_image(model=model, task=task_params):
+                        for response in generate_image(
+                            model=image_model, task=task_params
+                        ):
                            if (
                                shard_metadata.device_rank
                                == shard_metadata.world_size - 1
@@ -489,7 +522,7 @@ def main(
                RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
            )
            if isinstance(current_status, RunnerShutdown):
-                del model, tokenizer, group
+                del inference_model, image_model, tokenizer, group
                mx.clear_cache()
                import gc

--- a/src/exo/worker/runner/runner_supervisor.py
+++ b/src/exo/worker/runner/runner_supervisor.py
@@ -49,10 +49,12 @@ class RunnerSupervisor:
    _ev_recv: MpReceiver[Event]
    _task_sender: MpSender[Task]
    _event_sender: Sender[Event]
-    _tg: TaskGroup | None = field(default=None, init=False)
+    _cancel_sender: MpSender[TaskId]
+    _tg: TaskGroup = field(default_factory=create_task_group, init=False)
    status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
    pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
    completed: set[TaskId] = field(default_factory=set, init=False)
+    cancelled: set[TaskId] = field(default_factory=set, init=False)

    @classmethod
    def create(
@@ -63,8 +65,8 @@ class RunnerSupervisor:
        initialize_timeout: float = 400,
    ) -> Self:
        ev_send, ev_recv = mp_channel[Event]()
-        # A task is kind of a runner command
        task_sender, task_recv = mp_channel[Task]()
+        cancel_sender, cancel_recv = mp_channel[TaskId]()

        runner_process = Process(
            target=entrypoint,
@@ -72,6 +74,7 @@ class RunnerSupervisor:
                bound_instance,
                ev_send,
                task_recv,
+                cancel_recv,
                logger,
            ),
            daemon=True,
@@ -86,6 +89,7 @@ class RunnerSupervisor:
            initialize_timeout=initialize_timeout,
            _ev_recv=ev_recv,
            _task_sender=task_sender,
+            _cancel_sender=cancel_sender,
            _event_sender=event_sender,
        )

@@ -93,60 +97,64 @@ class RunnerSupervisor:

    async def run(self):
        self.runner_process.start()
-        async with create_task_group() as tg:
-            self._tg = tg
+        async with self._tg as tg:
            tg.start_soon(self._forward_events)

-        self._ev_recv.close()
-        self._task_sender.close()
-        self._event_sender.close()
-        await to_thread.run_sync(self.runner_process.join, 30)
-        if not self.runner_process.is_alive():
-            return
+        with anyio.CancelScope(shield=True), contextlib.suppress(ClosedResourceError):
+            await self._cancel_sender.send_async(TaskId("CANCEL_CURRENT_TASK"))

-        # This is overkill but it's not technically bad, just unnecessary.
-        logger.warning("Runner process didn't shutdown succesfully, terminating")
-        self.runner_process.terminate()
-        await to_thread.run_sync(self.runner_process.join, 5)
-        if not self.runner_process.is_alive():
-            return
+            self._ev_recv.close()
+            self._task_sender.close()
+            self._event_sender.close()
+            self._cancel_sender.close()

-        logger.critical("Runner process didn't respond to SIGTERM, killing")
-        self.runner_process.kill()
+            await to_thread.run_sync(self.runner_process.join, 10)
+            if not self.runner_process.is_alive():
+                return

-        await to_thread.run_sync(self.runner_process.join, 5)
-        if not self.runner_process.is_alive():
-            return
+            # This is overkill but it's not technically bad, just unnecessary.
+            logger.warning("Runner process didn't shutdown succesfully, terminating")
+            self.runner_process.terminate()
+            await to_thread.run_sync(self.runner_process.join, 5)
+            if not self.runner_process.is_alive():
+                return

-        logger.critical(
-            "Runner process didn't respond to SIGKILL. System resources may have leaked"
-        )
+            logger.critical("Runner process didn't respond to SIGTERM, killing")
+            self.runner_process.kill()

-    def shutdown(self):
-        assert self._tg
+            await to_thread.run_sync(self.runner_process.join, 5)
+            if not self.runner_process.is_alive():
+                return
+
+            logger.critical(
+                "Runner process didn't respond to SIGKILL. System resources may have leaked"
+            )
+
+    async def shutdown(self):
+        await self._cancel_sender.send_async(TaskId("CANCEL_CURRENT_TASK"))
        self._tg.cancel_scope.cancel()

    async def start_task(self, task: Task):
-        if task.task_id in self.pending:
-            logger.warning(
-                f"Skipping invalid task {task} as it has already been submitted"
-            )
-            return
-        if task.task_id in self.completed:
-            logger.warning(
-                f"Skipping invalid task {task} as it has already been completed"
-            )
+        if task.task_id in self.completed or task.task_id in self.pending:
+            logger.info(f"Skipping invalid task {task} as it has already been queued")
            return
        logger.info(f"Starting task {task}")
        event = anyio.Event()
        self.pending[task.task_id] = event
        try:
-            await self._task_sender.send_async(task)
+            self._task_sender.send(task)
        except ClosedResourceError:
            logger.warning(f"Task {task} dropped, runner closed communication.")
            return
        await event.wait()

+    async def cancel_task(self, task_id: TaskId):
+        if task_id in self.completed:
+            logger.info(f"Unable to cancel {task_id} as it has been completed")
+            return
+        self.cancelled.add(task_id)
+        await self._cancel_sender.send_async(task_id)
+
    async def _forward_events(self):
        with self._ev_recv as events:
            try:
@@ -211,4 +219,4 @@ class RunnerSupervisor:
                runner_status=RunnerFailed(error_message=f"Terminated ({cause})"),
            )
        )
-        self.shutdown()
+        await self.shutdown()
--- a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py
@@ -14,9 +14,9 @@ from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.cache import (
    KVPrefixCache,
-    cache_length,
+    _cache_length,
+    _get_prefix_length,
    encode_prompt,
-    get_prefix_length,
    make_kv_cache,
 )
 from exo.worker.engines.mlx.generator.generate import mlx_generate, prefill
@@ -35,47 +35,47 @@ class TestGetPrefixLength:
    def test_identical_arrays(self):
        a = mx.array([1, 2, 3, 4, 5])
        b = mx.array([1, 2, 3, 4, 5])
-        assert get_prefix_length(a, b) == 5
+        assert _get_prefix_length(a, b) == 5

    def test_no_common_prefix(self):
        a = mx.array([1, 2, 3])
        b = mx.array([4, 5, 6])
-        assert get_prefix_length(a, b) == 0
+        assert _get_prefix_length(a, b) == 0

    def test_partial_prefix(self):
        a = mx.array([1, 2, 3, 4, 5])
        b = mx.array([1, 2, 3, 7, 8])
-        assert get_prefix_length(a, b) == 3
+        assert _get_prefix_length(a, b) == 3

    def test_prompt_longer_than_cached(self):
        a = mx.array([1, 2, 3, 4, 5])
        b = mx.array([1, 2, 3])
-        assert get_prefix_length(a, b) == 3
+        assert _get_prefix_length(a, b) == 3

    def test_cached_longer_than_prompt(self):
        a = mx.array([1, 2, 3])
        b = mx.array([1, 2, 3, 4, 5])
-        assert get_prefix_length(a, b) == 3
+        assert _get_prefix_length(a, b) == 3

    def test_single_token_match(self):
        a = mx.array([1, 2, 3])
        b = mx.array([1, 5, 6])
-        assert get_prefix_length(a, b) == 1
+        assert _get_prefix_length(a, b) == 1

    def test_empty_prompt(self):
        a = mx.array([]).astype(mx.int32)
        b = mx.array([1, 2, 3])
-        assert get_prefix_length(a, b) == 0
+        assert _get_prefix_length(a, b) == 0

    def test_empty_cached(self):
        a = mx.array([1, 2, 3])
        b = mx.array([]).astype(mx.int32)
-        assert get_prefix_length(a, b) == 0
+        assert _get_prefix_length(a, b) == 0

    def test_both_empty(self):
        a = mx.array([]).astype(mx.int32)
        b = mx.array([]).astype(mx.int32)
-        assert get_prefix_length(a, b) == 0
+        assert _get_prefix_length(a, b) == 0


 class TestKVPrefix:
@@ -146,7 +146,7 @@ class TestKVPrefixCacheWithModel:
        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)

        # Cache should now hold the prompt tokens
-        assert cache_length(cache) == len(tokens)
+        assert _cache_length(cache) == len(tokens)

    def test_add_and_get_exact_match(self, model_and_tokenizer):
        model, tokenizer = model_and_tokenizer
@@ -166,7 +166,7 @@ class TestKVPrefixCacheWithModel:
        kv_prefix_cache.add_kv_cache(prompt, cache)

        assert len(kv_prefix_cache.prompts) == 1
-        stored_length = cache_length(kv_prefix_cache.caches[0])
+        stored_length = _cache_length(kv_prefix_cache.caches[0])
        assert stored_length > 0

        # Retrieve with same prompt: exact match
@@ -209,7 +209,7 @@ class TestKVPrefixCacheWithModel:
        long_tokens = encode_prompt(tokenizer, long_prompt)

        # The prompts share a prefix (chat template preamble + "Hi")
-        expected_prefix = get_prefix_length(long_tokens, short_tokens)
+        expected_prefix = _get_prefix_length(long_tokens, short_tokens)
        assert expected_prefix > 0, (
            "Prompts should share a prefix from the chat template"
        )
@@ -243,7 +243,7 @@ class TestKVPrefixCacheWithModel:
        kv_prefix_cache = KVPrefixCache(tokenizer)
        kv_prefix_cache.add_kv_cache(prompt, cache)

-        stored_length = cache_length(kv_prefix_cache.caches[0])
+        stored_length = _cache_length(kv_prefix_cache.caches[0])

        # Get cache and mutate it (simulating what generation does)
        result_cache, _, matched_index = kv_prefix_cache.get_kv_cache(model, prompt)
@@ -259,7 +259,7 @@ class TestKVPrefixCacheWithModel:
        mx.eval([c.keys for c in result_cache])

        # Stored cache must be unchanged
-        assert cache_length(kv_prefix_cache.caches[0]) == stored_length
+        assert _cache_length(kv_prefix_cache.caches[0]) == stored_length

    def test_stored_cache_survives_repeated_get_mutate_cycles(
        self, model_and_tokenizer
@@ -281,7 +281,7 @@ class TestKVPrefixCacheWithModel:
        kv_prefix_cache = KVPrefixCache(tokenizer)
        kv_prefix_cache.add_kv_cache(prompt, cache)

-        stored_length = cache_length(kv_prefix_cache.caches[0])
+        stored_length = _cache_length(kv_prefix_cache.caches[0])

        for i in range(3):
            result_cache, _, _ = kv_prefix_cache.get_kv_cache(model, prompt)
@@ -293,7 +293,7 @@ class TestKVPrefixCacheWithModel:
                layer_cache.update_and_fetch(extra, extra)
            mx.eval([c.keys for c in result_cache])

-            assert cache_length(kv_prefix_cache.caches[0]) == stored_length, (
+            assert _cache_length(kv_prefix_cache.caches[0]) == stored_length, (
                f"Failed on loop {i}"
            )

@@ -325,7 +325,7 @@ class TestKVPrefixCacheWithModel:
        assert len(kv_prefix_cache.caches) == 1
        # Cache should contain prompt + generated tokens
        expected_length = len(prompt_tokens) + generated_tokens
-        assert cache_length(kv_prefix_cache.caches[0]) == expected_length
+        assert _cache_length(kv_prefix_cache.caches[0]) == expected_length

    def test_mlx_generate_second_call_gets_prefix_hit(self, model_and_tokenizer):
        """Second mlx_generate call with same prompt should get a prefix hit from stored cache."""
@@ -400,7 +400,7 @@ class TestKVPrefixCacheWithModel:
        first_gen_time = time.perf_counter() - t0

        assert len(kv_prefix_cache.prompts) == 1
-        first_cache_length = cache_length(kv_prefix_cache.caches[0])
+        first_cache_length = _cache_length(kv_prefix_cache.caches[0])

        # Second generation: same long prompt + extra content (simulating multi-turn)
        task2 = ChatCompletionTaskParams(
@@ -416,7 +416,7 @@ class TestKVPrefixCacheWithModel:
        prompt2_tokens = encode_prompt(tokenizer, prompt2)

        # Verify the prompts share a long prefix
-        prefix_len = get_prefix_length(prompt2_tokens, prompt1_tokens)
+        prefix_len = _get_prefix_length(prompt2_tokens, prompt1_tokens)
        assert prefix_len > 1000, "Prompts must share > 1000 token prefix"

        # Second generation should reuse the cached prefix (only prefill new tokens)
@@ -440,7 +440,7 @@ class TestKVPrefixCacheWithModel:
        # With prefix_hit > 1000, should update in-place (not add a second entry)
        assert len(kv_prefix_cache.prompts) == 1
        # Updated cache should be longer (prompt2 + generated > prompt1 + generated)
-        updated_cache_length = cache_length(kv_prefix_cache.caches[0])
+        updated_cache_length = _cache_length(kv_prefix_cache.caches[0])
        assert updated_cache_length > first_cache_length

    def test_mlx_generate_stored_cache_not_mutated(self, model_and_tokenizer):
@@ -465,7 +465,7 @@ class TestKVPrefixCacheWithModel:
        ):
            pass

-        firstcache_length = cache_length(kv_prefix_cache.caches[0])
+        first_cache_length = _cache_length(kv_prefix_cache.caches[0])

        # Second generation gets the cache and mutates it during generation
        for _response in mlx_generate(
@@ -478,7 +478,7 @@ class TestKVPrefixCacheWithModel:
            pass

        # The first stored cache must not have been mutated by the second generation
-        assert cache_length(kv_prefix_cache.caches[0]) == firstcache_length
+        assert _cache_length(kv_prefix_cache.caches[0]) == first_cache_length

    def test_evicts_lru_entry_under_memory_pressure(self, model_and_tokenizer):
        """Under memory pressure, adding a new cache entry evicts the least recently used one."""
@@ -540,6 +540,6 @@ class TestKVPrefixCacheWithModel:
        assert len(kv_prefix_cache.prompts) == 1
        # The surviving entry should be the newly added one
        new_tokens = encode_prompt(tokenizer, prompt)
-        assert get_prefix_length(kv_prefix_cache.prompts[0], new_tokens) == len(
+        assert _get_prefix_length(kv_prefix_cache.prompts[0], new_tokens) == len(
            new_tokens
        )
--- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
@@ -109,8 +109,8 @@ def assert_events_equal(test_events: Iterable[Event], true_events: Iterable[Even

@pytest.fixture
 def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
-    # initialize_mlx returns a mock group
-    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup()))
+    # initialize_mlx returns a "group" equal to 1
+    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(1))
    monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer)))
    monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
    monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
@@ -147,14 +147,6 @@ class MockTokenizer:
    has_tool_calling = False


-class MockGroup:
-    def rank(self) -> int:
-        return 0
-
-    def size(self) -> int:
-        return 1
-
-
 def _run(tasks: Iterable[Task]):
    bound_instance = get_bound_mlx_ring_instance(
        instance_id=INSTANCE_1_ID,