Add MLX compute device backend selection (cpu/gpu/auto)

Add MlxDevice enum (Auto/Cpu/Gpu) that flows from the API through
PlaceInstance commands to instance creation and runner initialization.
The device is set via mx.set_default_device() in bootstrap.py before
the runner imports MLX, ensuring all subsequent operations use the
chosen device.

This enables Linux CPU-only machines to participate in inference and
allows forcing CPU mode for testing or when GPU has issues. Default
is Auto (current behavior) for full backward compatibility.

Closes #958

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

app/EXO/EXO/ExoProcessController.swift

@@ -126,11 +126,37 @@ final class ExoProcessController: ObservableObject {
             return
         }
         process.terminationHandler = nil
-        if process.isRunning {
-            process.terminate()
-        }
-        self.process = nil
         status = .stopped
+
+        guard process.isRunning else {
+            self.process = nil
+            return
+        }
+
+        let proc = process
+        self.process = nil
+
+        Task.detached {
+            proc.interrupt()
+
+            for _ in 0..<50 {
+                if !proc.isRunning { return }
+                try? await Task.sleep(nanoseconds: 100_000_000)
+            }
+
+            if proc.isRunning {
+                proc.terminate()
+            }
+
+            for _ in 0..<30 {
+                if !proc.isRunning { return }
+                try? await Task.sleep(nanoseconds: 100_000_000)
+            }
+
+            if proc.isRunning {
+                kill(proc.processIdentifier, SIGKILL)
+            }
+        }
     }
 
     func restart() {

pyproject.toml

@@ -132,7 +132,7 @@ markers = [
 env = [
   "EXO_TESTS=1"
 ]
-addopts = "-m 'not slow'"
+addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
 filterwarnings = [
     "ignore:builtin type Swig:DeprecationWarning",
 ]

src/exo/download/coordinator.py

@@ -14,6 +14,7 @@ from exo.download.download_utils import (
     map_repo_download_progress_to_download_progress_data,
 )
 from exo.download.shard_downloader import ShardDownloader
+from exo.shared.constants import EXO_MODELS_DIR
 from exo.shared.models.model_cards import ModelId
 from exo.shared.types.commands import (
     CancelDownload,
@@ -63,6 +64,9 @@ class DownloadCoordinator:
         self.event_sender, self.event_receiver = channel[Event]()
         self.shard_downloader.on_progress(self._download_progress_callback)
 
+    def _model_dir(self, model_id: ModelId) -> str:
+        return str(EXO_MODELS_DIR / model_id.normalize())
+
     async def _download_progress_callback(
         self, callback_shard: ShardMetadata, progress: RepoDownloadProgress
     ) -> None:
@@ -74,6 +78,7 @@ class DownloadCoordinator:
                 shard_metadata=callback_shard,
                 node_id=self.node_id,
                 total_bytes=progress.total_bytes,
+                model_directory=self._model_dir(model_id),
             )
             self.download_status[model_id] = completed
             await self.event_sender.send(
@@ -93,6 +98,7 @@ class DownloadCoordinator:
                 download_progress=map_repo_download_progress_to_download_progress_data(
                     progress
                 ),
+                model_directory=self._model_dir(model_id),
             )
             self.download_status[model_id] = ongoing
             await self.event_sender.send(
@@ -170,7 +176,11 @@ class DownloadCoordinator:
                 return
 
         # Emit pending status
-        progress = DownloadPending(shard_metadata=shard, node_id=self.node_id)
+        progress = DownloadPending(
+            shard_metadata=shard,
+            node_id=self.node_id,
+            model_directory=self._model_dir(model_id),
+        )
         self.download_status[model_id] = progress
         await self.event_sender.send(NodeDownloadProgress(download_progress=progress))
 
@@ -184,6 +194,7 @@ class DownloadCoordinator:
                 shard_metadata=shard,
                 node_id=self.node_id,
                 total_bytes=initial_progress.total_bytes,
+                model_directory=self._model_dir(model_id),
             )
             self.download_status[model_id] = completed
             await self.event_sender.send(
@@ -206,6 +217,7 @@ class DownloadCoordinator:
             download_progress=map_repo_download_progress_to_download_progress_data(
                 initial_progress
             ),
+            model_directory=self._model_dir(model_id),
         )
         self.download_status[model_id] = status
         self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
@@ -219,6 +231,7 @@ class DownloadCoordinator:
                     shard_metadata=shard,
                     node_id=self.node_id,
                     error_message=str(e),
+                    model_directory=self._model_dir(model_id),
                 )
                 self.download_status[model_id] = failed
                 await self.event_sender.send(
@@ -253,6 +266,7 @@ class DownloadCoordinator:
             pending = DownloadPending(
                 shard_metadata=current_status.shard_metadata,
                 node_id=self.node_id,
+                model_directory=self._model_dir(model_id),
             )
             await self.event_sender.send(
                 NodeDownloadProgress(download_progress=pending)
@@ -295,11 +309,18 @@ class DownloadCoordinator:
                             node_id=self.node_id,
                             shard_metadata=progress.shard,
                             total_bytes=progress.total_bytes,
+                            model_directory=self._model_dir(
+                                progress.shard.model_card.model_id
+                            ),
                         )
                     elif progress.status in ["in_progress", "not_started"]:
                         if progress.downloaded_bytes_this_session.in_bytes == 0:
                             status = DownloadPending(
-                                node_id=self.node_id, shard_metadata=progress.shard
+                                node_id=self.node_id,
+                                shard_metadata=progress.shard,
+                                model_directory=self._model_dir(
+                                    progress.shard.model_card.model_id
+                                ),
                             )
                         else:
                             status = DownloadOngoing(
@@ -308,6 +329,9 @@ class DownloadCoordinator:
                                 download_progress=map_repo_download_progress_to_download_progress_data(
                                     progress
                                 ),
+                                model_directory=self._model_dir(
+                                    progress.shard.model_card.model_id
+                                ),
                             )
                     else:
                         continue

src/exo/main.py

@@ -136,6 +136,8 @@ class Node:
 
     async def run(self):
         async with self._tg as tg:
+            signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
+            signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
             tg.start_soon(self.router.run)
             tg.start_soon(self.election.run)
             if self.download_coordinator:
@@ -147,8 +149,6 @@ class Node:
             if self.api:
                 tg.start_soon(self.api.run)
             tg.start_soon(self._elect_loop)
-            signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
-            signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
 
     def shutdown(self):
         # if this is our second call to shutdown, just sys.exit

src/exo/master/api.py

@@ -143,7 +143,12 @@ from exo.shared.types.openai_responses import (
     ResponsesResponse,
 )
 from exo.shared.types.state import State
-from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
+from exo.shared.types.worker.instances import (
+    Instance,
+    InstanceId,
+    InstanceMeta,
+    MlxDevice,
+)
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
 from exo.utils.channels import Receiver, Sender, channel
@@ -310,6 +315,7 @@ class API:
             sharding=payload.sharding,
             instance_meta=payload.instance_meta,
             min_nodes=payload.min_nodes,
+            mlx_device=payload.mlx_device,
         )
         await self._send(command)
 
@@ -350,6 +356,7 @@ class API:
         sharding: Sharding = Sharding.Pipeline,
         instance_meta: InstanceMeta = InstanceMeta.MlxRing,
         min_nodes: int = 1,
+        mlx_device: MlxDevice = MlxDevice.Auto,
     ) -> Instance:
         model_card = await ModelCard.load(model_id)
 
