Merge remote-tracking branch 'origin/main' into alexcheema/fix-graceful-process-shutdown

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

app/EXO/EXO/ExoProcessController.swift

@@ -118,11 +118,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() {

dashboard/src/lib/components/HuggingFaceResultItem.svelte

@@ -14,7 +14,6 @@
     isAdding: boolean;
     onAdd: () => void;
     onSelect: () => void;
-    downloadedOnNodes?: string[];
   };
 
   let {
@@ -23,7 +22,6 @@
     isAdding,
     onAdd,
     onSelect,
-    downloadedOnNodes = [],
   }: HuggingFaceResultItemProps = $props();
 
   function formatNumber(num: number): string {
@@ -47,28 +45,6 @@
       <span class="text-sm font-mono text-white truncate" title={model.id}
         >{modelName}</span
       >
-      {#if downloadedOnNodes.length > 0}
-        <span
-          class="flex-shrink-0"
-          title={`Downloaded on ${downloadedOnNodes.join(", ")}`}
-        >
-          <svg
-            class="w-4 h-4"
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            stroke-width="2"
-            stroke-linecap="round"
-            stroke-linejoin="round"
-          >
-            <path
-              class="text-white/40"
-              d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
-            />
-            <path class="text-green-400" d="m9 13 2 2 4-4" />
-          </svg>
-        </span>
-      {/if}
       {#if isAdded}
         <span
           class="px-1.5 py-0.5 text-[10px] font-mono bg-green-500/20 text-green-400 rounded"

dashboard/src/lib/components/ModelFilterPopover.svelte

@@ -5,7 +5,6 @@
   interface FilterState {
     capabilities: string[];
     sizeRange: { min: number; max: number } | null;
-    downloadedOnly: boolean;
   }
 
   type ModelFilterPopoverProps = {
@@ -149,36 +148,6 @@
       </div>
     </div>
 
-    <!-- Downloaded only -->
-    <div>
-      <h4 class="text-xs font-mono text-white/50 mb-2">Availability</h4>
-      <button
-        type="button"
-        class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
-          ? 'bg-green-500/20 text-green-400 border border-green-500/30'
-          : 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
-        onclick={() =>
-          onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
-      >
-        <svg
-          class="w-3.5 h-3.5 inline-block"
-          viewBox="0 0 24 24"
-          fill="none"
-          stroke="currentColor"
-          stroke-width="2"
-          stroke-linecap="round"
-          stroke-linejoin="round"
-        >
-          <path
-            class="text-white/40"
-            d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
-          />
-          <path class="text-green-400" d="m9 13 2 2 4-4" />
-        </svg>
-        <span class="ml-1">Downloaded</span>
-      </button>
-    </div>
-
     <!-- Size range -->
     <div>
       <h4 class="text-xs font-mono text-white/50 mb-2">Model Size</h4>

dashboard/src/lib/components/ModelPickerGroup.svelte

@@ -21,12 +21,6 @@
     hasMultipleVariants: boolean;
   }
 
-  type DownloadAvailability = {
-    available: boolean;
-    nodeNames: string[];
-    nodeIds: string[];
-  };
-
   type ModelPickerGroupProps = {
     group: ModelGroup;
     isExpanded: boolean;
@@ -37,7 +31,6 @@
     onSelectModel: (modelId: string) => void;
     onToggleFavorite: (baseModelId: string) => void;
     onShowInfo: (group: ModelGroup) => void;
-    downloadStatusMap?: Map<string, DownloadAvailability>;
   };
 
   let {
@@ -50,19 +43,8 @@
     onSelectModel,
     onToggleFavorite,
     onShowInfo,
-    downloadStatusMap,
   }: ModelPickerGroupProps = $props();
 
-  // Group-level download status: show if any variant is downloaded
-  const groupDownloadStatus = $derived.by(() => {
-    if (!downloadStatusMap || downloadStatusMap.size === 0) return undefined;
-    // Return the first available entry (prefer "available" ones)
-    for (const avail of downloadStatusMap.values()) {
-      if (avail.available) return avail;
-    }
-    return downloadStatusMap.values().next().value;
-  });
-
   // Format storage size
   function formatSize(mb: number | undefined): string {
     if (!mb) return "";
@@ -216,42 +198,10 @@
       </span>
     {/if}
 
-    <!-- Variant count with size range -->
+    <!-- Variant count -->
     {#if group.hasMultipleVariants}
-      {@const sizes = group.variants
-        .map((v) => v.storage_size_megabytes || 0)
-        .filter((s) => s > 0)
-        .sort((a, b) => a - b)}
       <span class="text-xs font-mono text-white/30 flex-shrink-0">
-        {group.variants.length} variants{#if sizes.length >= 2}{" "}({formatSize(
-            sizes[0],
-          )}-{formatSize(sizes[sizes.length - 1])}){/if}
-      </span>
-    {/if}
-
-    <!-- Download availability indicator -->
-    {#if groupDownloadStatus && groupDownloadStatus.nodeIds.length > 0}
-      <span
-        class="flex-shrink-0"
-        title={groupDownloadStatus.available
-          ? `Ready — downloaded on ${groupDownloadStatus.nodeNames.join(", ")}`
-          : `Downloaded on ${groupDownloadStatus.nodeNames.join(", ")} (may need more nodes)`}
-      >
-        <svg
-          class="w-4 h-4"
-          viewBox="0 0 24 24"
-          fill="none"
-          stroke="currentColor"
-          stroke-width="2"
-          stroke-linecap="round"
-          stroke-linejoin="round"
-        >
-          <path
-            class="text-white/40"
-            d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
-          />
-          <path class="text-green-400" d="m9 13 2 2 4-4" />
-        </svg>
+        {group.variants.length} variants
       </span>
     {/if}
 
@@ -355,33 +305,6 @@
             {formatSize(variant.storage_size_megabytes)}
           </span>
 
-          <!-- Download indicator for this variant -->
-          {#if downloadStatusMap?.get(variant.id)}
-            {@const variantDl = downloadStatusMap.get(variant.id)}
-            {#if variantDl}
-              <span
-                class="flex-shrink-0"
-                title={`Downloaded on ${variantDl.nodeNames.join(", ")}`}
-              >
-                <svg
-                  class="w-3.5 h-3.5"
-                  viewBox="0 0 24 24"
-                  fill="none"
-                  stroke="currentColor"
-                  stroke-width="2"
-                  stroke-linecap="round"
-                  stroke-linejoin="round"
-                >
-                  <path
-                    class="text-white/40"
-                    d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
-                  />
-                  <path class="text-green-400" d="m9 13 2 2 4-4" />
-                </svg>
-              </span>
-            {/if}
-          {/if}
-
           <!-- Check mark if selected -->
           {#if isSelected}
             <svg

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -5,7 +5,6 @@
   import ModelPickerGroup from "./ModelPickerGroup.svelte";
   import ModelFilterPopover from "./ModelFilterPopover.svelte";
   import HuggingFaceResultItem from "./HuggingFaceResultItem.svelte";
-  import { getNodesWithModelDownloaded } from "$lib/utils/downloads";
 
   interface ModelInfo {
     id: string;
@@ -34,7 +33,6 @@
   interface FilterState {
     capabilities: string[];
     sizeRange: { min: number; max: number } | null;
-    downloadedOnly: boolean;
   }
 
   interface HuggingFaceModel {
@@ -60,15 +58,6 @@
     onDeleteModel: (modelId: string) => Promise<void>;
     totalMemoryGB: number;
     usedMemoryGB: number;
-    downloadsData?: Record<string, unknown[]>;
-    topologyNodes?: Record<
-      string,
-      {
-        friendly_name?: string;
-        system_info?: { model_id?: string };
-        macmon_info?: { memory?: { ram_total?: number } };
-      }
-    >;
   };
 
   let {
@@ -85,8 +74,6 @@
     onDeleteModel,
     totalMemoryGB,
     usedMemoryGB,
-    downloadsData,
-    topologyNodes,
   }: ModelPickerModalProps = $props();
 
   // Local state
@@ -94,75 +81,9 @@
   let selectedFamily = $state<string | null>(null);
   let expandedGroups = $state<Set<string>>(new Set());
   let showFilters = $state(false);
-  let filters = $state<FilterState>({
-    capabilities: [],
-    sizeRange: null,
-    downloadedOnly: false,
-  });
+  let filters = $state<FilterState>({ capabilities: [], sizeRange: null });
   let infoGroup = $state<ModelGroup | null>(null);
 
-  // Download availability per model group
-  type DownloadAvailability = {
-    available: boolean;
-    nodeNames: string[];
-    nodeIds: string[];
-  };
-
-  function getNodeName(nodeId: string): string {
-    const node = topologyNodes?.[nodeId];
-    return (
-      node?.friendly_name || node?.system_info?.model_id || nodeId.slice(0, 8)
-    );
-  }
-
-  const modelDownloadAvailability = $derived.by(() => {
-    const result = new Map<string, DownloadAvailability>();
-    if (!downloadsData || !topologyNodes) return result;
-
-    for (const model of models) {
-      const nodeIds = getNodesWithModelDownloaded(downloadsData, model.id);
-      if (nodeIds.length === 0) continue;
-
-      // Sum total RAM across nodes that have the model
-      let totalRamBytes = 0;
-      for (const nodeId of nodeIds) {
-        const ramTotal = topologyNodes[nodeId]?.macmon_info?.memory?.ram_total;
-        if (typeof ramTotal === "number") totalRamBytes += ramTotal;
-      }
-
-      const modelSizeBytes = (model.storage_size_megabytes || 0) * 1024 * 1024;
-      result.set(model.id, {
-        available: modelSizeBytes > 0 && totalRamBytes >= modelSizeBytes,
-        nodeNames: nodeIds.map(getNodeName),
-        nodeIds,
-      });
-    }
-    return result;
-  });
-
-  // Aggregate download availability per group (available if ANY variant is available)
-  function getGroupDownloadAvailability(
-    group: ModelGroup,
-  ): DownloadAvailability | undefined {
-    for (const variant of group.variants) {
-      const avail = modelDownloadAvailability.get(variant.id);
-      if (avail && avail.nodeIds.length > 0) return avail;
-    }
-    return undefined;
-  }
-
-  // Get per-variant download map for a group
-  function getVariantDownloadMap(
-    group: ModelGroup,
-  ): Map<string, DownloadAvailability> {
-    const map = new Map<string, DownloadAvailability>();
-    for (const variant of group.variants) {
-      const avail = modelDownloadAvailability.get(variant.id);
-      if (avail && avail.nodeIds.length > 0) map.set(variant.id, avail);
-    }
-    return map;
-  }
-
   // HuggingFace Hub state
   let hfSearchQuery = $state("");
   let hfSearchResults = $state<HuggingFaceModel[]>([]);
@@ -174,12 +95,15 @@
   let manualModelId = $state("");
   let addModelError = $state<string | null>(null);
 
-  // Reset transient state when modal opens, but preserve tab selection
+  // Reset state when modal opens
   $effect(() => {
     if (isOpen) {
       searchQuery = "";
+      selectedFamily = null;
       expandedGroups = new Set();
       showFilters = false;
+      hfSearchQuery = "";
+      hfSearchResults = [];
       manualModelId = "";
       addModelError = null;
     }
@@ -415,16 +339,6 @@
       });
     }
 
-    // Filter to downloaded models only
-    if (filters.downloadedOnly) {
-      result = result.filter((g) =>
-        g.variants.some((v) => {
-          const avail = modelDownloadAvailability.get(v.id);
-          return avail && avail.nodeIds.length > 0;
-        }),
-      );
-    }
-
     // Sort: models that fit first, then by size (largest first)
     result.sort((a, b) => {
       const aFits = a.variants.some((v) => canModelFit(v.id));
@@ -471,13 +385,11 @@
   }
 
   function clearFilters() {
-    filters = { capabilities: [], sizeRange: null, downloadedOnly: false };
+    filters = { capabilities: [], sizeRange: null };
   }
 
   const hasActiveFilters = $derived(
-    filters.capabilities.length > 0 ||
-      filters.sizeRange !== null ||
-      filters.downloadedOnly,
+    filters.capabilities.length > 0 || filters.sizeRange !== null,
   );
 </script>
 
@@ -664,12 +576,6 @@
                     isAdding={addingModelId === model.id}
                     onAdd={() => handleAddModel(model.id)}
                     onSelect={() => handleSelectHfModel(model.id)}
-                    downloadedOnNodes={downloadsData
-                      ? getNodesWithModelDownloaded(
-                          downloadsData,
-                          model.id,
-                        ).map(getNodeName)
-                      : []}
                   />
                 {/each}
               {/if}
@@ -744,7 +650,6 @@
               onSelectModel={handleSelect}
               {onToggleFavorite}
               onShowInfo={(g) => (infoGroup = g)}
-              downloadStatusMap={getVariantDownloadMap(group)}
             />
           {/each}
         {/if}
@@ -762,11 +667,6 @@
             >{cap}</span
           >
         {/each}
-        {#if filters.downloadedOnly}
-          <span class="px-1.5 py-0.5 bg-green-500/20 text-green-400 rounded"
-            >Downloaded</span
-          >
-        {/if}
         {#if filters.sizeRange}
           <span class="px-1.5 py-0.5 bg-exo-yellow/20 text-exo-yellow rounded">
             {filters.sizeRange.min}GB - {filters.sizeRange.max}GB
@@ -842,40 +742,6 @@
             </div>
           </div>
         {/if}
-        {#if getGroupDownloadAvailability(infoGroup)?.nodeNames?.length}
-          {@const infoDownload = getGroupDownloadAvailability(infoGroup)}
-          {#if infoDownload}
-            <div class="mt-3 pt-3 border-t border-exo-yellow/10">
-              <div class="flex items-center gap-2 mb-1">
-                <svg
-                  class="w-3.5 h-3.5"
-                  viewBox="0 0 24 24"
-                  fill="none"
-                  stroke="currentColor"
-                  stroke-width="2"
-                  stroke-linecap="round"
-                  stroke-linejoin="round"
-                >
-                  <path
-                    class="text-white/40"
-                    d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
-                  />
-                  <path class="text-green-400" d="m9 13 2 2 4-4" />
-                </svg>
-                <span class="text-white/40">Downloaded on:</span>
-              </div>
-              <div class="flex flex-wrap gap-1 mt-1">
-                {#each infoDownload.nodeNames as nodeName}
-                  <span
-                    class="px-1.5 py-0.5 bg-green-500/10 text-green-400/80 border border-green-500/20 rounded text-[10px]"
-                  >
-                    {nodeName}
-                  </span>
-                {/each}
-              </div>
-            </div>
-          {/if}
-        {/if}
       </div>
     </div>
   {/if}

dashboard/src/lib/utils/downloads.ts

@@ -1,152 +0,0 @@
-/**
- * Shared utilities for parsing and querying download state.
- *
- * The download state from `/state` is shaped as:
- *   Record<NodeId, Array<TaggedDownloadEntry>>
- *
- * Each entry is a tagged union object like:
- *   { "DownloadCompleted": { shard_metadata: { "PipelineShardMetadata": { model_card: { model_id: "..." }, ... } }, ... } }
- */
-
-/** Unwrap one level of tagged-union envelope, returning [tag, payload]. */
-function unwrapTagged(
-  obj: Record<string, unknown>,
-): [string, Record<string, unknown>] | null {
-  const keys = Object.keys(obj);
-  if (keys.length !== 1) return null;
-  const tag = keys[0];
-  const payload = obj[tag];
-  if (!payload || typeof payload !== "object") return null;
-  return [tag, payload as Record<string, unknown>];
-}
-
-/** Extract the model ID string from a download entry's nested shard_metadata. */
-export function extractModelIdFromDownload(
-  downloadPayload: Record<string, unknown>,
-): string | null {
-  const shardMetadata =
-    downloadPayload.shard_metadata ?? downloadPayload.shardMetadata;
-  if (!shardMetadata || typeof shardMetadata !== "object") return null;
-
-  const unwrapped = unwrapTagged(shardMetadata as Record<string, unknown>);
-  if (!unwrapped) return null;
-  const [, shardData] = unwrapped;
-
-  const modelMeta = shardData.model_card ?? shardData.modelCard;
-  if (!modelMeta || typeof modelMeta !== "object") return null;
-
-  const meta = modelMeta as Record<string, unknown>;
-  return (meta.model_id as string) ?? (meta.modelId as string) ?? null;
-}
-
-/** Extract the shard_metadata object from a download entry payload. */
-export function extractShardMetadata(
-  downloadPayload: Record<string, unknown>,
-): Record<string, unknown> | null {
-  const shardMetadata =
-    downloadPayload.shard_metadata ?? downloadPayload.shardMetadata;
-  if (!shardMetadata || typeof shardMetadata !== "object") return null;
-  return shardMetadata as Record<string, unknown>;
-}
-
-/** Get the download tag (DownloadCompleted, DownloadOngoing, etc.) from a wrapped entry. */
-export function getDownloadTag(
-  entry: unknown,
-): [string, Record<string, unknown>] | null {
-  if (!entry || typeof entry !== "object") return null;
-  return unwrapTagged(entry as Record<string, unknown>);
-}
-
-/**
- * Iterate over all download entries for a given node, yielding [tag, payload, modelId].
- */
-function* iterNodeDownloads(
-  nodeDownloads: unknown[],
-): Generator<[string, Record<string, unknown>, string]> {
-  for (const entry of nodeDownloads) {
-    const tagged = getDownloadTag(entry);
-    if (!tagged) continue;
-    const [tag, payload] = tagged;
-    const modelId = extractModelIdFromDownload(payload);
-    if (!modelId) continue;
-    yield [tag, payload, modelId];
-  }
-}
-
-/** Check if a specific model is fully downloaded (DownloadCompleted) on a specific node. */
-export function isModelDownloadedOnNode(
-  downloadsData: Record<string, unknown[]>,
-  nodeId: string,
-  modelId: string,
-): boolean {
-  const nodeDownloads = downloadsData[nodeId];
-  if (!Array.isArray(nodeDownloads)) return false;
-
-  for (const [tag, , entryModelId] of iterNodeDownloads(nodeDownloads)) {
-    if (tag === "DownloadCompleted" && entryModelId === modelId) return true;
-  }
-  return false;
-}
-
-/** Get all node IDs where a model is fully downloaded (DownloadCompleted). */
-export function getNodesWithModelDownloaded(
-  downloadsData: Record<string, unknown[]>,
-  modelId: string,
-): string[] {
-  const result: string[] = [];
-  for (const nodeId of Object.keys(downloadsData)) {
-    if (isModelDownloadedOnNode(downloadsData, nodeId, modelId)) {
-      result.push(nodeId);
-    }
-  }
-  return result;
-}
-
-/**
- * Find shard metadata for a model from any download entry across all nodes.
- * Returns the first match found (completed entries are preferred).
- */
-export function getShardMetadataForModel(
-  downloadsData: Record<string, unknown[]>,
-  modelId: string,
-): Record<string, unknown> | null {
-  let fallback: Record<string, unknown> | null = null;
-
-  for (const nodeDownloads of Object.values(downloadsData)) {
-    if (!Array.isArray(nodeDownloads)) continue;
-
-    for (const [tag, payload, entryModelId] of iterNodeDownloads(
-      nodeDownloads,
-    )) {
-      if (entryModelId !== modelId) continue;
-      const shard = extractShardMetadata(payload);
-      if (!shard) continue;
-
-      if (tag === "DownloadCompleted") return shard;
-      if (!fallback) fallback = shard;
-    }
-  }
-  return fallback;
-}
-
-/**
- * Get the download status tag for a specific model on a specific node.
- * Returns the "best" status: DownloadCompleted > DownloadOngoing > others.
- */
-export function getModelDownloadStatus(
-  downloadsData: Record<string, unknown[]>,
-  nodeId: string,
-  modelId: string,
-): string | null {
-  const nodeDownloads = downloadsData[nodeId];
-  if (!Array.isArray(nodeDownloads)) return null;
-
-  let best: string | null = null;
-  for (const [tag, , entryModelId] of iterNodeDownloads(nodeDownloads)) {
-    if (entryModelId !== modelId) continue;
-    if (tag === "DownloadCompleted") return tag;
-    if (tag === "DownloadOngoing") best = tag;
-    else if (!best) best = tag;
-  }
-  return best;
-}

dashboard/src/routes/+page.svelte

@@ -3264,6 +3264,4 @@
   onDeleteModel={deleteCustomModel}
   totalMemoryGB={clusterMemory().total / (1024 * 1024 * 1024)}
   usedMemoryGB={clusterMemory().used / (1024 * 1024 * 1024)}
-  {downloadsData}
-  topologyNodes={data?.nodes}
 />

flake.nix

@@ -118,10 +118,9 @@
             {
               metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
               mlx = pkgs.callPackage ./nix/mlx.nix {
-                inherit (self'.packages) metal-toolchain;
+                metal-toolchain = self'.packages.metal-toolchain;
                 inherit uvLockMlxVersion;
               };
-              default = self'.packages.exo;
             }
           );
 

resources/inference_model_cards/mlx-community--Qwen3-Coder-Next-4bit.toml

@@ -1,8 +0,0 @@
-model_id = "mlx-community/Qwen3-Coder-Next-4bit"
-n_layers = 48
-hidden_size = 2048
-supports_tensor = true
-tasks = ["TextGeneration"]
-
-[storage_size]
-in_bytes = 45644286500

resources/inference_model_cards/mlx-community--Qwen3-Coder-Next-5bit.toml

@@ -1,8 +0,0 @@
-model_id = "mlx-community/Qwen3-Coder-Next-5bit"
-n_layers = 48
-hidden_size = 2048
-supports_tensor = true
-tasks = ["TextGeneration"]
-
-[storage_size]
-in_bytes = 57657697020

resources/inference_model_cards/mlx-community--Qwen3-Coder-Next-6bit.toml

@@ -1,8 +0,0 @@
-model_id = "mlx-community/Qwen3-Coder-Next-6bit"
-n_layers = 48
-hidden_size = 2048
-supports_tensor = true
-tasks = ["TextGeneration"]
-
-[storage_size]
-in_bytes = 68899327465

resources/inference_model_cards/mlx-community--Qwen3-Coder-Next-8bit.toml

@@ -1,8 +0,0 @@
-model_id = "mlx-community/Qwen3-Coder-Next-8bit"
-n_layers = 48
-hidden_size = 2048
-supports_tensor = true
-tasks = ["TextGeneration"]
-
-[storage_size]
-in_bytes = 89357758772

resources/inference_model_cards/mlx-community--Qwen3-Coder-Next-bf16.toml

@@ -1,8 +0,0 @@
-model_id = "mlx-community/Qwen3-Coder-Next-bf16"
-n_layers = 48
-hidden_size = 2048
-supports_tensor = true
-tasks = ["TextGeneration"]
-
-[storage_size]
-in_bytes = 157548627945

src/exo/download/coordinator.py

@@ -53,10 +53,11 @@ class DownloadCoordinator:
     # Internal event channel for forwarding (initialized in __post_init__)
     event_sender: Sender[Event] = field(init=False)
     event_receiver: Receiver[Event] = field(init=False)
-    _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
+    _tg: TaskGroup = field(init=False)
 
     def __post_init__(self) -> None:
         self.event_sender, self.event_receiver = channel[Event]()
+        self._tg = anyio.create_task_group()
 
     async def run(self) -> None:
         logger.info("Starting DownloadCoordinator")

src/exo/main.py

@@ -27,6 +27,7 @@ from exo.utils.pydantic_ext import CamelCaseModel
 from exo.worker.main import Worker
 
 
+# I marked this as a dataclass as I want trivial constructors.
 @dataclass
 class Node:
     router: Router
@@ -135,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:
@@ -146,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

@@ -1320,40 +1320,29 @@ class API:
         ]
 
     async def run(self):
-        shutdown_ev = anyio.Event()
-
-        try:
-            async with create_task_group() as tg:
-                self._tg = tg
-                logger.info("Starting API")
-                tg.start_soon(self._apply_state)
-                tg.start_soon(self._pause_on_new_election)
-                tg.start_soon(self._cleanup_expired_images)
-                print_startup_banner(self.port)
-                tg.start_soon(self.run_api, shutdown_ev)
-                try:
-                    await anyio.sleep_forever()
-                finally:
-                    with anyio.CancelScope(shield=True):
-                        shutdown_ev.set()
-        finally:
-            self.command_sender.close()
-            self.global_event_receiver.close()
-
-    async def run_api(self, ev: anyio.Event):
         cfg = Config()
-        cfg.bind = [f"0.0.0.0:{self.port}"]
+        cfg.bind = f"0.0.0.0:{self.port}"
         # nb: shared.logging needs updating if any of this changes
         cfg.accesslog = None
         cfg.errorlog = "-"
         cfg.logger_class = InterceptLogger
-        with anyio.CancelScope(shield=True):
+
+        async with create_task_group() as tg:
+            self._tg = tg
+            logger.info("Starting API")
+            tg.start_soon(self._apply_state)
+            tg.start_soon(self._pause_on_new_election)
+            tg.start_soon(self._cleanup_expired_images)
+            print_startup_banner(self.port)
             await serve(
                 cast(ASGIFramework, self.app),
                 cfg,
-                shutdown_trigger=ev.wait,
+                shutdown_trigger=lambda: anyio.sleep_forever(),
             )
 
+        self.command_sender.close()
+        self.global_event_receiver.close()
+
     async def _apply_state(self):
         with self.global_event_receiver as events:
             async for f_event in events:

src/exo/master/main.py

@@ -96,18 +96,16 @@ class Master:
     async def run(self):
         logger.info("Starting Master")
 
-        try:
-            async with self._tg as tg:
-                tg.start_soon(self._event_processor)
-                tg.start_soon(self._command_processor)
-                tg.start_soon(self._loopback_processor)
-                tg.start_soon(self._plan)
-        finally:
-            self.global_event_sender.close()
-            self.local_event_receiver.close()
-            self.command_receiver.close()
-            self._loopback_event_sender.close()
-            self._loopback_event_receiver.close()
+        async with self._tg as tg:
+            tg.start_soon(self._event_processor)
+            tg.start_soon(self._command_processor)
+            tg.start_soon(self._loopback_processor)
+            tg.start_soon(self._plan)
+        self.global_event_sender.close()
+        self.local_event_receiver.close()
+        self.command_receiver.close()
+        self._loopback_event_sender.close()
+        self._loopback_event_receiver.close()
 
     async def shutdown(self):
         logger.info("Stopping Master")

src/exo/routing/router.py

@@ -9,7 +9,6 @@ from anyio import (
     BrokenResourceError,
     ClosedResourceError,
     create_task_group,
-    move_on_after,
     sleep_forever,
 )
 from anyio.abc import TaskGroup
@@ -147,21 +146,18 @@ class Router:
 
     async def run(self):
         logger.debug("Starting Router")
-        try:
-            async with create_task_group() as tg:
-                self._tg = tg
-                for topic in self.topic_routers:
-                    router = self.topic_routers[topic]
-                    tg.start_soon(router.run)
-                tg.start_soon(self._networking_recv)
-                tg.start_soon(self._networking_recv_connection_messages)
-                tg.start_soon(self._networking_publish)
-                # Router only shuts down if you cancel it.
-                await sleep_forever()
-        finally:
-            with move_on_after(1, shield=True):
-                for topic in self.topic_routers:
-                    await self._networking_unsubscribe(str(topic))
+        async with create_task_group() as tg:
+            self._tg = tg
+            for topic in self.topic_routers:
+                router = self.topic_routers[topic]
+                tg.start_soon(router.run)
+            tg.start_soon(self._networking_recv)
+            tg.start_soon(self._networking_recv_connection_messages)
+            tg.start_soon(self._networking_publish)
+            # Router only shuts down if you cancel it.
+            await sleep_forever()
+        for topic in self.topic_routers:
+            await self._networking_unsubscribe(str(topic))
 
     async def shutdown(self):
         logger.debug("Shutting down Router")
@@ -170,12 +166,12 @@ class Router:
         self._tg.cancel_scope.cancel()
 
     async def _networking_subscribe(self, topic: str):
+        logger.info(f"Subscribing to {topic}")
         await self._net.gossipsub_subscribe(topic)
-        logger.info(f"Subscribed to {topic}")
 
     async def _networking_unsubscribe(self, topic: str):
+        logger.info(f"Unsubscribing from {topic}")
         await self._net.gossipsub_unsubscribe(topic)
-        logger.info(f"Unsubscribed from {topic}")
 
     async def _networking_recv(self):
         while True:

src/exo/shared/election.py

@@ -86,29 +86,28 @@ class Election:
 
     async def run(self):
         logger.info("Starting Election")
-        try:
-            async with create_task_group() as tg:
-                self._tg = tg
-                tg.start_soon(self._election_receiver)
-                tg.start_soon(self._connection_receiver)
-                tg.start_soon(self._command_counter)
+        async with create_task_group() as tg:
+            self._tg = tg
+            tg.start_soon(self._election_receiver)
+            tg.start_soon(self._connection_receiver)
+            tg.start_soon(self._command_counter)
 
-                # And start an election immediately, that instantly resolves
-                candidates: list[ElectionMessage] = []
-                logger.debug("Starting initial campaign")
-                self._candidates = candidates
-                await self._campaign(candidates, campaign_timeout=0.0)
-                logger.debug("Initial campaign finished")
-        finally:
-            # Cancel and wait for the last election to end
-            if self._campaign_cancel_scope is not None:
-                logger.debug("Cancelling campaign")
-                self._campaign_cancel_scope.cancel()
-            if self._campaign_done is not None:
-                logger.debug("Waiting for campaign to finish")
-                await self._campaign_done.wait()
-            logger.debug("Campaign cancelled and finished")
-            logger.info("Election shutdown")
+            # And start an election immediately, that instantly resolves
+            candidates: list[ElectionMessage] = []
+            logger.debug("Starting initial campaign")
+            self._candidates = candidates
+            await self._campaign(candidates, campaign_timeout=0.0)
+            logger.debug("Initial campaign finished")
+
+        # Cancel and wait for the last election to end
+        if self._campaign_cancel_scope is not None:
+            logger.debug("Cancelling campaign")
+            self._campaign_cancel_scope.cancel()
+        if self._campaign_done is not None:
+            logger.debug("Waiting for campaign to finish")
+            await self._campaign_done.wait()
+        logger.debug("Campaign cancelled and finished")
+        logger.info("Election finished")
 
     async def elect(self, em: ElectionMessage) -> None:
         logger.debug(f"Electing: {em}")

src/exo/utils/channels.py

@@ -194,10 +194,9 @@ class MpReceiver[T]:
                 raise EndOfStream from None
             return item
 
+    # nb: this function will not cancel particularly well
     async def receive_async(self) -> T:
-        return await to_thread.run_sync(
-            self.receive, limiter=CapacityLimiter(1), abandon_on_cancel=True
-        )
+        return await to_thread.run_sync(self.receive, limiter=CapacityLimiter(1))
 
     def close(self) -> None:
         if not self._state.closed.is_set():

src/exo/worker/main.py

@@ -98,23 +98,21 @@ class Worker:
         info_send, info_recv = channel[GatheredInfo]()
         info_gatherer: InfoGatherer = InfoGatherer(info_send)
 
-        try:
-            async with self._tg as tg:
-                tg.start_soon(info_gatherer.run)
-                tg.start_soon(self._forward_info, info_recv)
-                tg.start_soon(self.plan_step)
-                tg.start_soon(self._resend_out_for_delivery)
-                tg.start_soon(self._event_applier)
-                tg.start_soon(self._forward_events)
-                tg.start_soon(self._poll_connection_updates)
-        finally:
-            # Actual shutdown code - waits for all tasks to complete before executing.
-            logger.info("Stopping 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 self._tg as tg:
+            tg.start_soon(info_gatherer.run)
+            tg.start_soon(self._forward_info, info_recv)
+            tg.start_soon(self.plan_step)
+            tg.start_soon(self._resend_out_for_delivery)
+            tg.start_soon(self._event_applier)
+            tg.start_soon(self._forward_events)
+            tg.start_soon(self._poll_connection_updates)
+
+        # Actual shutdown code - waits for all tasks to complete before executing.
+        self.local_event_sender.close()
+        self.command_sender.close()
+        self.download_command_sender.close()
+        for runner in self.runners.values():
+            runner.shutdown()
 
     async def _forward_info(self, recv: Receiver[GatheredInfo]):
         with recv as info_stream:

src/exo/worker/runner/runner_supervisor.py

@@ -8,8 +8,10 @@ import anyio
 from anyio import (
     BrokenResourceError,
     ClosedResourceError,
+    create_task_group,
     to_thread,
 )
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.shared.types.events import (
@@ -47,6 +49,7 @@ class RunnerSupervisor:
     _ev_recv: MpReceiver[Event]
     _task_sender: MpSender[Task]
     _event_sender: Sender[Event]
+    _tg: TaskGroup | None = field(default=None, 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)
@@ -90,29 +93,28 @@ class RunnerSupervisor:
 
     async def run(self):
         self.runner_process.start()
-        await self._forward_events()
+        async with create_task_group() as tg:
+            self._tg = tg
+            tg.start_soon(self._forward_events)
 
-    def shutdown(self):
-        logger.info("Runner supervisor shutting down")
         self._ev_recv.close()
         self._task_sender.close()
         self._event_sender.close()
-        self.runner_process.join(1)
+        await to_thread.run_sync(self.runner_process.join, 30)
         if not self.runner_process.is_alive():
-            logger.info("Runner process succesfully terminated")
             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()
-        self.runner_process.join(1)
+        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 SIGTERM, killing")
         self.runner_process.kill()
 
-        self.runner_process.join(1)
+        await to_thread.run_sync(self.runner_process.join, 5)
         if not self.runner_process.is_alive():
             return
 
@@ -120,6 +122,10 @@ class RunnerSupervisor:
             "Runner process didn't respond to SIGKILL. System resources may have leaked"
         )
 
+    def shutdown(self):
+        assert self._tg
+        self._tg.cancel_scope.cancel()
+
     async def start_task(self, task: Task):
         if task.task_id in self.pending:
             logger.warning(

tests/run_exo_on.sh

@@ -22,7 +22,7 @@ echo "Deploying $commit to $# hosts..."
 hosts=("$@")
 cleanup() {
   for host in "${hosts[@]}"; do
-    ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" &
+    ssh -T -o BatchMode=yes "$host@$host" "pkill -SIGINT -of exo-env" &
   done
   wait
   jobs -pr | xargs -r kill 2>/dev/null || true
@@ -34,13 +34,21 @@ reset=$'\e[0m'
 i=0
 for host; do
   colour=${colours[i++ % 4]}
-  ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
-    "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" |&
-    awk -v p="${colour}[${host}]${reset}" '{ print p $0; fflush() }' &
+  {
+    ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
+      "/nix/var/nix/profiles/default/bin/nix shell nixpkgs#git -c bash -s -- '$commit'" \
+      2>&1 | awk -v p="${colour}[${host}]${reset}" '{ print p $0; fflush() }' &
+  } <<'EOF'
+        set -euo pipefail
+        cd exo
+        git fetch -q origin
+        git checkout -q "$1"
+        EXO_LIBP2P_NAMESPACE="$1" /nix/var/nix/profiles/default/bin/nix run .#exo
+EOF
 done
 
 for host; do
   echo "Waiting for $host..."
-  until curl -sf "http://$host:52415/models" &>/dev/null; do sleep 1; done
+  until curl -sf "http://$host:52415/models"; do sleep 1; done
 done
 wait