fix: keep TRUST_REMOTE_CODE=True for built-in models

The constant is the default for built-in models with known model cards,
which are trusted. Custom models added via API already default to
trust_remote_code=False in ModelCard.fetch_from_hf(). The CLI flag
overrides custom models only.

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

app/EXO/EXO/ExoProcessController.swift

@@ -5,7 +5,6 @@ import Foundation
 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
 private let enableImageModelsKey = "EXOEnableImageModels"
-private let offlineModeKey = "EXOOfflineMode"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
 
 @MainActor
@@ -61,14 +60,6 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
         }
     }
-    @Published var offlineMode: Bool = {
-        return UserDefaults.standard.bool(forKey: offlineModeKey)
-    }()
-    {
-        didSet {
-            UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
-        }
-    }
 
     /// Fires once when EXO transitions to `.running` for the very first time (fresh install).
     @Published private(set) var isFirstLaunchReady = false
@@ -276,9 +267,6 @@ final class ExoProcessController: ObservableObject {
         if enableImageModels {
             environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
         }
-        if offlineMode {
-            environment["EXO_OFFLINE"] = "true"
-        }
 
         var paths: [String] = []
         if let existing = environment["PATH"], !existing.isEmpty {

app/EXO/EXO/Views/SettingsView.swift

@@ -13,7 +13,6 @@ struct SettingsView: View {
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
     @State private var pendingEnableImageModels = false
-    @State private var pendingOfflineMode = false
     @State private var needsRestart = false
     @State private var bugReportInFlight = false
     @State private var bugReportMessage: String?
@@ -43,7 +42,6 @@ struct SettingsView: View {
             pendingNamespace = controller.customNamespace
             pendingHFToken = controller.hfToken
             pendingEnableImageModels = controller.enableImageModels
-            pendingOfflineMode = controller.offlineMode
             needsRestart = false
         }
     }
@@ -74,13 +72,6 @@ struct SettingsView: View {
                     .foregroundColor(.secondary)
             }
 
-            Section {
-                Toggle("Offline Mode", isOn: $pendingOfflineMode)
-                Text("Skip internet checks and use only locally available models.")
-                    .font(.caption)
-                    .foregroundColor(.secondary)
-            }
-
             Section {
                 HStack {
                     Spacer()
@@ -454,7 +445,6 @@ struct SettingsView: View {
 
     private var hasGeneralChanges: Bool {
         pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
-            || pendingOfflineMode != controller.offlineMode
     }
 
     private var hasModelChanges: Bool {
@@ -464,7 +454,6 @@ struct SettingsView: View {
     private func applyGeneralSettings() {
         controller.customNamespace = pendingNamespace
         controller.hfToken = pendingHFToken
-        controller.offlineMode = pendingOfflineMode
         restartIfRunning()
     }
 

dashboard/src/routes/+page.svelte

@@ -42,7 +42,6 @@
     setSelectedChatModel,
     selectedChatModel,
     sendMessage,
-    messages,
     debugMode,
     toggleDebugMode,
     topologyOnlyMode,
@@ -890,7 +889,6 @@
       availableModels.some((m) => m.id === defaults.modelId)
     ) {
       selectPreviewModel(defaults.modelId);
-      setSelectedChatModel(defaults.modelId);
     }
   }
 
@@ -1326,7 +1324,6 @@
 
   function handleModelPickerSelect(modelId: string) {
     selectPreviewModel(modelId);
-    setSelectedChatModel(modelId);
     saveLaunchDefaults();
     isModelPickerOpen = false;
   }
@@ -2287,8 +2284,7 @@
     selectedChatCategory = null;
     pendingAutoMessage = null;
     userForcedIdle = true;
-    // Restore chat model from the sidebar preview selection so both selectors stay in sync
-    setSelectedChatModel(selectedModelId ?? "");
+    setSelectedChatModel("");
     clearChat();
   }
 
@@ -2518,11 +2514,6 @@
         return;
       }
     }
-
-    // Fallthrough: model exists but has no active instance/download/loading state
-    chatLaunchState = "idle";
-    pendingChatModelId = null;
-    selectedChatCategory = null;
   });
 
   // Suggested prompts per category
@@ -2646,11 +2637,7 @@
   }
 
   // Launch a model for seamless chat