@@ -360,6 +367,7 @@ class API:
                     sharding=sharding,
                     instance_meta=instance_meta,
                     min_nodes=min_nodes,
+                    mlx_device=mlx_device,
                 ),
                 node_memory=self.state.node_memory,
                 node_network=self.state.node_network,

src/exo/master/placement.py

@@ -159,6 +159,7 @@ def place_instance(
                 shard_assignments=shard_assignments,
                 jaccl_devices=mlx_jaccl_devices,
                 jaccl_coordinators=mlx_jaccl_coordinators,
+                mlx_device=command.mlx_device,
             )
         case InstanceMeta.MlxRing:
             ephemeral_port = random_ephemeral_port()
@@ -173,6 +174,7 @@ def place_instance(
                 shard_assignments=shard_assignments,
                 hosts_by_node=hosts_by_node,
                 ephemeral_port=ephemeral_port,
+                mlx_device=command.mlx_device,
             )
 
     return target_instances

src/exo/shared/types/api.py

@@ -8,7 +8,12 @@ from pydantic import BaseModel, Field
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
+from exo.shared.types.worker.instances import (
+    Instance,
+    InstanceId,
+    InstanceMeta,
+    MlxDevice,
+)
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel
 
@@ -226,6 +231,7 @@ class PlaceInstanceParams(BaseModel):
     sharding: Sharding = Sharding.Pipeline
     instance_meta: InstanceMeta = InstanceMeta.MlxRing
     min_nodes: int = 1
+    mlx_device: MlxDevice = MlxDevice.Auto
 
 
 class CreateInstanceParams(BaseModel):

src/exo/shared/types/commands.py

@@ -8,7 +8,12 @@ from exo.shared.types.api import (
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.text_generation import TextGenerationTaskParams
-from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
+from exo.shared.types.worker.instances import (
+    Instance,
+    InstanceId,
+    InstanceMeta,
+    MlxDevice,
+)
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
@@ -38,6 +43,7 @@ class PlaceInstance(BaseCommand):
     sharding: Sharding
     instance_meta: InstanceMeta
     min_nodes: int
+    mlx_device: MlxDevice = MlxDevice.Auto
 
 
 class CreateInstance(BaseCommand):

src/exo/shared/types/worker/downloads.py

@@ -26,6 +26,7 @@ class DownloadProgressData(CamelCaseModel):
 class BaseDownloadProgress(TaggedModel):
     node_id: NodeId
     shard_metadata: ShardMetadata
+    model_directory: str = ""
 
 
 class DownloadPending(BaseDownloadProgress):

src/exo/shared/types/worker/instances.py

@@ -16,9 +16,16 @@ class InstanceMeta(str, Enum):
     MlxJaccl = "MlxJaccl"
 
 
+class MlxDevice(str, Enum):
+    Auto = "Auto"
+    Cpu = "Cpu"
+    Gpu = "Gpu"
+
+
 class BaseInstance(TaggedModel):
     instance_id: InstanceId
     shard_assignments: ShardAssignments
+    mlx_device: MlxDevice = MlxDevice.Auto
 
     def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
         return self.shard_assignments.runner_to_shard.get(runner_id, None)

src/exo/worker/runner/bootstrap.py

@@ -4,7 +4,7 @@ import loguru
 
 from exo.shared.types.events import Event, RunnerStatusUpdated
 from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance
+from exo.shared.types.worker.instances import BoundInstance, MlxDevice, MlxJacclInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 
@@ -35,6 +35,15 @@ def entrypoint(
 
     logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
 
+    # Set MLX compute device before importing runner (which imports mlx.core at module scope)
+    mlx_device = bound_instance.instance.mlx_device
+    if mlx_device != MlxDevice.Auto:
+        import mlx.core as mx
+
+        device = mx.cpu if mlx_device == MlxDevice.Cpu else mx.gpu
+        mx.set_default_device(device)
+        logger.info(f"MLX device set to: {mlx_device}")
+
     # Import main after setting global logger - this lets us just import logger from this module
     try:
         from exo.worker.runner.runner import main