-  async function launchModelForChat(
-    modelId: string,
-    category: string,
-    skipCreate = false,
-  ) {
+  async function launchModelForChat(modelId: string, category: string) {
     userForcedIdle = false;
     pendingChatModelId = modelId;
     selectedChatCategory = category;
@@ -2658,7 +2645,7 @@
     // Check if already running — skip straight to chat
     if (hasRunningInstance(modelId)) {
       setSelectedChatModel(modelId);
-      if (!skipCreate) createConversation();
+      createConversation();
       chatLaunchState = "ready";
       return;
     }
@@ -2667,7 +2654,7 @@
     if (hasExistingInstance(modelId)) {
       setSelectedChatModel(modelId);
       pendingChatModelId = modelId;
-      if (!skipCreate) createConversation();
+      createConversation();
       const dlStatus = getModelDownloadStatus(modelId);
       if (dlStatus.isDownloading) {
         chatLaunchState = "downloading";
@@ -2720,7 +2707,7 @@
 
       setSelectedChatModel(modelId);
       recordRecentLaunch(modelId);
-      if (!skipCreate) createConversation();
+      createConversation();
       chatLaunchState = "downloading";
     } catch (error) {
       addToast({ type: "error", message: `Network error: ${error}` });
@@ -3006,7 +2993,6 @@
   // Handle model selection from the picker when opened from chat context
   function handleChatPickerSelect(modelId: string) {
     setSelectedChatModel(modelId);
-    selectPreviewModel(modelId);
     userForcedIdle = false;
     isModelPickerOpen = false;
   }
@@ -3026,7 +3012,6 @@
 
     // Model is selected and running — send directly
     if (model && hasRunningInstance(model)) {
-      chatLaunchState = "ready";
       sendMessage(content, files, null);
       return;
     }
@@ -3035,7 +3020,7 @@
     if (model) {
       pendingAutoMessage = { content, files };
       userForcedIdle = false;
-      launchModelForChat(model, "picker", messages().length > 0);
+      launchModelForChat(model, "picker");
       return;
     }
 
@@ -5819,7 +5804,43 @@
           class="flex-1 flex flex-col min-w-0 overflow-hidden"
           in:fade={{ duration: 300, delay: 100 }}
         >
-          {#if chatLaunchState !== "idle" && chatLaunchState !== "ready"}
+          {#if chatLaunchState === "idle"}
+            <!-- No running instance: show model selector -->
+            <div
+              class="flex-1 overflow-y-auto flex items-center justify-center px-8 py-6"
+            >
+              <ChatModelSelector
+                models={models.map((m) => ({
+                  id: m.id,
+                  name: m.name ?? "",
+                  base_model: m.base_model ?? "",
+                  storage_size_megabytes: m.storage_size_megabytes ?? 0,
+                  capabilities: m.capabilities ?? [],
+                  family: m.family ?? "",
+                  quantization: m.quantization ?? "",
+                }))}
+                clusterLabel={chatClusterLabel}
+                totalMemoryGB={availableMemoryGB()}
+                onSelect={handleChatModelSelect}
+                onAddModel={handleChatAddModel}
+              />
+            </div>
+            <div
+              class="flex-shrink-0 px-8 pb-6 pt-4 bg-gradient-to-t from-exo-black via-exo-black to-transparent"
+            >
+              <div class="max-w-7xl mx-auto">
+                <ChatForm
+                  placeholder="Ask anything — we'll pick the best model automatically"
+                  showModelSelector={!!bestRunningModelId}
+                  modelDisplayOverride={bestRunningModelId ?? undefined}
+                  modelTasks={modelTasks()}
+                  modelCapabilities={modelCapabilities()}
+                  onAutoSend={handleAutoSend}
+                  onOpenModelPicker={openChatModelPicker}
+                />
+              </div>
+            </div>
+          {:else if chatLaunchState !== "idle" && chatLaunchState !== "ready"}
             <!-- Model launching/downloading/loading: show progress -->
             <div class="flex-1 flex items-center justify-center px-8 py-6">
               <div class="flex flex-col items-center gap-6 max-w-md w-full">
@@ -5926,8 +5947,8 @@
                 />
               </div>
             </div>
-          {:else if messages().length > 0 || chatLaunchState === "ready"}
-            <!-- Normal chat: show messages -->
+          {:else}
+            <!-- Normal chat: model is running -->
             <div
               class="flex-1 overflow-y-auto px-8 py-6"
               bind:this={chatScrollRef}
@@ -5983,42 +6004,6 @@
                 />
               </div>
             </div>
-          {:else}
-            <!-- No running instance, no messages: show model selector -->
-            <div
-              class="flex-1 overflow-y-auto flex items-center justify-center px-8 py-6"
-            >
-              <ChatModelSelector
-                models={models.map((m) => ({
-                  id: m.id,
-                  name: m.name ?? "",
-                  base_model: m.base_model ?? "",
-                  storage_size_megabytes: m.storage_size_megabytes ?? 0,
-                  capabilities: m.capabilities ?? [],
-                  family: m.family ?? "",
-                  quantization: m.quantization ?? "",
-                }))}
-                clusterLabel={chatClusterLabel}
-                totalMemoryGB={availableMemoryGB()}
-                onSelect={handleChatModelSelect}
-                onAddModel={handleChatAddModel}
-              />
-            </div>
-            <div
-              class="flex-shrink-0 px-8 pb-6 pt-4 bg-gradient-to-t from-exo-black via-exo-black to-transparent"
-            >
-              <div class="max-w-7xl mx-auto">
-                <ChatForm
-                  placeholder="Ask anything — we'll pick the best model automatically"
-                  showModelSelector={!!bestRunningModelId}
-                  modelDisplayOverride={bestRunningModelId ?? undefined}
-                  modelTasks={modelTasks()}
-                  modelCapabilities={modelCapabilities()}
-                  onAutoSend={handleAutoSend}
-                  onOpenModelPicker={openChatModelPicker}
-                />
-              </div>
-            </div>
           {/if}
         </div>
 

nix/mlx.nix

@@ -41,7 +41,7 @@ let
 
   mlx = stdenv.mkDerivation rec {
     pname = "mlx";
-    version = let v = "0.30.7.dev20260225+257d5692"; in
+    version = let v = "0.30.7.dev20260220+13998a05"; in
       assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
       v;
     pyproject = true;
@@ -49,8 +49,8 @@ let
     src = fetchFromGitHub {
       owner = "rltakashige";
       repo = "mlx-jaccl-fix-small-recv";
-      rev = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
-      hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
+      rev = "13998a054715edcdc93618fb1496c79c7c25ff7c";
+      hash = "sha256-fAqA3hFwNBx7FcoGnhQsIFpAIRbC2EerACm4Fvne0Cc=";
     };
 
     patches = [

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "exo"
-version = "0.3.68"
+version = "0.3.0"
 description = "Exo"
 readme = "README.md"
 requires-python = ">=3.13"

src/exo/download/coordinator.py

@@ -1,4 +1,5 @@
 import asyncio
+import socket
 from dataclasses import dataclass, field
 from random import random
 
@@ -72,6 +73,8 @@ class DownloadCoordinator:
 
     def __post_init__(self) -> None:
         self.event_sender, self.event_receiver = channel[Event]()
+        if self.offline:
+            self.shard_downloader.set_internet_connection(False)
         self.shard_downloader.on_progress(self._download_progress_callback)
 
     def _model_dir(self, model_id: ModelId) -> str:
@@ -120,6 +123,8 @@ class DownloadCoordinator:
         logger.info(
             f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
         )
+        if not self.offline:
+            self._test_internet_connection()
         try:
             async with self._tg as tg:
                 tg.start_soon(self._command_processor)
@@ -127,10 +132,40 @@ class DownloadCoordinator:
                 tg.start_soon(self._emit_existing_download_progress)
                 tg.start_soon(self._resend_out_for_delivery)
                 tg.start_soon(self._clear_ofd)
+                if not self.offline:
+                    tg.start_soon(self._check_internet_connection)
         finally:
             for task in self.active_downloads.values():
                 task.cancel()
 
+    def _test_internet_connection(self) -> None:
+        # Try multiple endpoints since some ISPs/networks block specific IPs
+        for host in ("1.1.1.1", "8.8.8.8", "1.0.0.1"):
+            try:
+                socket.create_connection((host, 443), timeout=3).close()
+                self.shard_downloader.set_internet_connection(True)
+                logger.debug(f"Internet connectivity: True (via {host})")
+                return
+            except OSError:
+                continue
+        self.shard_downloader.set_internet_connection(False)
+        logger.debug("Internet connectivity: False")
+
+    async def _check_internet_connection(self) -> None:
+        first_connection = True
+        while True:
+            await asyncio.sleep(10)
+
+            # Assume that internet connection is set to False on 443 errors.
+            if self.shard_downloader.internet_connection:
+                continue
+
+            self._test_internet_connection()
+
+            if first_connection and self.shard_downloader.internet_connection:
+                first_connection = False
+                self._tg.start_soon(self._emit_existing_download_progress)
+
     def shutdown(self) -> None:
         self._tg.cancel_tasks()
 

src/exo/download/impl_shard_downloader.py

@@ -15,13 +15,9 @@ from exo.shared.types.worker.shards import (
 )
 
 
-def exo_shard_downloader(
-    max_parallel_downloads: int = 8, offline: bool = False
-) -> ShardDownloader:
+def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
     return SingletonShardDownloader(
-        CachedShardDownloader(
-            ResumableShardDownloader(max_parallel_downloads, offline=offline)
-        )
+        CachedShardDownloader(ResumableShardDownloader(max_parallel_downloads))
     )
 
 
@@ -54,6 +50,10 @@ class SingletonShardDownloader(ShardDownloader):
         self.shard_downloader = shard_downloader
         self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}
 
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+        self.shard_downloader.set_internet_connection(value)
+
     def on_progress(
         self,
         callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
@@ -90,6 +90,10 @@ class CachedShardDownloader(ShardDownloader):
         self.shard_downloader = shard_downloader
         self.cache: dict[tuple[str, ShardMetadata], Path] = {}
 
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+        self.shard_downloader.set_internet_connection(value)
+
     def on_progress(
         self,
         callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
@@ -119,9 +123,8 @@ class CachedShardDownloader(ShardDownloader):
 
 
 class ResumableShardDownloader(ShardDownloader):
-    def __init__(self, max_parallel_downloads: int = 8, offline: bool = False):
+    def __init__(self, max_parallel_downloads: int = 8):
         self.max_parallel_downloads = max_parallel_downloads
-        self.offline = offline
         self.on_progress_callbacks: list[
             Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
         ] = []
@@ -148,7 +151,8 @@ class ResumableShardDownloader(ShardDownloader):
             self.on_progress_wrapper,
             max_parallel_downloads=self.max_parallel_downloads,
             allow_patterns=allow_patterns,
-            skip_internet=self.offline,
+            skip_internet=not self.internet_connection,
+            on_connection_lost=lambda: self.set_internet_connection(False),
         )
         return target_dir
 
@@ -164,7 +168,8 @@ class ResumableShardDownloader(ShardDownloader):
                 shard,
                 self.on_progress_wrapper,
                 skip_download=True,
-                skip_internet=self.offline,
+                skip_internet=not self.internet_connection,
+                on_connection_lost=lambda: self.set_internet_connection(False),
             )
 
         semaphore = asyncio.Semaphore(self.max_parallel_downloads)
@@ -193,6 +198,7 @@ class ResumableShardDownloader(ShardDownloader):
             shard,
             self.on_progress_wrapper,
             skip_download=True,
-            skip_internet=self.offline,
+            skip_internet=not self.internet_connection,
+            on_connection_lost=lambda: self.set_internet_connection(False),
         )
         return progress

src/exo/download/shard_downloader.py

@@ -16,6 +16,11 @@ from exo.shared.types.worker.shards import (
 
 # TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?
 class ShardDownloader(ABC):
+    internet_connection: bool = False
+
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+
     @abstractmethod
     async def ensure_shard(
         self, shard: ShardMetadata, config_only: bool = False

src/exo/main.py

@@ -60,7 +60,7 @@ class Node:
             download_coordinator = DownloadCoordinator(
                 node_id,
                 session_id,
-                exo_shard_downloader(offline=args.offline),
+                exo_shard_downloader(),
                 download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
                 local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 offline=args.offline,
@@ -211,7 +211,7 @@ class Node:
                         self.download_coordinator = DownloadCoordinator(
                             self.node_id,
                             result.session_id,
-                            exo_shard_downloader(offline=self.offline),
+                            exo_shard_downloader(),
                             download_command_receiver=self.router.receiver(
                                 topics.DOWNLOAD_COMMANDS
                             ),
@@ -261,6 +261,13 @@ def main():
     if args.offline:
         logger.info("Running in OFFLINE mode — no internet checks, local models only")
 
+    # Set trust_remote_code override env var for runner subprocesses
+    if args.trust_remote_code:
+        os.environ["EXO_TRUST_REMOTE_CODE"] = "1"
+        logger.warning(
+            "--trust-remote-code enabled: models may execute arbitrary code during loading"
+        )
+
     # Set FAST_SYNCH override env var for runner subprocesses
     if args.fast_synch is True:
         os.environ["EXO_FAST_SYNCH"] = "on"
@@ -270,16 +277,9 @@ def main():
         logger.info("FAST_SYNCH forced OFF")
 
     node = anyio.run(Node.create, args)
-    try:
-        anyio.run(node.run)
-    except BaseException as exception:
-        logger.opt(exception=exception).critical(
-            "EXO terminated due to unhandled exception"
-        )
-        raise
-    finally:
-        logger.info("EXO Shutdown complete")
-        logger_cleanup()
+    anyio.run(node.run)
+    logger.info("EXO Shutdown complete")
+    logger_cleanup()
 
 
 class Args(CamelCaseModel):
@@ -290,8 +290,9 @@ class Args(CamelCaseModel):
     tb_only: bool = False
     no_worker: bool = False
     no_downloads: bool = False
-    offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
+    offline: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
+    trust_remote_code: bool = False
 
     @classmethod
     def parse(cls) -> Self:
@@ -341,9 +342,13 @@ class Args(CamelCaseModel):
         parser.add_argument(
             "--offline",
             action="store_true",
-            default=os.getenv("EXO_OFFLINE", "false").lower() == "true",
             help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
         )
+        parser.add_argument(
+            "--trust-remote-code",
+            action="store_true",
+            help="Allow models to execute custom code during tokenizer loading (security-sensitive, CLI-only)",
+        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/api.py

@@ -524,15 +524,15 @@ class API:
 
             if (
                 model_card.model_id,
-                instance.sharding(),
-                instance.instance_meta(),
+                sharding,
+                instance_meta,
                 len(placement_node_ids),
             ) not in seen:
                 previews.append(
                     PlacementPreview(
                         model_id=model_card.model_id,
-                        sharding=instance.sharding(),
-                        instance_meta=instance.instance_meta(),
+                        sharding=sharding,
+                        instance_meta=instance_meta,
                         instance=instance,
                         memory_delta_by_node=memory_delta_by_node or None,
                         error=None,
@@ -541,8 +541,8 @@ class API:
             seen.add(
                 (
                     model_card.model_id,
-                    instance.sharding(),
-                    instance.instance_meta(),
+                    sharding,
+                    instance_meta,
                     len(placement_node_ids),
                 )
             )

src/exo/routing/router.py

@@ -174,42 +174,28 @@ class Router:
         logger.info(f"Unsubscribed from {topic}")
 
     async def _networking_recv(self):
-        try:
-            while True:
-                topic, data = await self._net.gossipsub_recv()
-                logger.trace(f"Received message on {topic} with payload {data}")
-                if topic not in self.topic_routers:
-                    logger.warning(
-                        f"Received message on unknown or inactive topic {topic}"
-                    )
-                    continue
+        while True:
+            topic, data = await self._net.gossipsub_recv()
+            logger.trace(f"Received message on {topic} with payload {data}")
+            if topic not in self.topic_routers:
+                logger.warning(f"Received message on unknown or inactive topic {topic}")
+                continue
 
-                router = self.topic_routers[topic]
-                await router.publish_bytes(data)
-        except Exception as exception:
-            logger.opt(exception=exception).error(
-                "Gossipsub receive loop terminated unexpectedly"
-            )
-            raise
+            router = self.topic_routers[topic]
+            await router.publish_bytes(data)
 
     async def _networking_recv_connection_messages(self):
-        try:
-            while True:
-                update = await self._net.connection_update_recv()
-                message = ConnectionMessage.from_update(update)
-                logger.trace(
-                    f"Received message on connection_messages with payload {message}"
-                )
-                if CONNECTION_MESSAGES.topic in self.topic_routers:
-                    router = self.topic_routers[CONNECTION_MESSAGES.topic]
-                    assert router.topic.model_type == ConnectionMessage
-                    router = cast(TopicRouter[ConnectionMessage], router)
-                    await router.publish(message)
-        except Exception as exception:
-            logger.opt(exception=exception).error(
-                "Connection update receive loop terminated unexpectedly"
+        while True:
+            update = await self._net.connection_update_recv()
+            message = ConnectionMessage.from_update(update)
+            logger.trace(
+                f"Received message on connection_messages with payload {message}"
             )
-            raise
+            if CONNECTION_MESSAGES.topic in self.topic_routers:
+                router = self.topic_routers[CONNECTION_MESSAGES.topic]
+                assert router.topic.model_type == ConnectionMessage
+                router = cast(TopicRouter[ConnectionMessage], router)
+                await router.publish(message)
 
     async def _networking_publish(self):
         with self.networking_receiver as networked_items:

src/exo/shared/constants.py

@@ -78,6 +78,4 @@ EXO_ENABLE_IMAGE_MODELS = (
     os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true"
 )
 
-EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
-
 EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"

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

@@ -4,13 +4,7 @@ from pydantic import model_validator
 
 from exo.shared.models.model_cards import ModelTask
 from exo.shared.types.common import Host, Id, NodeId
-from exo.shared.types.worker.runners import RunnerId, ShardAssignments
-from exo.shared.types.worker.shards import (
-    PipelineShardMetadata,
-    Sharding,
-    ShardMetadata,
-    TensorShardMetadata,
-)
+from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
 
@@ -30,40 +24,16 @@ class BaseInstance(TaggedModel):
     def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
         return self.shard_assignments.runner_to_shard.get(runner_id, None)
 
-    @staticmethod
-    def instance_meta() -> InstanceMeta: ...
-
-    def sharding(self) -> Sharding:
-        if all(
-            isinstance(sm, PipelineShardMetadata)
-            for sm in self.shard_assignments.runner_to_shard.values()
-        ):
-            return Sharding.Pipeline
-        if all(
-            isinstance(sm, TensorShardMetadata)
-            for sm in self.shard_assignments.runner_to_shard.values()
-        ):
-            return Sharding.Tensor
-        raise ValueError("shard metadata malformed")
-
 
 class MlxRingInstance(BaseInstance):
     hosts_by_node: dict[NodeId, list[Host]]
     ephemeral_port: int
 
-    @staticmethod
-    def instance_meta() -> InstanceMeta:
-        return InstanceMeta.MlxRing
-
 
 class MlxJacclInstance(BaseInstance):
     jaccl_devices: list[list[str | None]]
     jaccl_coordinators: dict[NodeId, str]
 
-    @staticmethod
-    def instance_meta() -> InstanceMeta:
-        return InstanceMeta.MlxJaccl
-
 
 # TODO: Single node instance
 Instance = MlxRingInstance | MlxJacclInstance

src/exo/worker/engines/mlx/auto_parallel.py

@@ -128,11 +128,11 @@ class PipelineFirstLayer(CustomMlxLayer):
 
     def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
         if self.r != 0:
-            # We want to avoid GPU timeout errors by evalling the distributed operation
-            # so that it stays on CPU, which does not have a timeout.
-            mx.eval(x)
             x = mx.distributed.recv_like(x, (self.r - 1), group=self.group)
-            mx.eval(x)
+            if self.is_prefill:
+                # We want to avoid GPU timeout errors by evalling the distributed operation
+                # so that it stays on CPU, which does not have a timeout.
+                mx.eval(x)
         return self.original_layer(x, *args, **kwargs)
 
 
@@ -158,10 +158,6 @@ class PipelineLastLayer(CustomMlxLayer):
 
         output: mx.array = self.original_layer(x, *args, **kwargs)
 
-        # Eval layer output to materialize it before send — this splits the graph
-        # so the send is isolated and the receiving rank's recv can complete.
-        mx.eval(output)
-
         if self.r != self.s - 1:
             output = mx.distributed.send(
                 output, (self.r + 1) % self.s, group=self.group
@@ -171,15 +167,15 @@ class PipelineLastLayer(CustomMlxLayer):
                 # doesn't have .keys directly; access via first sub-cache.
                 _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
                 _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
-            mx.eval(output)
-            if cache is not None:
-                mx.eval(_cache.keys)  # type: ignore
+            if self.is_prefill:
+                mx.eval(output)
+                if cache is not None:
+                    mx.eval(_cache.keys)  # type: ignore
 
         if not self.is_prefill:
             output = mx.distributed.all_gather(output, group=self.group)[
                 -output.shape[0] :
             ]
-            mx.eval(output)
 
         return output
 

src/exo/worker/engines/mlx/constants.py

@@ -13,5 +13,6 @@ KV_CACHE_BITS: int | None = None
 
 DEFAULT_TOP_LOGPROBS: int = 5
 
-# TODO: We should really make this opt-in, but Kimi requires trust_remote_code=True
+# True for built-in models with known model cards; custom models added via API default to False
+# and can be overridden with the --trust-remote-code CLI flag.
 TRUST_REMOTE_CODE: bool = True

src/exo/worker/engines/mlx/generator/generate.py

@@ -1,4 +1,3 @@
-import math
 import time
 from copy import deepcopy
 from typing import Callable, Generator, cast, get_args
@@ -249,9 +248,6 @@ def extract_top_logprobs(
     for i in range(top_logprobs):
         token_id = int(top_indices[i].item())
         token_logprob = float(top_values[i].item())
-        if math.isnan(token_logprob):
-            continue
-
         # Decode token ID to string
         token_str = tokenizer.decode([token_id])
         # Get byte representation

src/exo/worker/engines/mlx/utils_mlx.py

@@ -2,7 +2,6 @@ import json
 import os
 import re
 import sys
-import tempfile
 import time
 from pathlib import Path
 from typing import Any, cast
@@ -99,13 +98,14 @@ def mlx_distributed_init(
     rank = bound_instance.bound_shard.device_rank
     logger.info(f"Starting initialization for rank {rank}")
 
-    with tempfile.TemporaryDirectory() as tmpdir:
-        coordination_file = str(
-            Path(tmpdir) / f"hosts_{bound_instance.instance.instance_id}_{rank}.json"
-        )
+    coordination_file = None
+    try:
         # TODO: singleton instances
         match bound_instance.instance:
             case MlxRingInstance(hosts_by_node=hosts_by_node, ephemeral_port=_):
+                coordination_file = (
+                    f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
+                )
                 hosts_for_node = hosts_by_node[bound_instance.bound_node_id]
                 hosts_json = HostList.from_hosts(hosts_for_node).model_dump_json()
 
@@ -128,6 +128,9 @@ def mlx_distributed_init(
                     jaccl_devices[i][i] is None for i in range(len(jaccl_devices))
                 )
                 # Use RDMA connectivity matrix
+                coordination_file = (
+                    f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
+                )
                 jaccl_devices_json = json.dumps(jaccl_devices)
 
                 with open(coordination_file, "w") as f:
@@ -147,6 +150,10 @@ def mlx_distributed_init(
         logger.info(f"Rank {rank} mlx distributed initialization complete")
 
         return group
+    finally:
+        with contextlib.suppress(FileNotFoundError):
+            if coordination_file:
+                os.remove(coordination_file)
 
 
 def initialize_mlx(
@@ -207,8 +214,6 @@ def load_mlx_items(
 
     set_wired_limit_for_model(get_weights_size(bound_instance.bound_shard))
 
-    mx.clear_cache()
-
     return cast(Model, model), tokenizer
 
 
@@ -286,10 +291,14 @@ def shard_and_load(
 
 def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
     """Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
+    trust_remote_code = (
+        shard_metadata.model_card.trust_remote_code
+        or os.environ.get("EXO_TRUST_REMOTE_CODE") == "1"
+    )
     return load_tokenizer_for_model_id(
         shard_metadata.model_card.model_id,
         model_path,
-        trust_remote_code=shard_metadata.model_card.trust_remote_code,
+        trust_remote_code=trust_remote_code,
     )
 
 

src/exo/worker/runner/llm_inference/runner.py

@@ -97,7 +97,6 @@ def main(
         bound_instance.bound_runner_id,
         bound_instance.bound_shard,
     )
-    model_id = shard_metadata.model_card.model_id
     device_rank = shard_metadata.device_rank
     logger.info("hello from the runner")
     if getattr(shard_metadata, "immediate_exception", False):
@@ -282,7 +281,7 @@ def main(
                                 ChunkGenerated(
                                     command_id=command_id,
                                     chunk=PrefillProgressChunk(
-                                        model=model_id,
+                                        model=shard_metadata.model_card.model_id,
                                         processed_tokens=processed,
                                         total_tokens=total,
                                     ),
@@ -326,10 +325,7 @@ def main(
                         # Model-specific output parsing for tool calls.
                         if isinstance(inference_model, GptOssModel):
                             mlx_generator = parse_gpt_oss(mlx_generator)
-                        elif (
-                            isinstance(inference_model, DeepseekV32Model)
-                            and "deepseek" in model_id.normalize().lower()
-                        ):
+                        elif isinstance(inference_model, DeepseekV32Model):
                             mlx_generator = parse_deepseek_v32(mlx_generator)
                         elif tool_parser:
                             mlx_generator = parse_tool_calls(mlx_generator, tool_parser)
@@ -359,7 +355,7 @@ def main(
                                                 command_id=command_id,
                                                 chunk=ErrorChunk(
                                                     error_message=response.text,
-                                                    model=model_id,
+                                                    model=shard_metadata.model_card.model_id,
                                                 ),
                                             )
                                         )
@@ -374,7 +370,7 @@ def main(
                                             ChunkGenerated(
                                                 command_id=command_id,
                                                 chunk=TokenChunk(
-                                                    model=model_id,
+                                                    model=shard_metadata.model_card.model_id,
                                                     text=response.text,
                                                     token_id=response.token,
                                                     usage=response.usage,
@@ -393,7 +389,7 @@ def main(
                                                 command_id=command_id,
                                                 chunk=ToolCallChunk(
                                                     tool_calls=response.tool_calls,
-                                                    model=model_id,
+                                                    model=shard_metadata.model_card.model_id,
                                                     usage=response.usage,
                                                     stats=response.stats,
                                                 ),
@@ -409,7 +405,7 @@ def main(
                                 ChunkGenerated(
                                     command_id=command_id,
                                     chunk=ErrorChunk(
-                                        model=model_id,
+                                        model=shard_metadata.model_card.model_id,
                                         finish_reason="error",
                                         error_message=str(e),
                                     ),
@@ -731,7 +727,7 @@ def parse_tool_calls(
     in_tool_call = False
     tool_call_text_parts: list[str] = []
     for response in responses:
-        if not in_tool_call and response.text.startswith(tool_parser.start_parsing):
+        if response.text.startswith(tool_parser.start_parsing):
             in_tool_call = True
 
         if in_tool_call:
@@ -769,9 +765,10 @@ def parse_tool_calls(
                 )
                 yield response
 
-        else:
-            # fallthrough
-            yield response
+            continue
+
+        # fallthrough
+        yield response
 
 
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"

src/exo/worker/runner/llm_inference/tool_parsers.py

@@ -58,16 +58,15 @@ def _flatten(p: dict[str, Any]) -> dict[str, str]:
     }
 
 
-def make_json_parser() -> ToolParser:
-    return ToolParser(
-        start_parsing="<tool_call>",
-        end_parsing="</tool_call>",
-        parse_tool_calls=_parse_json_calls,
-    )
+json_tool_parser = ToolParser(
+    start_parsing="<tool_call>",
+    end_parsing="</tool_call>",
+    parse_tool_calls=_parse_json_calls,
+)
 
 
 def infer_tool_parser(chat_template: str) -> ToolParser | None:
     """Attempt to auto-infer a tool parser from the chat template."""
     if "<tool_call>" in chat_template and "tool_call.name" in chat_template:
-        return make_json_parser()
+        return json_tool_parser
     return None

tests/auto_bench.sh

@@ -17,13 +17,6 @@ git branch -r --contains "$commit" | grep -qE '^\s*origin/' || {
   exit 1
 }
 hosts=("$@")
-
-for host; do
-  ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
-    "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
-done
-wait
-
 cleanup() {
   for host in "${hosts[@]}"; do
     ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" &
@@ -33,6 +26,11 @@ cleanup() {
 }
 trap 'cleanup' EXIT INT TERM
 
+for host; do
+  ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
+    "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
+done
+wait
 for host; do
   ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
     "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null &

uv.lock

@@ -363,7 +363,7 @@ wheels = [
 
 [[package]]
 name = "exo"
-version = "0.3.68"
+version = "0.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -378,7 +378,7 @@ dependencies = [
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260220+13998a05", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }, marker = "sys_platform == 'darwin'" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1025,7 +1025,7 @@ dependencies = [
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260220+13998a05", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1072,8 +1072,8 @@ cuda13 = [
 
 [[package]]
 name = "mlx"
-version = "0.30.7.dev20260225+257d5692"
-source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }
+version = "0.30.7.dev20260220+13998a05"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }
 resolution-markers = [
     "sys_platform == 'darwin'",
 ]
@@ -1108,7 +1108,7 @@ version = "0.30.7"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260220+13998a05", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },