mlx: build with Nix

In order to make testing and deployment simpler and more reproducible,
we want to provide a Nix derivation for our macOS .app build. We already
build the Rust and dashboard with Nix, but so far the Python has been
blocked because we haven't had an MLX build.

This change adds a Metal compiler derivation that uses `requireFile` to
be provided a NAR of the unfree macOS Metal compiler. It is documented
how to get this file, but effectively you have to trigger the download,
mount the DMG, and NAR the result. Once this is added to the store by
hash we can build MLX using it. The MLX build itself is quite self
explanatory.

Test plan:
- CI. We follow the instructions to grab the Metal compiler. Once this
  is in Cachix we should really never do this again, and I can pin the
  path too to ensure it doesn't leave.
- MLX tests run as part of the MLX derivation's build. They pass.
- `NIXPKGS_ALLOW_UNFREE=1 nix build .#mlx.passthru.tests.mlxTest --impure --option sandbox false`

.github/workflows/pipeline.yml

@@ -123,6 +123,63 @@ jobs:
           name: exo
           authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
 
+      - name: Build Metal packages (macOS only)
+        if: runner.os == 'macOS'
+        run: |
+          # Try to build metal-toolchain first (may succeed via cachix cache hit)
+          if nix build .#metal-toolchain 2>/dev/null; then
+            echo "metal-toolchain built successfully (likely cache hit)"
+          else
+            echo "metal-toolchain build failed, extracting from Xcode..."
+
+            NAR_HASH="sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw="
+            NAR_NAME="metal-toolchain-17C48.nar"
+
+            # Use RUNNER_TEMP to avoid /tmp symlink issues on macOS
+            WORK_DIR="${RUNNER_TEMP}/metal-work"
+            mkdir -p "$WORK_DIR"
+
+            # Download the Metal toolchain component
+            xcodebuild -downloadComponent MetalToolchain
+
+            # Find and mount the DMG
+            DMG_PATH=$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' 2>/dev/null | head -1)
+            if [ -z "$DMG_PATH" ]; then
+              echo "Error: Could not find Metal toolchain DMG"
+              exit 1
+            fi
+
+            echo "Found DMG at: $DMG_PATH"
+            hdiutil attach "$DMG_PATH" -mountpoint "${WORK_DIR}/metal-dmg"
+
+            # Copy the toolchain
+            cp -R "${WORK_DIR}/metal-dmg/Metal.xctoolchain" "${WORK_DIR}/metal-export"
+            hdiutil detach "${WORK_DIR}/metal-dmg"
+
+            # Create NAR and add to store
+            nix nar pack "${WORK_DIR}/metal-export" > "${WORK_DIR}/${NAR_NAME}"
+            STORE_PATH=$(nix store add --mode flat "${WORK_DIR}/${NAR_NAME}")
+            echo "Added NAR to store: $STORE_PATH"
+
+            # Verify the hash matches
+            ACTUAL_HASH=$(nix hash file "${WORK_DIR}/${NAR_NAME}")
+            if [ "$ACTUAL_HASH" != "$NAR_HASH" ]; then
+              echo "Warning: NAR hash mismatch!"
+              echo "Expected: $NAR_HASH"
+              echo "Actual:   $ACTUAL_HASH"
+              echo "The metal-toolchain.nix may need updating"
+            fi
+
+            # Clean up
+            rm -rf "$WORK_DIR"
+
+            # Retry the build now that NAR is in store
+            nix build .#metal-toolchain
+          fi
+
+          # Build mlx (depends on metal-toolchain)
+          nix build .#mlx
+
       - name: Build all Nix outputs
         run: |
           nix flake show --json | jq -r '

app/EXO/EXO/EXOApp.swift

@@ -45,8 +45,8 @@ struct EXOApp: App {
         let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
         _thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
         enableLaunchAtLoginIfNeeded()
-        // Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
-        NetworkSetupHelper.promptAndInstallIfNeeded()
+        // Remove old LaunchDaemon components if they exist (from previous versions)
+        cleanupLegacyNetworkSetup()
         // Check local network access periodically (warning disappears when user grants permission)
         localNetwork.startPeriodicChecking(interval: 10)
         controller.scheduleLaunch(after: 15)
@@ -136,6 +136,36 @@ struct EXOApp: App {
         }
     }
 
+    private func cleanupLegacyNetworkSetup() {
+        guard NetworkSetupHelper.hasInstalledComponents() else { return }
+        // Dispatch async to ensure app is ready before showing alert
+        DispatchQueue.main.async {
+            let alert = NSAlert()
+            alert.messageText = "EXO Network Configuration"
+            alert.informativeText =
+                "EXO needs to configure local network discovery on your device. This requires granting permission once."
+            alert.alertStyle = .informational
+            alert.addButton(withTitle: "Continue")
+            alert.addButton(withTitle: "Later")
+
+            let response = alert.runModal()
+            guard response == .alertFirstButtonReturn else {
+                Logger().info("User deferred legacy network setup cleanup")
+                return
+            }
+
+            do {
+                try NetworkSetupHelper.uninstall()
+                Logger().info("Cleaned up legacy network setup components")
+            } catch {
+                // Non-fatal: user may have cancelled admin prompt or cleanup may have
+                // partially succeeded. The app will continue normally.
+                Logger().warning(
+                    "Could not clean up legacy network setup (non-fatal): \(error.localizedDescription)"
+                )
+            }
+        }
+    }
 }
 
 /// Helper for managing EXO's launch-at-login registration

app/EXO/EXO/Services/NetworkSetupHelper.swift

@@ -11,97 +11,6 @@ enum NetworkSetupHelper {
     private static let legacyScriptDestination =
         "/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
     private static let plistDestination = "/Library/LaunchDaemons/io.exo.networksetup.plist"
-    private static let requiredStartInterval: Int = 1786
-
-    private static let setupScript = """
-        #!/usr/bin/env bash
-
-        set -euo pipefail
-
-        PREFS="/Library/Preferences/SystemConfiguration/preferences.plist"
-
-        # Remove bridge0 interface
-        ifconfig bridge0 &>/dev/null && {
-          ifconfig bridge0 | grep -q 'member' && {
-            ifconfig bridge0 | awk '/member/ {print $2}' | xargs -n1 ifconfig bridge0 deletem 2>/dev/null || true
-          }
-          ifconfig bridge0 destroy 2>/dev/null || true
-        }
-
-        # Remove Thunderbolt Bridge from VirtualNetworkInterfaces in preferences.plist
-        /usr/libexec/PlistBuddy -c "Delete :VirtualNetworkInterfaces:Bridge:bridge0" "$PREFS" 2>/dev/null || true
-
-        networksetup -listlocations | grep -q exo || {
-          networksetup -createlocation exo
-        }
-
-        networksetup -switchtolocation exo
-        networksetup -listallhardwareports \\
-          | awk -F': ' '/Hardware Port: / {print $2}' \\
-          | while IFS=":" read -r name; do
-              case "$name" in
-                "Ethernet Adapter"*)
-                        ;;
-                "Thunderbolt Bridge")
-                        ;;
-                "Thunderbolt "*)
-                  networksetup -listallnetworkservices \\
-                    | grep -q "EXO $name" \\
-                      || networksetup -createnetworkservice "EXO $name" "$name" 2>/dev/null \\
-                      || continue
-                  networksetup -setdhcp "EXO $name"
-                        ;;
-                *)
-                  networksetup -listallnetworkservices \\
-                    | grep -q "$name" \\
-                      || networksetup -createnetworkservice "$name" "$name" 2>/dev/null \\
-                      || continue
-                        ;;
-              esac
-            done
-
-        networksetup -listnetworkservices | grep -q "Thunderbolt Bridge" && {
-          networksetup -setnetworkserviceenabled "Thunderbolt Bridge" off
-        } || true
-        """
-
-    /// Prompts user and installs the LaunchDaemon if not already installed.
-    /// Shows an alert explaining what will be installed before requesting admin privileges.
-    static func promptAndInstallIfNeeded() {
-        // Use .utility priority to match NSAppleScript's internal QoS and avoid priority inversion
-        Task.detached(priority: .utility) {
-            // If already correctly installed, skip
-            if daemonAlreadyInstalled() {
-                return
-            }
-
-            // Show alert on main thread
-            let shouldInstall = await MainActor.run {
-                let alert = NSAlert()
-                alert.messageText = "EXO Network Configuration"
-                alert.informativeText =
-                    "EXO needs to install a system service to automatically disable Thunderbolt Bridge on startup. This prevents network loops when connecting multiple Macs via Thunderbolt.\n\nYou will be prompted for your administrator password."
-                alert.alertStyle = .informational
-                alert.addButton(withTitle: "Install")
-                alert.addButton(withTitle: "Not Now")
-                return alert.runModal() == .alertFirstButtonReturn
-            }
-
-            guard shouldInstall else {
-                logger.info("User deferred network setup daemon installation")
-                return
-            }
-
-            do {
-                try installLaunchDaemon()
-                logger.info("Network setup launch daemon installed and started")
-            } catch {
-                logger.error(
-                    "Network setup launch daemon failed: \(error.localizedDescription, privacy: .public)"
-                )
-            }
-        }
-    }
 
     /// Removes all EXO network setup components from the system.
     /// This includes the LaunchDaemon, scripts, logs, and network location.
@@ -121,100 +30,6 @@ enum NetworkSetupHelper {
         return scriptExists || legacyScriptExists || plistExists
     }
 
-    private static func daemonAlreadyInstalled() -> Bool {
-        let manager = FileManager.default
-        let scriptExists = manager.fileExists(atPath: scriptDestination)
-        let plistExists = manager.fileExists(atPath: plistDestination)
-        guard scriptExists, plistExists else { return false }
-        guard
-            let installedScript = try? String(contentsOfFile: scriptDestination, encoding: .utf8),
-            installedScript.trimmingCharacters(in: .whitespacesAndNewlines)
-                == setupScript.trimmingCharacters(in: .whitespacesAndNewlines)
-        else {
-            return false
-        }
-        guard
-            let data = try? Data(contentsOf: URL(fileURLWithPath: plistDestination)),
-            let plist = try? PropertyListSerialization.propertyList(
-                from: data, options: [], format: nil) as? [String: Any]
-        else {
-            return false
-        }
-        guard
-            let interval = plist["StartInterval"] as? Int,
-            interval == requiredStartInterval
-        else {
-            return false
-        }
-        if let programArgs = plist["ProgramArguments"] as? [String],
-            programArgs.contains(scriptDestination) == false
-        {
-            return false
-        }
-        return true
-    }
-
-    private static func installLaunchDaemon() throws {
-        let installerScript = makeInstallerScript()
-        try runShellAsAdmin(installerScript)
-    }
-
-    private static func makeInstallerScript() -> String {
-        """
-        set -euo pipefail
-
-        LABEL="\(daemonLabel)"
-        SCRIPT_DEST="\(scriptDestination)"
-        LEGACY_SCRIPT_DEST="\(legacyScriptDestination)"
-        PLIST_DEST="\(plistDestination)"
-        LOG_OUT="/var/log/\(daemonLabel).log"
-        LOG_ERR="/var/log/\(daemonLabel).err.log"
-
-        # First, completely remove any existing installation
-        launchctl bootout system/"$LABEL" 2>/dev/null || true
-        rm -f "$PLIST_DEST"
-        rm -f "$SCRIPT_DEST"
-        rm -f "$LEGACY_SCRIPT_DEST"
-        rm -f "$LOG_OUT" "$LOG_ERR"
-
-        # Install fresh
-        mkdir -p "$(dirname "$SCRIPT_DEST")"
-
-        cat > "$SCRIPT_DEST" <<'EOF_SCRIPT'
-        \(setupScript)
-        EOF_SCRIPT
-        chmod 755 "$SCRIPT_DEST"
-
-        cat > "$PLIST_DEST" <<'EOF_PLIST'
-        <?xml version="1.0" encoding="UTF-8"?>
-        <!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
-        <plist version="1.0">
-        <dict>
-          <key>Label</key>
-          <string>\(daemonLabel)</string>
-          <key>ProgramArguments</key>
-          <array>
-            <string>/bin/bash</string>
-            <string>\(scriptDestination)</string>
-          </array>
-          <key>StartInterval</key>
-          <integer>\(requiredStartInterval)</integer>
-          <key>RunAtLoad</key>
-          <true/>
-          <key>StandardOutPath</key>
-          <string>/var/log/\(daemonLabel).log</string>
-          <key>StandardErrorPath</key>
-          <string>/var/log/\(daemonLabel).err.log</string>
-        </dict>
-        </plist>
-        EOF_PLIST
-
-        launchctl bootstrap system "$PLIST_DEST"
-        launchctl enable system/"$LABEL"
-        launchctl kickstart -k system/"$LABEL"
-        """
-    }
-
     private static func makeUninstallScript() -> String {
         """
         set -euo pipefail

dashboard/parts.nix

@@ -3,61 +3,6 @@
   perSystem =
     { pkgs, lib, ... }:
     let
-      # Filter source to ONLY include package.json and package-lock.json
-      # This ensures prettier-svelte only rebuilds when lockfiles change
-      dashboardLockfileSrc = lib.cleanSourceWith {
-        src = inputs.self;
-        filter =
-          path: type:
-          let
-            baseName = builtins.baseNameOf path;
-            isDashboardDir = baseName == "dashboard" && type == "directory";
-            isPackageFile =
-              (lib.hasInfix "/dashboard/" path || lib.hasSuffix "/dashboard" (builtins.dirOf path))
-              && (baseName == "package.json" || baseName == "package-lock.json");
-          in
-          isDashboardDir || isPackageFile;
-      };
-
-      # Stub source with lockfiles and minimal files for build to succeed
-      # This allows prettier-svelte to avoid rebuilding when dashboard source changes
-      dashboardStubSrc = pkgs.runCommand "dashboard-stub-src" { } ''
-        mkdir -p $out
-        cp ${dashboardLockfileSrc}/dashboard/package.json $out/
-        cp ${dashboardLockfileSrc}/dashboard/package-lock.json $out/
-        # Minimal files so vite build succeeds (produces empty output)
-        echo '<!DOCTYPE html><html><head></head><body></body></html>' > $out/index.html
-        mkdir -p $out/src
-        touch $out/src/app.html
-      '';
-
-      # Deps-only build using stub source (for prettier-svelte)
-      # Only rebuilds when package.json or package-lock.json change
-      dashboardDeps = inputs.dream2nix.lib.evalModules {
-        packageSets.nixpkgs = pkgs;
-        modules = [
-          ./dashboard.nix
-          {
-            paths.projectRoot = inputs.self;
-            paths.projectRootFile = "flake.nix";
-            paths.package = inputs.self + "/dashboard";
-          }
-          {
-            deps.dashboardSrc = lib.mkForce dashboardStubSrc;
-          }
-          # Override build phases to skip the actual build - just need node_modules
-          {
-            mkDerivation = {
-              buildPhase = lib.mkForce "true";
-              installPhase = lib.mkForce ''
-                runHook preInstall
-                runHook postInstall
-              '';
-            };
-          }
-        ];
-      };
-
       # Filter source to only include dashboard directory
       dashboardSrc = lib.cleanSourceWith {
         src = inputs.self;
@@ -97,12 +42,11 @@
       '';
 
       # Prettier with svelte plugin for treefmt
-      # Uses dashboardDeps instead of dashboardFull to avoid rebuilding on source changes
       packages.prettier-svelte = pkgs.writeShellScriptBin "prettier-svelte" ''
-        export NODE_PATH="${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules"
+        export NODE_PATH="${dashboardFull}/lib/node_modules/exo-dashboard/node_modules"
         exec ${pkgs.nodejs}/bin/node \
-          ${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules/prettier/bin/prettier.cjs \
-          --plugin "${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules/prettier-plugin-svelte/plugin.js" \
+          ${dashboardFull}/lib/node_modules/exo-dashboard/node_modules/prettier/bin/prettier.cjs \
+          --plugin "${dashboardFull}/lib/node_modules/exo-dashboard/node_modules/prettier-plugin-svelte/plugin.js" \
           "$@"
       '';
     };

dashboard/src/lib/components/ChatForm.svelte

@@ -89,10 +89,7 @@
 
   const isImageModel = $derived(() => {
     if (!currentModel) return false;
-    return (
-      modelSupportsTextToImage(currentModel) ||
-      modelSupportsImageEditing(currentModel)
-    );
+    return modelSupportsTextToImage(currentModel);
   });
 
   const isEditOnlyWithoutImage = $derived(
@@ -649,23 +646,6 @@
             </svg>
             <span>EDIT</span>
           </span>
-        {:else if isEditOnlyWithoutImage}
-          <span class="inline-flex items-center gap-1.5">
-            <svg
-              class="w-3.5 h-3.5"
-              fill="none"
-              viewBox="0 0 24 24"
-              stroke="currentColor"
-              stroke-width="2"
-            >
-              <path
-                stroke-linecap="round"
-                stroke-linejoin="round"
-                d="M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z"
-              />
-            </svg>
-            <span>EDIT</span>
-          </span>
         {:else if isImageModel()}
           <span class="inline-flex items-center gap-1.5">
             <svg

dashboard/src/lib/components/ImageParamsPanel.svelte

@@ -110,36 +110,6 @@
     setImageGenerationParams({ negativePrompt: value || null });
   }
 
-  function handleNumImagesChange(event: Event) {
-    const input = event.target as HTMLInputElement;
-    const value = input.value.trim();
-    if (value === "") {
-      setImageGenerationParams({ numImages: 1 });
-    } else {
-      const num = parseInt(value, 10);
-      if (!isNaN(num) && num >= 1) {
-        setImageGenerationParams({ numImages: num });
-      }
-    }
-  }
-
-  function handleStreamChange(enabled: boolean) {
-    setImageGenerationParams({ stream: enabled });
-  }
-
-  function handlePartialImagesChange(event: Event) {
-    const input = event.target as HTMLInputElement;
-    const value = input.value.trim();
-    if (value === "") {
-      setImageGenerationParams({ partialImages: 0 });
-    } else {
-      const num = parseInt(value, 10);
-      if (!isNaN(num) && num >= 0) {
-        setImageGenerationParams({ partialImages: num });
-      }
-    }
-  }
-
   function clearSteps() {
     setImageGenerationParams({ numInferenceSteps: null });
   }
@@ -164,92 +134,90 @@
 <div class="border-b border-exo-medium-gray/30 px-3 py-2">
   <!-- Basic params row -->
   <div class="flex items-center gap-3 flex-wrap">
-    <!-- Size (hidden in edit mode - output size comes from input image) -->
-    {#if !isEditMode}
-      <div class="flex items-center gap-1.5">
-        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
-          >SIZE:</span
+    <!-- Size -->
+    <div class="flex items-center gap-1.5">
+      <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+        >SIZE:</span
+      >
+      <div class="relative">
+        <button
+          bind:this={sizeButtonRef}
+          type="button"
+          onclick={() => (isSizeDropdownOpen = !isSizeDropdownOpen)}
+          class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen
+            ? 'border-exo-yellow/70'
+            : ''}"
         >
-        <div class="relative">
-          <button
-            bind:this={sizeButtonRef}
-            type="button"
-            onclick={() => (isSizeDropdownOpen = !isSizeDropdownOpen)}
-            class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen
-              ? 'border-exo-yellow/70'
-              : ''}"
+          {params.size}
+        </button>
+        <div
+          class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen
+            ? 'rotate-180'
+            : ''}"
+        >
+          <svg
+            class="w-3 h-3 text-exo-yellow/60"
+            fill="none"
+            viewBox="0 0 24 24"
+            stroke="currentColor"
           >
-            {params.size}
-          </button>
-          <div
-            class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen
-              ? 'rotate-180'
-              : ''}"
-          >
-            <svg
-              class="w-3 h-3 text-exo-yellow/60"
-              fill="none"
-              viewBox="0 0 24 24"
-              stroke="currentColor"
-            >
-              <path
-                stroke-linecap="round"
-                stroke-linejoin="round"
-                stroke-width="2"
-                d="M19 9l-7 7-7-7"
-              />
-            </svg>
+            <path
+              stroke-linecap="round"
+              stroke-linejoin="round"
+              stroke-width="2"
+              d="M19 9l-7 7-7-7"
+            />
+          </svg>
+        </div>
+      </div>
+
+      {#if isSizeDropdownOpen}
+        <!-- Backdrop to close dropdown -->
+        <button
+          type="button"
+          class="fixed inset-0 z-[9998] cursor-default"
+          onclick={() => (isSizeDropdownOpen = false)}
+          aria-label="Close dropdown"
+        ></button>
+
+        <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
+        <div
+          class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto min-w-max"
+          style="bottom: calc(100vh - {sizeDropdownPosition()
+            .top}px + 4px); left: {sizeDropdownPosition().left}px;"
+        >
+          <div class="py-1">
+            {#each sizeOptions as size}
+              <button
+                type="button"
+                onclick={() => selectSize(size)}
+                class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {params.size ===
+                size
+                  ? 'bg-transparent text-exo-yellow'
+                  : 'text-exo-light-gray hover:text-exo-yellow'}"
+              >
+                {#if params.size === size}
+                  <svg
+                    class="w-3 h-3 flex-shrink-0"
+                    fill="currentColor"
+                    viewBox="0 0 20 20"
+                  >
+                    <path
+                      fill-rule="evenodd"
+                      d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
+                      clip-rule="evenodd"
+                    />
+                  </svg>
+                {:else}
+                  <span class="w-3"></span>
+                {/if}
+                <span>{size}</span>
+              </button>
+            {/each}
           </div>
         </div>
-
-        {#if isSizeDropdownOpen}
-          <!-- Backdrop to close dropdown -->
-          <button
-            type="button"
-            class="fixed inset-0 z-[9998] cursor-default"
-            onclick={() => (isSizeDropdownOpen = false)}
-            aria-label="Close dropdown"
-          ></button>
-
-          <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
-          <div
-            class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto min-w-max"
-            style="bottom: calc(100vh - {sizeDropdownPosition()
-              .top}px + 4px); left: {sizeDropdownPosition().left}px;"
-          >
-            <div class="py-1">
-              {#each sizeOptions as size}
-                <button
-                  type="button"
-                  onclick={() => selectSize(size)}
-                  class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {params.size ===
-                  size
-                    ? 'bg-transparent text-exo-yellow'
-                    : 'text-exo-light-gray hover:text-exo-yellow'}"
-                >
-                  {#if params.size === size}
-                    <svg
-                      class="w-3 h-3 flex-shrink-0"
-                      fill="currentColor"
-                      viewBox="0 0 20 20"
-                    >
-                      <path
-                        fill-rule="evenodd"
-                        d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
-                        clip-rule="evenodd"
-                      />
-                    </svg>
-                  {:else}
-                    <span class="w-3"></span>
-                  {/if}
-                  <span>{size}</span>
-                </button>
-              {/each}
-            </div>
-          </div>
-        {/if}
-      </div>
-    {/if}
+      {/if}
+    </div>
 
     <!-- Quality -->
     <div class="flex items-center gap-1.5">
@@ -357,59 +325,6 @@
       </div>
     </div>
 
-    <!-- Number of Images (not in edit mode) -->
-    {#if !isEditMode}
-      <div class="flex items-center gap-1.5">
-        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
-          >IMAGES:</span
-        >
-        <input
-          type="number"
-          min="1"
-          value={params.numImages}
-          oninput={handleNumImagesChange}
-          class="w-12 bg-exo-medium-gray/50 border border-exo-yellow/30 rounded px-2 py-1 text-xs font-mono text-exo-yellow text-center transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
-        />
-      </div>
-    {/if}
-
-    <!-- Stream toggle -->
-    <div class="flex items-center gap-1.5">
-      <span class="text-xs text-exo-light-gray uppercase tracking-wider"
-        >STREAM:</span
-      >
-      <button
-        type="button"
-        onclick={() => handleStreamChange(!params.stream)}
-        class="w-8 h-4 rounded-full transition-all duration-200 cursor-pointer relative {params.stream
-          ? 'bg-exo-yellow'
-          : 'bg-exo-medium-gray/50 border border-exo-yellow/30'}"
-        title={params.stream ? "Streaming enabled" : "Streaming disabled"}
-      >
-        <div
-          class="absolute top-0.5 w-3 h-3 rounded-full transition-all duration-200 {params.stream
-            ? 'right-0.5 bg-exo-black'
-            : 'left-0.5 bg-exo-light-gray'}"
-        ></div>
-      </button>
-    </div>
-
-    <!-- Partial Images (only when streaming) -->
-    {#if params.stream}
-      <div class="flex items-center gap-1.5">
-        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
-          >PARTIALS:</span
-        >
-        <input
-          type="number"
-          min="0"
-          value={params.partialImages}
-          oninput={handlePartialImagesChange}
-          class="w-12 bg-exo-medium-gray/50 border border-exo-yellow/30 rounded px-2 py-1 text-xs font-mono text-exo-yellow text-center transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
-        />
-      </div>
-    {/if}
-
     <!-- Input Fidelity (edit mode only) -->
     {#if isEditMode}
       <div class="flex items-center gap-1.5">

dashboard/src/routes/downloads/+page.svelte

@@ -6,8 +6,6 @@
     type DownloadProgress,
     refreshState,
     lastUpdate as lastUpdateStore,
-    startDownload,
-    deleteDownload,
   } from "$lib/stores/app.svelte";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
 
@@ -30,7 +28,6 @@
     etaMs: number;
     status: "completed" | "downloading";
     files: FileProgress[];
-    shardMetadata?: Record<string, unknown>;
   };
 
   type NodeEntry = {
@@ -272,12 +269,6 @@
             }
           }
 
-          // Extract shard_metadata for use with download actions
-          const shardMetadata = (downloadPayload.shard_metadata ??
-            downloadPayload.shardMetadata) as
-            | Record<string, unknown>
-            | undefined;
-
           const entry: ModelEntry = {
             modelId,
             prettyName,
@@ -294,7 +285,6 @@
                 ? "completed"
                 : "downloading",
             files,
-            shardMetadata,
           };
 
           const existing = modelMap.get(modelId);
@@ -479,52 +469,6 @@
                     >
                       {pct.toFixed(1)}%
                     </span>
-                    {#if model.status !== "completed" && model.shardMetadata}
-                      <button
-                        type="button"
-                        class="text-exo-light-gray hover:text-exo-yellow transition-colors"
-                        onclick={() =>
-                          startDownload(node.nodeId, model.shardMetadata!)}
-                        title="Start download"
-                      >
-                        <svg
-                          class="w-4 h-4"
-                          viewBox="0 0 20 20"
-                          fill="none"
-                          stroke="currentColor"
-                          stroke-width="2"
-                        >
-                          <path
-                            d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                            stroke-linecap="round"
-                            stroke-linejoin="round"
-                          ></path>
-                        </svg>
-                      </button>
-                    {/if}
-                    {#if model.status === "completed"}
-                      <button
-                        type="button"
-                        class="text-exo-light-gray hover:text-red-400 transition-colors"
-                        onclick={() =>
-                          deleteDownload(node.nodeId, model.modelId)}
-                        title="Delete download"
-                      >
-                        <svg
-                          class="w-4 h-4"
-                          viewBox="0 0 20 20"
-                          fill="none"
-                          stroke="currentColor"
-                          stroke-width="2"
-                        >
-                          <path
-                            d="M4 6h12M8 6V4h4v2m1 0v10a1 1 0 01-1 1H8a1 1 0 01-1-1V6h6"
-                            stroke-linecap="round"
-                            stroke-linejoin="round"
-                          ></path>
-                        </svg>
-                      </button>
-                    {/if}
                     <button
                       type="button"
                       class="text-exo-light-gray hover:text-exo-yellow transition-colors"

flake.nix

@@ -58,6 +58,11 @@
           pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
         in
         {
+          # Allow unfree for metal-toolchain (needed for Darwin Metal packages)
+          _module.args.pkgs = import inputs.nixpkgs {
+            inherit system;
+            config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain";
+          };
           treefmt = {
             projectRootFile = "flake.nix";
             programs = {
@@ -88,6 +93,13 @@
             touch $out
           '';
 
+          packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
+            metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
+            mlx = pkgs.callPackage ./nix/mlx.nix {
+              metal-toolchain = self'.packages.metal-toolchain;
+            };
+          };
+
           devShells.default = with pkgs; pkgs.mkShell {
             inputsFrom = [ self'.checks.cargo-build ];
 

nix/darwin-build-fixes.patch

@@ -0,0 +1,79 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index 0ed30932..d8528132 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -177,11 +177,7 @@ if(MLX_BUILD_METAL)
+     add_compile_definitions(MLX_METAL_DEBUG)
+   endif()
+
+-  # Throw an error if xcrun not found
+-  execute_process(
+-    COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
+-    OUTPUT_VARIABLE MACOS_SDK_VERSION
+-    OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY)
++  set(MACOS_SDK_VERSION @sdkVersion@)
+
+   if(${MACOS_SDK_VERSION} LESS 14.0)
+     message(
+@@ -199,11 +195,8 @@ if(MLX_BUILD_METAL)
+     endif()
+     set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+   endif()
+-  execute_process(
+-    COMMAND
+-      zsh "-c"
+-      "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
+-    OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
++  set(
++    MLX_METAL_VERSION @metalVersion@)
+   FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL})
+   FetchContent_MakeAvailable(metal_cpp)
+   target_include_directories(
+diff --git a/cmake/extension.cmake b/cmake/extension.cmake
+index 13db804a..5b385132 100644
+--- a/cmake/extension.cmake
++++ b/cmake/extension.cmake
+@@ -36,7 +36,7 @@ macro(mlx_build_metallib)
+   add_custom_command(
+     OUTPUT ${MTLLIB_BUILD_TARGET}
+     COMMAND
+-      xcrun -sdk macosx metal
++      metal -fmodules-cache-path=${CMAKE_BINARY_DIR}/metal-cache
+       "$<LIST:TRANSFORM,${MTLLIB_INCLUDE_DIRS},PREPEND,-I>"
+       ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET}
+     DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES}
+diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt
+index 262b0495..5c7446ad 100644
+--- a/mlx/backend/metal/kernels/CMakeLists.txt
++++ b/mlx/backend/metal/kernels/CMakeLists.txt
+@@ -29,7 +29,7 @@ function(build_kernel_base TARGET SRCFILE DEPS)
+                     "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+   endif()
+   add_custom_command(
+-    COMMAND xcrun -sdk macosx metal ${METAL_FLAGS} -c ${SRCFILE}
++    COMMAND metal -fmodules-cache-path=${CMAKE_BINARY_DIR}/metal-cache ${METAL_FLAGS} -c ${SRCFILE}
+             -I${PROJECT_SOURCE_DIR} -o ${TARGET}.air
+     DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS}
+     OUTPUT ${TARGET}.air
+@@ -170,7 +170,7 @@ endif()
+
+ add_custom_command(
+   OUTPUT ${MLX_METAL_PATH}/mlx.metallib
+-  COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o
++  COMMAND metallib ${KERNEL_AIR} -o
+           ${MLX_METAL_PATH}/mlx.metallib
+   DEPENDS ${KERNEL_AIR}
+   COMMENT "Building mlx.metallib"
+diff --git a/mlx/backend/metal/make_compiled_preamble.sh b/mlx/backend/metal/make_compiled_preamble.sh
+index bb55ed3a..94ea7dd7 100644
+--- a/mlx/backend/metal/make_compiled_preamble.sh
++++ b/mlx/backend/metal/make_compiled_preamble.sh
+@@ -31,7 +31,7 @@ OUTPUT_FILE=${OUTPUT_DIR}/${SRC_NAME}.cpp
+ mkdir -p "$OUTPUT_DIR"
+
+ # Use the metal compiler to get a list of headers (with depth)
+-CCC="xcrun -sdk macosx metal -x metal"
++CCC="metal -x metal -fmodules-cache-path=${OUTPUT_DIR}/metal-cache"
+ HDRS=$( $CCC -I"$SRC_DIR" -I"$JIT_INCLUDES" -DMLX_METAL_JIT -E -P -CC -C -H "$INPUT_FILE" $CFLAGS -w 2>&1 1>/dev/null )
+
+ # Remove any included system frameworks (for MetalPerformancePrimitive headers)

nix/metal-toolchain.nix

@@ -0,0 +1,56 @@
+{ lib, stdenvNoCC, requireFile, nix }:
+
+let
+  narFile = requireFile {
+    name = "metal-toolchain-17C48.nar";
+    message = ''
+      The Metal Toolchain NAR must be available.
+
+      If you have cachix configured for exo.cachix.org, this should be automatic.
+
+      Otherwise:
+        1. Install Xcode 26+ from the App Store
+        2. Run: xcodebuild -downloadComponent MetalToolchain
+        3. Export the toolchain:
+           hdiutil attach "$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' | head -1)" -mountpoint /tmp/metal-dmg
+           cp -R /tmp/metal-dmg/Metal.xctoolchain /tmp/metal-export
+           hdiutil detach /tmp/metal-dmg
+        4. Create NAR and add to store:
+           nix nar pack /tmp/metal-export > /tmp/metal-toolchain-17C48.nar
+           nix store add --mode flat /tmp/metal-toolchain-17C48.nar
+    '';
+    hash = "sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw=";
+  };
+in
+stdenvNoCC.mkDerivation {
+  pname = "metal-toolchain";
+  version = "17C48";
+
+  dontUnpack = true;
+  dontBuild = true;
+  dontFixup = true;
+
+  nativeBuildInputs = [ nix ];
+
+  installPhase = ''
+    runHook preInstall
+
+    nix-store --restore $out < ${narFile}
+
+    # Create bin directory with symlinks for PATH
+    mkdir -p $out/bin
+    ln -s $out/usr/bin/metal $out/bin/metal
+    ln -s $out/usr/bin/metallib $out/bin/metallib
+
+    runHook postInstall
+  '';
+
+  # Metal language version for CMake (from: echo __METAL_VERSION__ | metal -E -x metal -P -)
+  passthru.metalVersion = "400";
+
+  meta = {
+    description = "Apple Metal compiler toolchain";
+    platforms = [ "aarch64-darwin" ];
+    license = lib.licenses.unfree;
+  };
+}

nix/mlx.nix

@@ -0,0 +1,155 @@
+{ stdenv
+, lib
+, fetchFromGitHub
+, replaceVars
+, fetchzip
+, cmake
+, nlohmann_json
+, apple-sdk_26
+, metal-toolchain
+, runCommand
+, fmt
+, python313Packages
+}:
+
+assert stdenv.isDarwin;
+
+let
+  python = python313Packages.python;
+
+  # Static dependencies included directly during compilation
+  gguf-tools = fetchFromGitHub {
+    owner = "antirez";
+    repo = "gguf-tools";
+    rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
+    hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
+  };
+
+  metal_cpp = fetchzip {
+    url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
+    hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
+  };
+
+  nanobind = fetchFromGitHub {
+    owner = "wjakob";
+    repo = "nanobind";
+    rev = "v2.10.2";
+    hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
+    fetchSubmodules = true;
+  };
+
+  mlx = stdenv.mkDerivation rec {
+    pname = "mlx";
+    version = "0.30.3";
+    pyproject = true;
+
+    src = fetchFromGitHub {
+      owner = "ml-explore";
+      repo = "mlx";
+      tag = "v${version}";
+      hash = "sha256-Y4RTkGcDCZ9HLyflN0qYhPt/oVOsBhF1mHnKM4n1/ys=";
+    };
+
+    patches = [
+      (replaceVars ./darwin-build-fixes.patch {
+        sdkVersion = apple-sdk_26.version;
+        metalVersion = metal-toolchain.metalVersion;
+      })
+    ];
+
+    postPatch = ''
+      substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
+        --replace-fail "g++" "$CXX"
+    '';
+
+    dontUseCmakeConfigure = true;
+
+    enableParallelBuilding = true;
+
+    # Allows multiple cores to be used in Python builds.
+    postUnpack = ''
+      export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
+    '';
+
+    # Updates the wrong fetcher rev attribute
+    passthru.skipBulkUpdate = true;
+
+    env = {
+      DEV_RELEASE = 1;
+      CMAKE_ARGS = toString [
+        (lib.cmakeBool "USE_SYSTEM_FMT" true)
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
+        (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
+        (lib.cmakeBool "MLX_BUILD_METAL" true)
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
+        (lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${apple-sdk_26.version}")
+        (lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${apple-sdk_26.passthru.sdkroot}")
+      ];
+      SDKROOT = apple-sdk_26.passthru.sdkroot;
+      MACOSX_DEPLOYMENT_TARGET = apple-sdk_26.version;
+    };
+
+    build-system = [
+      python313Packages.setuptools
+    ];
+
+    nativeBuildInputs = [
+      cmake
+      metal-toolchain
+      python313Packages.pypaBuildHook
+      python313Packages.pypaInstallHook
+      python313Packages.setuptools
+      python313Packages.typing-extensions
+      python313Packages.wheel
+      python313Packages.cmake
+      python313Packages.ninja
+    ];
+
+    buildInputs = [
+      fmt
+      gguf-tools
+      python313Packages.nanobind
+      python313Packages.pybind11
+      apple-sdk_26
+    ];
+
+    # Tests require Metal GPU access which isn't available in the Nix sandbox.
+    # To run tests, build with: nix build --option sandbox false .#mlx.passthru.tests.mlxTest
+    doCheck = false;
+
+    pythonImportsCheck = [ "mlx" ];
+
+    passthru.tests = {
+      # Runs example scripts to verify MLX works. Requires --option sandbox false
+      # since Metal GPU access is needed.
+      mlxTest =
+        runCommand "run-mlx-examples"
+          {
+            buildInputs = [ mlx ];
+            nativeBuildInputs = [ python ];
+          }
+          ''
+            cp ${src}/examples/python/logistic_regression.py .
+            ${python.interpreter} logistic_regression.py
+            rm logistic_regression.py
+
+            cp ${src}/examples/python/linear_regression.py .
+            ${python.interpreter} linear_regression.py
+            rm linear_regression.py
+
+            touch $out
+          '';
+    };
+
+    meta = {
+      homepage = "https://github.com/ml-explore/mlx";
+      description = "Array framework for Apple silicon";
+      changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}";
+      license = lib.licenses.mit;
+      platforms = [ "aarch64-darwin" ];
+    };
+  };
+in
+mlx

pyproject.toml

@@ -19,14 +19,14 @@ dependencies = [
     "anyio==4.11.0",
     "mlx==0.30.3; sys_platform == 'darwin'",
     "mlx[cpu]==0.30.3; sys_platform == 'linux'",
-    "mlx-lm==0.30.5",
+    "mlx-lm @ git+https://github.com/AlexCheema/mlx-lm.git@fix-transformers-5.0.0rc2",
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",
     "openai-harmony>=0.0.8",
     "httpx>=0.28.1",
     "tomlkit>=0.14.0",
     "pillow>=11.0,<12.0", # compatibility with mflux
-    "mflux==0.15.4",
+    "mflux>=0.14.2",
     "python-multipart>=0.0.21",
 ]
 

src/exo/download/coordinator.py

@@ -1,284 +0,0 @@
-import asyncio
-from dataclasses import dataclass, field
-from typing import Iterator
-
-import anyio
-from anyio import current_time
-from anyio.abc import TaskGroup
-from loguru import logger
-
-from exo.download.download_utils import (
-    RepoDownloadProgress,
-    delete_model,
-    map_repo_download_progress_to_download_progress_data,
-)
-from exo.download.shard_downloader import ShardDownloader
-from exo.shared.models.model_cards import ModelId
-from exo.shared.types.commands import (
-    DeleteDownload,
-    ForwarderDownloadCommand,
-    StartDownload,
-)
-from exo.shared.types.common import NodeId, SessionId
-from exo.shared.types.events import (
-    Event,
-    ForwarderEvent,
-    NodeDownloadProgress,
-)
-from exo.shared.types.worker.downloads import (
-    DownloadCompleted,
-    DownloadFailed,
-    DownloadOngoing,
-    DownloadPending,
-    DownloadProgress,
-)
-from exo.shared.types.worker.shards import ShardMetadata
-from exo.utils.channels import Receiver, Sender, channel
-
-
-@dataclass
-class DownloadCoordinator:
-    node_id: NodeId
-    session_id: SessionId
-    shard_downloader: ShardDownloader
-    download_command_receiver: Receiver[ForwarderDownloadCommand]
-    local_event_sender: Sender[ForwarderEvent]
-    event_index_counter: Iterator[int]
-
-    # Local state
-    download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
-    active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
-
-    # 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)
-
-    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")
-        async with self._tg as tg:
-            tg.start_soon(self._command_processor)
-            tg.start_soon(self._forward_events)
-            tg.start_soon(self._emit_existing_download_progress)
-
-    def shutdown(self) -> None:
-        self._tg.cancel_scope.cancel()
-
-    async def _command_processor(self) -> None:
-        with self.download_command_receiver as commands:
-            async for cmd in commands:
-                # Only process commands targeting this node
-                if cmd.command.target_node_id != self.node_id:
-                    continue
-
-                match cmd.command:
-                    case StartDownload(shard_metadata=shard):
-                        await self._start_download(shard)
-                    case DeleteDownload(model_id=model_id):
-                        await self._delete_download(model_id)
-
-    async def _start_download(self, shard: ShardMetadata) -> None:
-        model_id = shard.model_card.model_id
-
-        # Check if already downloading or complete
-        if model_id in self.download_status:
-            status = self.download_status[model_id]
-            if isinstance(status, (DownloadOngoing, DownloadCompleted)):
-                logger.debug(
-                    f"Download for {model_id} already in progress or complete, skipping"
-                )
-                return
-
-        # Emit pending status
-        progress = DownloadPending(shard_metadata=shard, node_id=self.node_id)
-        self.download_status[model_id] = progress
-        await self.event_sender.send(NodeDownloadProgress(download_progress=progress))
-
-        # Check initial status from downloader
-        initial_progress = (
-            await self.shard_downloader.get_shard_download_status_for_shard(shard)
-        )
-
-        if initial_progress.status == "complete":
-            completed = DownloadCompleted(
-                shard_metadata=shard,
-                node_id=self.node_id,
-                total_bytes=initial_progress.total_bytes,
-            )
-            self.download_status[model_id] = completed
-            await self.event_sender.send(
-                NodeDownloadProgress(download_progress=completed)
-            )
-            return
-
-        # Start actual download
-        self._start_download_task(shard, initial_progress)
-
-    def _start_download_task(
-        self, shard: ShardMetadata, initial_progress: RepoDownloadProgress
-    ) -> None:
-        model_id = shard.model_card.model_id
-
-        # Emit ongoing status
-        status = DownloadOngoing(
-            node_id=self.node_id,
-            shard_metadata=shard,
-            download_progress=map_repo_download_progress_to_download_progress_data(
-                initial_progress
-            ),
-        )
-        self.download_status[model_id] = status
-        self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
-
-        last_progress_time = 0.0
-        throttle_interval_secs = 1.0
-
-        async def download_progress_callback(
-            callback_shard: ShardMetadata, progress: RepoDownloadProgress
-        ) -> None:
-            nonlocal last_progress_time
-
-            if progress.status == "complete":
-                completed = DownloadCompleted(
-                    shard_metadata=callback_shard,
-                    node_id=self.node_id,
-                    total_bytes=progress.total_bytes,
-                )
-                self.download_status[callback_shard.model_card.model_id] = completed
-                await self.event_sender.send(
-                    NodeDownloadProgress(download_progress=completed)
-                )
-                # Clean up active download tracking
-                if callback_shard.model_card.model_id in self.active_downloads:
-                    del self.active_downloads[callback_shard.model_card.model_id]
-            elif (
-                progress.status == "in_progress"
-                and current_time() - last_progress_time > throttle_interval_secs
-            ):
-                ongoing = DownloadOngoing(
-                    node_id=self.node_id,
-                    shard_metadata=callback_shard,
-                    download_progress=map_repo_download_progress_to_download_progress_data(
-                        progress
-                    ),
-                )
-                self.download_status[callback_shard.model_card.model_id] = ongoing
-                await self.event_sender.send(
-                    NodeDownloadProgress(download_progress=ongoing)
-                )
-                last_progress_time = current_time()
-
-        self.shard_downloader.on_progress(download_progress_callback)
-
-        async def download_wrapper() -> None:
-            try:
-                await self.shard_downloader.ensure_shard(shard)
-            except Exception as e:
-                logger.error(f"Download failed for {model_id}: {e}")
-                failed = DownloadFailed(
-                    shard_metadata=shard,
-                    node_id=self.node_id,
-                    error_message=str(e),
-                )
-                self.download_status[model_id] = failed
-                await self.event_sender.send(
-                    NodeDownloadProgress(download_progress=failed)
-                )
-            finally:
-                if model_id in self.active_downloads:
-                    del self.active_downloads[model_id]
-
-        task = asyncio.create_task(download_wrapper())
-        self.active_downloads[model_id] = task
-
-    async def _delete_download(self, model_id: ModelId) -> None:
-        # Cancel if active
-        if model_id in self.active_downloads:
-            logger.info(f"Cancelling active download for {model_id} before deletion")
-            self.active_downloads[model_id].cancel()
-            del self.active_downloads[model_id]
-
-        # Delete from disk
-        logger.info(f"Deleting model files for {model_id}")
-        deleted = await delete_model(model_id)
-
-        if deleted:
-            logger.info(f"Successfully deleted model {model_id}")
-        else:
-            logger.warning(f"Model {model_id} was not found on disk")
-
-        # Emit pending status to reset UI state, then remove from local tracking
-        if model_id in self.download_status:
-            current_status = self.download_status[model_id]
-            pending = DownloadPending(
-                shard_metadata=current_status.shard_metadata,
-                node_id=self.node_id,
-            )
-            await self.event_sender.send(
-                NodeDownloadProgress(download_progress=pending)
-            )
-            del self.download_status[model_id]
-
-    async def _forward_events(self) -> None:
-        with self.event_receiver as events:
-            async for event in events:
-                idx = next(self.event_index_counter)
-                fe = ForwarderEvent(
-                    origin_idx=idx,
-                    origin=self.node_id,
-                    session=self.session_id,
-                    event=event,
-                )
-                logger.debug(
-                    f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
-                )
-                await self.local_event_sender.send(fe)
-
-    async def _emit_existing_download_progress(self) -> None:
-        try:
-            while True:
-                logger.info(
-                    "DownloadCoordinator: Fetching and emitting existing download progress..."
-                )
-                async for (
-                    _,
-                    progress,
-                ) in self.shard_downloader.get_shard_download_status():
-                    if progress.status == "complete":
-                        status: DownloadProgress = DownloadCompleted(
-                            node_id=self.node_id,
-                            shard_metadata=progress.shard,
-                            total_bytes=progress.total_bytes,
-                        )
-                    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
-                            )
-                        else:
-                            status = DownloadOngoing(
-                                node_id=self.node_id,
-                                shard_metadata=progress.shard,
-                                download_progress=map_repo_download_progress_to_download_progress_data(
-                                    progress
-                                ),
-                            )
-                    else:
-                        continue
-
-                    self.download_status[progress.shard.model_card.model_id] = status
-                    await self.event_sender.send(
-                        NodeDownloadProgress(download_progress=status)
-                    )
-                logger.info(
-                    "DownloadCoordinator: Done emitting existing download progress."
-                )
-                await anyio.sleep(5 * 60)  # 5 minutes
-        except Exception as e:
-            logger.error(
-                f"DownloadCoordinator: Error emitting existing download progress: {e}"
-            )

src/exo/download/tests/test_download_verification.py

@@ -1,451 +0,0 @@
-"""Tests for download verification and cache behavior."""
-
-import time
-from collections.abc import AsyncIterator
-from datetime import timedelta
-from pathlib import Path
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import aiofiles
-import aiofiles.os as aios
-import pytest
-from pydantic import TypeAdapter
-
-from exo.download.download_utils import (
-    delete_model,
-    fetch_file_list_with_cache,
-)
-from exo.shared.types.common import ModelId
-from exo.shared.types.memory import Memory
-from exo.shared.types.worker.downloads import FileListEntry, RepoFileDownloadProgress
-
-
-@pytest.fixture
-def model_id() -> ModelId:
-    return ModelId("test-org/test-model")
-
-
-@pytest.fixture
-async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
-    """Set up a temporary models directory for testing."""
-    models_dir = tmp_path / "models"
-    await aios.makedirs(models_dir, exist_ok=True)
-    with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
-        yield models_dir
-
-
-class TestFileVerification:
-    """Tests for file size verification in _download_file."""
-
-    async def test_redownload_when_file_size_changes_upstream(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """Test that files with mismatched sizes are re-downloaded."""
-        # Import inside test to allow patching
-        from exo.download.download_utils import (
-            _download_file,  # pyright: ignore[reportPrivateUsage]
-        )
-
-        target_dir = tmp_path / "downloads"
-        await aios.makedirs(target_dir, exist_ok=True)
-
-        # Create a local file with wrong size
-        local_file = target_dir / "test.safetensors"
-        async with aiofiles.open(local_file, "wb") as f:
-            await f.write(b"local content")  # 13 bytes
-
-        remote_size = 1000  # Different from local
-        remote_hash = "abc123"
-
-        with (
-            patch(
-                "exo.download.download_utils.file_meta",
-                new_callable=AsyncMock,
-                return_value=(remote_size, remote_hash),
-            ) as mock_file_meta,
-            patch(
-                "exo.download.download_utils.create_http_session"
-            ) as mock_session_factory,
-        ):
-            # Set up mock HTTP response for re-download
-            mock_response = MagicMock()
-            mock_response.status = 200
-            mock_response.content.read = AsyncMock(  # pyright: ignore[reportAny]
-                side_effect=[b"x" * remote_size, b""]
-            )
-
-            mock_session = MagicMock()
-            mock_session.get.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
-                return_value=mock_response
-            )
-            mock_session.get.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
-                return_value=None
-            )
-            mock_session_factory.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
-                return_value=mock_session
-            )
-            mock_session_factory.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
-                return_value=None
-            )
-
-            # Mock calc_hash to return the expected hash
-            with patch(
-                "exo.download.download_utils.calc_hash",
-                new_callable=AsyncMock,
-                return_value=remote_hash,
-            ):
-                await _download_file(model_id, "main", "test.safetensors", target_dir)
-
-            # file_meta should be called twice: once for verification, once for download
-            assert mock_file_meta.call_count == 2
-
-    async def test_skip_download_when_file_size_matches(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """Test that files with matching sizes are not re-downloaded."""
-        from exo.download.download_utils import (
-            _download_file,  # pyright: ignore[reportPrivateUsage]
-        )
-
-        target_dir = tmp_path / "downloads"
-        await aios.makedirs(target_dir, exist_ok=True)
-
-        # Create a local file
-        local_file = target_dir / "test.safetensors"
-        local_content = b"local content"
-        async with aiofiles.open(local_file, "wb") as f:
-            await f.write(local_content)
-
-        remote_size = len(local_content)  # Same as local
-        remote_hash = "abc123"
-
-        with (
-            patch(
-                "exo.download.download_utils.file_meta",
-                new_callable=AsyncMock,
-                return_value=(remote_size, remote_hash),
-            ) as mock_file_meta,
-            patch(
-                "exo.download.download_utils.create_http_session"
-            ) as mock_session_factory,
-        ):
-            result = await _download_file(
-                model_id, "main", "test.safetensors", target_dir
-            )
-
-            # Should return immediately without downloading
-            assert result == local_file
-            mock_file_meta.assert_called_once()
-            mock_session_factory.assert_not_called()
-
-    async def test_offline_fallback_uses_local_file(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """Test that local files are used when network is unavailable."""
-        from exo.download.download_utils import (
-            _download_file,  # pyright: ignore[reportPrivateUsage]
-        )
-
-        target_dir = tmp_path / "downloads"
-        await aios.makedirs(target_dir, exist_ok=True)
-
-        # Create a local file
-        local_file = target_dir / "test.safetensors"
-        async with aiofiles.open(local_file, "wb") as f:
-            await f.write(b"local content")
-
-        with (
-            patch(
-                "exo.download.download_utils.file_meta",
-                new_callable=AsyncMock,
-                side_effect=Exception("Network error"),
-            ),
-            patch(
-                "exo.download.download_utils.create_http_session"
-            ) as mock_session_factory,
-        ):
-            result = await _download_file(
-                model_id, "main", "test.safetensors", target_dir
-            )
-
-            # Should return local file without attempting download
-            assert result == local_file
-            mock_session_factory.assert_not_called()
-
-
-class TestFileListCache:
-    """Tests for file list caching behavior."""
-
-    async def test_fetch_fresh_and_update_cache(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """Test that fresh data is fetched and cache is updated."""
-        models_dir = tmp_path / "models"
-
-        file_list = [
-            FileListEntry(type="file", path="model.safetensors", size=1000),
-            FileListEntry(type="file", path="config.json", size=100),
-        ]
-
-        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
-            patch(
-                "exo.download.download_utils.fetch_file_list_with_retry",
-                new_callable=AsyncMock,
-                return_value=file_list,
-            ) as mock_fetch,
-        ):
-            result = await fetch_file_list_with_cache(model_id, "main")
-
-            assert result == file_list
-            mock_fetch.assert_called_once()
-
-            # Verify cache was written
-            cache_file = (
-                models_dir
-                / "caches"
-                / model_id.normalize()
-                / f"{model_id.normalize()}--main--file_list.json"
-            )
-            assert await aios.path.exists(cache_file)
-
-            async with aiofiles.open(cache_file, "r") as f:
-                cached_data = TypeAdapter(list[FileListEntry]).validate_json(
-                    await f.read()
-                )
-            assert cached_data == file_list
-
-    async def test_fallback_to_cache_when_fetch_fails(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """Test that cached data is used when fetch fails."""
-        models_dir = tmp_path / "models"
-        cache_dir = models_dir / "caches" / model_id.normalize()
-        await aios.makedirs(cache_dir, exist_ok=True)
-
-        # Create cache file
-        cached_file_list = [
-            FileListEntry(type="file", path="model.safetensors", size=1000),
-        ]
-        cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
-        async with aiofiles.open(cache_file, "w") as f:
-            await f.write(
-                TypeAdapter(list[FileListEntry]).dump_json(cached_file_list).decode()
-            )
-
-        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
-            patch(
-                "exo.download.download_utils.fetch_file_list_with_retry",
-                new_callable=AsyncMock,
-                side_effect=Exception("Network error"),
-            ),
-        ):
-            result = await fetch_file_list_with_cache(model_id, "main")
-
-            assert result == cached_file_list
-
-    async def test_error_propagates_when_no_cache(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """Test that errors propagate when fetch fails and no cache exists."""
-        models_dir = tmp_path / "models"
-
-        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
-            patch(
-                "exo.download.download_utils.fetch_file_list_with_retry",
-                new_callable=AsyncMock,
-                side_effect=Exception("Network error"),
-            ),
-            pytest.raises(Exception, match="Network error"),
-        ):
-            await fetch_file_list_with_cache(model_id, "main")
-
-
-class TestModelDeletion:
-    """Tests for model deletion including cache cleanup."""
-
-    async def test_delete_model_clears_cache(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """Test that deleting a model also deletes its cache."""
-        models_dir = tmp_path / "models"
-        model_dir = models_dir / model_id.normalize()
-        cache_dir = models_dir / "caches" / model_id.normalize()
-
-        # Create model and cache directories
-        await aios.makedirs(model_dir, exist_ok=True)
-        await aios.makedirs(cache_dir, exist_ok=True)
-
-        # Add some files
-        async with aiofiles.open(model_dir / "model.safetensors", "w") as f:
-            await f.write("model data")
-        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
-            await f.write("[]")
-
-        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
-            result = await delete_model(model_id)
-
-            assert result is True
-            assert not await aios.path.exists(model_dir)
-            assert not await aios.path.exists(cache_dir)
-
-    async def test_delete_model_only_cache_exists(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """Test deleting when only cache exists (model already deleted)."""
-        models_dir = tmp_path / "models"
-        cache_dir = models_dir / "caches" / model_id.normalize()
-
-        # Only create cache directory
-        await aios.makedirs(cache_dir, exist_ok=True)
-        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
-            await f.write("[]")
-
-        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
-            result = await delete_model(model_id)
-
-            # Returns False because model dir didn't exist
-            assert result is False
-            # But cache should still be cleaned up
-            assert not await aios.path.exists(cache_dir)
-
-    async def test_delete_nonexistent_model(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        """Test deleting a model that doesn't exist."""
-        models_dir = tmp_path / "models"
-        await aios.makedirs(models_dir, exist_ok=True)
-
-        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
-            result = await delete_model(model_id)
-
-            assert result is False
-
-
-class TestProgressResetOnRedownload:
-    """Tests for progress tracking when files are re-downloaded."""
-
-    async def test_progress_resets_correctly_on_redownload(
-        self, model_id: ModelId
-    ) -> None:
-        """Test that progress tracking resets when a file is re-downloaded.
-
-        When a file is deleted and re-downloaded (due to size mismatch),
-        the progress tracking should reset rather than calculating negative
-        downloaded_this_session values.
-        """
-        # Simulate file_progress dict as it exists in download_shard
-        file_progress: dict[str, RepoFileDownloadProgress] = {}
-
-        # Initialize with old file progress (simulating existing large file)
-        old_file_size = 1_500_000_000  # 1.5 GB
-        file_progress["model.safetensors"] = RepoFileDownloadProgress(
-            repo_id=model_id,
-            repo_revision="main",
-            file_path="model.safetensors",
-            downloaded=Memory.from_bytes(old_file_size),
-            downloaded_this_session=Memory.from_bytes(0),
-            total=Memory.from_bytes(old_file_size),
-            speed=0,
-            eta=timedelta(0),
-            status="not_started",
-            start_time=time.time() - 10,  # Started 10 seconds ago
-        )
-
-        # Simulate the logic from on_progress_wrapper after re-download starts
-        # This is the exact logic from the fixed on_progress_wrapper
-        curr_bytes = 100_000  # 100 KB - new download just started
-        previous_progress = file_progress.get("model.safetensors")
-
-        # Detect re-download: curr_bytes < previous downloaded
-        is_redownload = (
-            previous_progress is not None
-            and curr_bytes < previous_progress.downloaded.in_bytes
-        )
-
-        if is_redownload or previous_progress is None:
-            # Fresh download or re-download: reset tracking
-            start_time = time.time()
-            downloaded_this_session = curr_bytes
-        else:
-            # Continuing download: accumulate
-            start_time = previous_progress.start_time
-            downloaded_this_session = (
-                previous_progress.downloaded_this_session.in_bytes
-                + (curr_bytes - previous_progress.downloaded.in_bytes)
-            )
-
-        # Key assertions
-        assert is_redownload is True, "Should detect re-download scenario"
-        assert downloaded_this_session == curr_bytes, (
-            "downloaded_this_session should equal curr_bytes on re-download"
-        )
-        assert downloaded_this_session > 0, (
-            "downloaded_this_session should be positive, not negative"
-        )
-
-        # Calculate speed (should be positive)
-        elapsed = time.time() - start_time
-        speed = downloaded_this_session / elapsed if elapsed > 0 else 0
-        assert speed >= 0, "Speed should be non-negative"
-
-    async def test_progress_accumulates_on_continuing_download(
-        self, model_id: ModelId
-    ) -> None:
-        """Test that progress accumulates correctly for continuing downloads.
-
-        When a download continues from where it left off (resume),
-        the progress should accumulate correctly.
-        """
-        file_progress: dict[str, RepoFileDownloadProgress] = {}
-
-        # Initialize with partial download progress
-        initial_downloaded = 500_000  # 500 KB already downloaded
-        start_time = time.time() - 5  # Started 5 seconds ago
-        file_progress["model.safetensors"] = RepoFileDownloadProgress(
-            repo_id=model_id,
-            repo_revision="main",
-            file_path="model.safetensors",
-            downloaded=Memory.from_bytes(initial_downloaded),
-            downloaded_this_session=Memory.from_bytes(initial_downloaded),
-            total=Memory.from_bytes(1_000_000),
-            speed=100_000,
-            eta=timedelta(seconds=5),
-            status="in_progress",
-            start_time=start_time,
-        )
-
-        # Progress callback with more bytes downloaded
-        curr_bytes = 600_000  # 600 KB - continuing download
-        previous_progress = file_progress.get("model.safetensors")
-
-        # This is NOT a re-download (curr_bytes > previous downloaded)
-        is_redownload = (
-            previous_progress is not None
-            and curr_bytes < previous_progress.downloaded.in_bytes
-        )
-
-        if is_redownload or previous_progress is None:
-            downloaded_this_session = curr_bytes
-            used_start_time = time.time()
-        else:
-            used_start_time = previous_progress.start_time
-            downloaded_this_session = (
-                previous_progress.downloaded_this_session.in_bytes
-                + (curr_bytes - previous_progress.downloaded.in_bytes)
-            )
-
-        # Key assertions
-        assert is_redownload is False, (
-            "Should NOT detect re-download for continuing download"
-        )
-        assert used_start_time == start_time, "Should preserve original start_time"
-        expected_session = initial_downloaded + (curr_bytes - initial_downloaded)
-        assert downloaded_this_session == expected_session, (
-            f"Should accumulate: {downloaded_this_session} == {expected_session}"
-        )
-        assert downloaded_this_session == 600_000, (
-            "downloaded_this_session should equal total downloaded so far"
-        )

src/exo/main.py

@@ -1,11 +1,10 @@
 import argparse
-import itertools
 import multiprocessing as mp
 import os
 import resource
 import signal
 from dataclasses import dataclass, field
-from typing import Iterator, Self
+from typing import Self
 
 import anyio
 from anyio.abc import TaskGroup
@@ -13,8 +12,6 @@ from loguru import logger
 from pydantic import PositiveInt
 
 import exo.routing.topics as topics
-from exo.download.coordinator import DownloadCoordinator
-from exo.download.impl_shard_downloader import exo_shard_downloader
 from exo.master.api import API  # TODO: should API be in master?
 from exo.master.main import Master
 from exo.routing.router import Router, get_node_id_keypair
@@ -24,6 +21,7 @@ from exo.shared.logging import logger_cleanup, logger_setup
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, channel
 from exo.utils.pydantic_ext import CamelCaseModel
+from exo.worker.download.impl_shard_downloader import exo_shard_downloader
 from exo.worker.main import Worker
 
 
@@ -31,7 +29,6 @@ from exo.worker.main import Worker
 @dataclass
 class Node:
     router: Router
-    download_coordinator: DownloadCoordinator | None
     worker: Worker | None
     election: Election  # Every node participates in election, as we do want a node to become master even if it isn't a master candidate if no master candidates are present.
     election_result_receiver: Receiver[ElectionResult]
@@ -39,7 +36,6 @@ class Node:
     api: API | None
 
     node_id: NodeId
-    event_index_counter: Iterator[int]
     _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
 
     @classmethod
@@ -53,26 +49,8 @@ class Node:
         await router.register_topic(topics.COMMANDS)
         await router.register_topic(topics.ELECTION_MESSAGES)
         await router.register_topic(topics.CONNECTION_MESSAGES)
-        await router.register_topic(topics.DOWNLOAD_COMMANDS)
 
         logger.info(f"Starting node {node_id}")
-
-        # Create shared event index counter for Worker and DownloadCoordinator
-        event_index_counter = itertools.count()
-
-        # Create DownloadCoordinator (unless --no-downloads)
-        if not args.no_downloads:
-            download_coordinator = DownloadCoordinator(
-                node_id,
-                session_id,
-                exo_shard_downloader(),
-                download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
-                local_event_sender=router.sender(topics.LOCAL_EVENTS),
-                event_index_counter=event_index_counter,
-            )
-        else:
-            download_coordinator = None
-
         if args.spawn_api:
             api = API(
                 node_id,
@@ -80,7 +58,6 @@ class Node:
                 port=args.api_port,
                 global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
                 command_sender=router.sender(topics.COMMANDS),
-                download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
             )
         else:
@@ -90,12 +67,11 @@ class Node:
             worker = Worker(
                 node_id,
                 session_id,
+                exo_shard_downloader(),
                 connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES),
                 global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
                 local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 command_sender=router.sender(topics.COMMANDS),
-                download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
-                event_index_counter=event_index_counter,
             )
         else:
             worker = None
@@ -123,25 +99,13 @@ class Node:
             election_result_sender=er_send,
         )
 
-        return cls(
-            router,
-            download_coordinator,
-            worker,
-            election,
-            er_recv,
-            master,
-            api,
-            node_id,
-            event_index_counter,
-        )
+        return cls(router, worker, election, er_recv, master, api, node_id)
 
     async def run(self):
         async with self._tg as tg:
             signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
             tg.start_soon(self.router.run)
             tg.start_soon(self.election.run)
-            if self.download_coordinator:
-                tg.start_soon(self.download_coordinator.run)
             if self.worker:
                 tg.start_soon(self.worker.run)
             if self.master:
@@ -206,27 +170,13 @@ class Node:
                     )
                 if result.is_new_master:
                     await anyio.sleep(0)
-                    # Fresh counter for new session (buffer expects indices from 0)
-                    self.event_index_counter = itertools.count()
-                    if self.download_coordinator:
-                        self.download_coordinator.shutdown()
-                        self.download_coordinator = DownloadCoordinator(
-                            self.node_id,
-                            result.session_id,
-                            exo_shard_downloader(),
-                            download_command_receiver=self.router.receiver(
-                                topics.DOWNLOAD_COMMANDS
-                            ),
-                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
-                            event_index_counter=self.event_index_counter,
-                        )
-                        self._tg.start_soon(self.download_coordinator.run)
                     if self.worker:
                         self.worker.shutdown()
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
                             result.session_id,
+                            exo_shard_downloader(),
                             connection_message_receiver=self.router.receiver(
                                 topics.CONNECTION_MESSAGES
                             ),
@@ -235,10 +185,6 @@ class Node:
                             ),
                             local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
                             command_sender=self.router.sender(topics.COMMANDS),
-                            download_command_sender=self.router.sender(
-                                topics.DOWNLOAD_COMMANDS
-                            ),
-                            event_index_counter=self.event_index_counter,
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
@@ -280,7 +226,6 @@ class Args(CamelCaseModel):
     api_port: PositiveInt = 52415
     tb_only: bool = False
     no_worker: bool = False
-    no_downloads: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
@@ -323,11 +268,6 @@ class Args(CamelCaseModel):
             "--no-worker",
             action="store_true",
         )
-        parser.add_argument(
-            "--no-downloads",
-            action="store_true",
-            help="Disable the download coordinator (node won't download models)",
-        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/api.py

@@ -1,5 +1,4 @@
 import base64
-import contextlib
 import json
 import time
 from collections.abc import AsyncGenerator
@@ -34,7 +33,6 @@ from exo.shared.models.model_cards import (
     ModelId,
 )
 from exo.shared.types.api import (
-    AdvancedImageParams,
     BenchChatCompletionResponse,
     BenchChatCompletionTaskParams,
     BenchImageGenerationResponse,
@@ -44,7 +42,6 @@ from exo.shared.types.api import (
     ChatCompletionResponse,
     CreateInstanceParams,
     CreateInstanceResponse,
-    DeleteDownloadResponse,
     DeleteInstanceResponse,
     ErrorInfo,
     ErrorResponse,
@@ -62,8 +59,6 @@ from exo.shared.types.api import (
     PlaceInstanceParams,
     PlacementPreview,
     PlacementPreviewResponse,
-    StartDownloadParams,
-    StartDownloadResponse,
     StreamingChoiceResponse,
     ToolCall,
 )
@@ -78,16 +73,12 @@ from exo.shared.types.commands import (
     ChatCompletion,
     Command,
     CreateInstance,
-    DeleteDownload,
     DeleteInstance,
-    DownloadCommand,
     ForwarderCommand,
-    ForwarderDownloadCommand,
     ImageEdits,
     ImageGeneration,
     PlaceInstance,
     SendInputChunk,
-    StartDownload,
     TaskFinished,
 )
 from exo.shared.types.common import CommandId, Id, NodeId, SessionId
@@ -163,14 +154,12 @@ class API:
         # Ideally this would be a MasterForwarderEvent but type system says no :(
         global_event_receiver: Receiver[ForwarderEvent],
         command_sender: Sender[ForwarderCommand],
-        download_command_sender: Sender[ForwarderDownloadCommand],
         # This lets us pause the API if an election is running
         election_receiver: Receiver[ElectionMessage],
     ) -> None:
         self.state = State()
         self._event_log: list[Event] = []
         self.command_sender = command_sender
-        self.download_command_sender = download_command_sender
         self.global_event_receiver = global_event_receiver
         self.election_receiver = election_receiver
         self.event_buffer: OrderedBuffer[Event] = OrderedBuffer[Event]()
@@ -269,8 +258,6 @@ class API:
         self.app.get("/images/{image_id}")(self.get_image)
         self.app.get("/state")(lambda: self.state)
         self.app.get("/events")(lambda: self._event_log)
-        self.app.post("/download/start")(self.start_download)
-        self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
 
     async def place_instance(self, payload: PlaceInstanceParams):
         command = PlaceInstance(
@@ -356,9 +343,14 @@ class API:
     ) -> PlacementPreviewResponse:
         seen: set[tuple[ModelId, Sharding, InstanceMeta, int]] = set()
         previews: list[PlacementPreview] = []
-        required_nodes = set(node_ids) if node_ids else None
 
-        if len(list(self.state.topology.list_nodes())) == 0:
+        # Create filtered topology if node_ids specified
+        if node_ids and len(node_ids) > 0:
+            topology = self.state.topology.get_subgraph_from_nodes(node_ids)
+        else:
+            topology = self.state.topology
+
+        if len(list(topology.list_nodes())) == 0:
             return PlacementPreviewResponse(previews=[])
 
         cards = [card for card in MODEL_CARDS.values() if card.model_id == model_id]
@@ -371,9 +363,7 @@ class API:
                 instance_combinations.extend(
                     [
                         (sharding, instance_meta, i)
-                        for i in range(
-                            1, len(list(self.state.topology.list_nodes())) + 1
-                        )
+                        for i in range(1, len(list(topology.list_nodes())) + 1)
                     ]
                 )
         # TODO: PDD
@@ -391,9 +381,8 @@ class API:
                         ),
                         node_memory=self.state.node_memory,
                         node_network=self.state.node_network,
-                        topology=self.state.topology,
+                        topology=topology,
                         current_instances=self.state.instances,
-                        required_nodes=required_nodes,
                     )
                 except ValueError as exc:
                     if (model_card.model_id, sharding, instance_meta, 0) not in seen:
@@ -432,16 +421,14 @@ class API:
 
                 instance = new_instances[0]
                 shard_assignments = instance.shard_assignments
-                placement_node_ids = list(shard_assignments.node_to_runner.keys())
+                node_ids = list(shard_assignments.node_to_runner.keys())
 
                 memory_delta_by_node: dict[str, int] = {}
-                if placement_node_ids:
+                if node_ids:
                     total_bytes = model_card.storage_size.in_bytes
-                    per_node = total_bytes // len(placement_node_ids)
-                    remainder = total_bytes % len(placement_node_ids)
-                    for index, node_id in enumerate(
-                        sorted(placement_node_ids, key=str)
-                    ):
+                    per_node = total_bytes // len(node_ids)
+                    remainder = total_bytes % len(node_ids)
+                    for index, node_id in enumerate(sorted(node_ids, key=str)):
                         extra = 1 if index < remainder else 0
                         memory_delta_by_node[str(node_id)] = per_node + extra
 
@@ -449,7 +436,7 @@ class API:
                     model_card.model_id,
                     sharding,
                     instance_meta,
-                    len(placement_node_ids),
+                    len(node_ids),
                 ) not in seen:
                     previews.append(
                         PlacementPreview(
@@ -461,14 +448,7 @@ class API:
                             error=None,
                         )
                     )
-                seen.add(
-                    (
-                        model_card.model_id,
-                        sharding,
-                        instance_meta,
-                        len(placement_node_ids),
-                    )
-                )
+                seen.add((model_card.model_id, sharding, instance_meta, len(node_ids)))
 
         return PlacementPreviewResponse(previews=previews)
 
@@ -855,7 +835,6 @@ class API:
                             # Yield partial image event (always use b64_json for partials)
                             event_data = {
                                 "type": "partial",
-                                "image_index": chunk.image_index,
                                 "partial_index": partial_idx,
                                 "total_partials": total_partials,
                                 "format": str(chunk.format),
@@ -1045,9 +1024,6 @@ class API:
         stream: bool,
         partial_images: int,
         bench: bool,
-        quality: Literal["high", "medium", "low"],
-        output_format: Literal["png", "jpeg", "webp"],
-        advanced_params: AdvancedImageParams | None,
     ) -> ImageEdits:
         """Prepare and send an image edits command with chunked image upload."""
         resolved_model = await self._validate_image_model(model)
@@ -1076,9 +1052,6 @@ class API:
                 stream=stream,
                 partial_images=partial_images,
                 bench=bench,
-                quality=quality,
-                output_format=output_format,
-                advanced_params=advanced_params,
             ),
         )
 
@@ -1113,22 +1086,12 @@ class API:
         input_fidelity: Literal["low", "high"] = Form("low"),
         stream: str = Form("false"),
         partial_images: str = Form("0"),
-        quality: Literal["high", "medium", "low"] = Form("medium"),
-        output_format: Literal["png", "jpeg", "webp"] = Form("png"),
-        advanced_params: str | None = Form(None),
     ) -> ImageGenerationResponse | StreamingResponse:
         """Handle image editing requests (img2img)."""
         # Parse string form values to proper types
         stream_bool = stream.lower() in ("true", "1", "yes")
         partial_images_int = int(partial_images) if partial_images.isdigit() else 0
 
-        parsed_advanced_params: AdvancedImageParams | None = None
-        if advanced_params:
-            with contextlib.suppress(Exception):
-                parsed_advanced_params = AdvancedImageParams.model_validate_json(
-                    advanced_params
-                )
-
         command = await self._send_image_edits_command(
             image=image,
             prompt=prompt,
@@ -1140,9 +1103,6 @@ class API:
             stream=stream_bool,
             partial_images=partial_images_int,
             bench=False,
-            quality=quality,
-            output_format=output_format,
-            advanced_params=parsed_advanced_params,
         )
 
         if stream_bool and partial_images_int > 0:
@@ -1173,18 +1133,8 @@ class API:
         size: str = Form("1024x1024"),
         response_format: Literal["url", "b64_json"] = Form("b64_json"),
         input_fidelity: Literal["low", "high"] = Form("low"),
-        quality: Literal["high", "medium", "low"] = Form("medium"),
-        output_format: Literal["png", "jpeg", "webp"] = Form("png"),
-        advanced_params: str | None = Form(None),
     ) -> BenchImageGenerationResponse:
         """Handle benchmark image editing requests with generation stats."""
-        parsed_advanced_params: AdvancedImageParams | None = None
-        if advanced_params:
-            with contextlib.suppress(Exception):
-                parsed_advanced_params = AdvancedImageParams.model_validate_json(
-                    advanced_params
-                )
-
         command = await self._send_image_edits_command(
             image=image,
             prompt=prompt,
@@ -1196,9 +1146,6 @@ class API:
             stream=False,
             partial_images=0,
             bench=True,
-            quality=quality,
-            output_format=output_format,
-            advanced_params=parsed_advanced_params,
         )
 
         return await self._collect_image_generation_with_stats(
@@ -1310,28 +1257,3 @@ class API:
         await self.command_sender.send(
             ForwarderCommand(origin=self.node_id, command=command)
         )
-
-    async def _send_download(self, command: DownloadCommand):
-        await self.download_command_sender.send(
-            ForwarderDownloadCommand(origin=self.node_id, command=command)
-        )
-
-    async def start_download(
-        self, payload: StartDownloadParams
-    ) -> StartDownloadResponse:
-        command = StartDownload(
-            target_node_id=payload.target_node_id,
-            shard_metadata=payload.shard_metadata,
-        )
-        await self._send_download(command)
-        return StartDownloadResponse(command_id=command.command_id)
-
-    async def delete_download(
-        self, node_id: NodeId, model_id: ModelId
-    ) -> DeleteDownloadResponse:
-        command = DeleteDownload(
-            target_node_id=node_id,
-            model_id=ModelId(model_id),
-        )
-        await self._send_download(command)
-        return DeleteDownloadResponse(command_id=command.command_id)

src/exo/master/placement.py

@@ -35,7 +35,7 @@ from exo.shared.types.worker.shards import Sharding
 
 def random_ephemeral_port() -> int:
     port = random.randint(49153, 65535)
-    return port - 1 if port <= 52415 else port
+    return port - 1 if port <= 52415 else 52414
 
 
 def add_instance_to_placements(
@@ -54,18 +54,9 @@ def place_instance(
     current_instances: Mapping[InstanceId, Instance],
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
-    required_nodes: set[NodeId] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
-
-    # Filter to cycles containing all required nodes (subset matching)
-    if required_nodes:
-        candidate_cycles = [
-            cycle
-            for cycle in candidate_cycles
-            if required_nodes.issubset(cycle.node_ids)
-        ]
     cycles_with_sufficient_memory = filter_cycles_by_memory(
         candidate_cycles, node_memory, command.model_card.storage_size
     )

src/exo/master/placement_utils.py

@@ -257,13 +257,7 @@ def _find_ip_prioritised(
     ip_to_type = {
         iface.ip_address: iface.interface_type for iface in other_network.interfaces
     }
-    priority = {
-        "ethernet": 0,
-        "wifi": 1,
-        "unknown": 2,
-        "maybe_ethernet": 3,
-        "thunderbolt": 4,
-    }
+    priority = {"ethernet": 0, "wifi": 1, "unknown": 2, "thunderbolt": 3}
     return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
 
 

src/exo/routing/topics.py

@@ -3,7 +3,7 @@ from enum import Enum
 
 from exo.routing.connection_message import ConnectionMessage
 from exo.shared.election import ElectionMessage
-from exo.shared.types.commands import ForwarderCommand, ForwarderDownloadCommand
+from exo.shared.types.commands import ForwarderCommand
 from exo.shared.types.events import (
     ForwarderEvent,
 )
@@ -45,6 +45,3 @@ ELECTION_MESSAGES = TypedTopic(
 CONNECTION_MESSAGES = TypedTopic(
     "connection_messages", PublishPolicy.Never, ConnectionMessage
 )
-DOWNLOAD_COMMANDS = TypedTopic(
-    "download_commands", PublishPolicy.Always, ForwarderDownloadCommand
-)

src/exo/shared/models/model_cards.py

@@ -413,9 +413,9 @@ MODEL_CARDS: dict[str, ModelCard] = {
     ),
 }
 
-_IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
+_IMAGE_MODEL_CARDS: dict[str, ModelCard] = {
     "flux1-schnell": ModelCard(
-        model_id=ModelId("exolabs/FLUX.1-schnell"),
+        model_id=ModelId("black-forest-labs/FLUX.1-schnell"),
         storage_size=Memory.from_bytes(23782357120 + 9524621312),
         n_layers=57,
         hidden_size=1,
@@ -428,7 +428,7 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
                 storage_size=Memory.from_kb(0),
                 n_layers=12,
                 can_shard=False,
-                safetensors_index_filename=None,
+                safetensors_index_filename=None,  # Single file
             ),
             ComponentInfo(
                 component_name="text_encoder_2",
@@ -442,7 +442,7 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
                 component_name="transformer",
                 component_path="transformer/",
                 storage_size=Memory.from_bytes(23782357120),
-                n_layers=57,
+                n_layers=57,  # 19 transformer_blocks + 38 single_transformer_blocks
                 can_shard=True,
                 safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
             ),
@@ -457,7 +457,7 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
         ],
     ),
     "flux1-dev": ModelCard(
-        model_id=ModelId("exolabs/FLUX.1-dev"),
+        model_id=ModelId("black-forest-labs/FLUX.1-dev"),
         storage_size=Memory.from_bytes(23782357120 + 9524621312),
         n_layers=57,
         hidden_size=1,
@@ -470,7 +470,7 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
                 storage_size=Memory.from_kb(0),
                 n_layers=12,
                 can_shard=False,
-                safetensors_index_filename=None,
+                safetensors_index_filename=None,  # Single file
             ),
             ComponentInfo(
                 component_name="text_encoder_2",
@@ -484,49 +484,7 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
                 component_name="transformer",
                 component_path="transformer/",
                 storage_size=Memory.from_bytes(23802816640),
-                n_layers=57,
-                can_shard=True,
-                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
-            ),
-            ComponentInfo(
-                component_name="vae",
-                component_path="vae/",
-                storage_size=Memory.from_kb(0),
-                n_layers=None,
-                can_shard=False,
-                safetensors_index_filename=None,
-            ),
-        ],
-    ),
-    "flux1-krea-dev": ModelCard(
-        model_id=ModelId("exolabs/FLUX.1-Krea-dev"),
-        storage_size=Memory.from_bytes(23802816640 + 9524621312),  # Same as dev
-        n_layers=57,
-        hidden_size=1,
-        supports_tensor=False,
-        tasks=[ModelTask.TextToImage],
-        components=[
-            ComponentInfo(
-                component_name="text_encoder",
-                component_path="text_encoder/",
-                storage_size=Memory.from_kb(0),
-                n_layers=12,
-                can_shard=False,
-                safetensors_index_filename=None,
-            ),
-            ComponentInfo(
-                component_name="text_encoder_2",
-                component_path="text_encoder_2/",
-                storage_size=Memory.from_bytes(9524621312),
-                n_layers=24,
-                can_shard=False,
-                safetensors_index_filename="model.safetensors.index.json",
-            ),
-            ComponentInfo(
-                component_name="transformer",
-                component_path="transformer/",
-                storage_size=Memory.from_bytes(23802816640),
-                n_layers=57,
+                n_layers=57,  # 19 transformer_blocks + 38 single_transformer_blocks
                 can_shard=True,
                 safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
             ),
@@ -541,9 +499,9 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
         ],
     ),
     "qwen-image": ModelCard(
-        model_id=ModelId("exolabs/Qwen-Image"),
+        model_id=ModelId("Qwen/Qwen-Image"),
         storage_size=Memory.from_bytes(16584333312 + 40860802176),
-        n_layers=60,
+        n_layers=60,  # Qwen has 60 transformer blocks (all joint-style)
         hidden_size=1,
         supports_tensor=False,
         tasks=[ModelTask.TextToImage],
@@ -551,10 +509,10 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
             ComponentInfo(
                 component_name="text_encoder",
                 component_path="text_encoder/",
-                storage_size=Memory.from_bytes(16584333312),
+                storage_size=Memory.from_kb(16584333312),
                 n_layers=12,
                 can_shard=False,
-                safetensors_index_filename=None,
+                safetensors_index_filename=None,  # Single file
             ),
             ComponentInfo(
                 component_name="transformer",
@@ -575,9 +533,9 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
         ],
     ),
     "qwen-image-edit-2509": ModelCard(
-        model_id=ModelId("exolabs/Qwen-Image-Edit-2509"),
+        model_id=ModelId("Qwen/Qwen-Image-Edit-2509"),
         storage_size=Memory.from_bytes(16584333312 + 40860802176),
-        n_layers=60,
+        n_layers=60,  # Qwen has 60 transformer blocks (all joint-style)
         hidden_size=1,
         supports_tensor=False,
         tasks=[ModelTask.ImageToImage],
@@ -585,10 +543,10 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
             ComponentInfo(
                 component_name="text_encoder",
                 component_path="text_encoder/",
-                storage_size=Memory.from_bytes(16584333312),
+                storage_size=Memory.from_kb(16584333312),
                 n_layers=12,
                 can_shard=False,
-                safetensors_index_filename=None,
+                safetensors_index_filename=None,  # Single file
             ),
             ComponentInfo(
                 component_name="transformer",
@@ -610,92 +568,6 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
     ),
 }
 
-
-def _generate_image_model_quant_variants(
-    base_name: str,
-    base_card: ModelCard,
-) -> dict[str, ModelCard]:
-    """Create quantized variants of an image model card.
-
-    Only the transformer component is quantized; text encoders stay at bf16.
-    Sizes are calculated exactly from the base card's component sizes.
-    """
-    if base_card.components is None:
-        raise ValueError(f"Image model {base_name} must have components defined")
-
-    # quantizations = [8, 6, 5, 4, 3]
-    quantizations = [8, 4]
-
-    num_transformer_bytes = next(
-        c.storage_size.in_bytes
-        for c in base_card.components
-        if c.component_name == "transformer"
-    )
-
-    transformer_bytes = Memory.from_bytes(num_transformer_bytes)
-
-    remaining_bytes = Memory.from_bytes(
-        sum(
-            c.storage_size.in_bytes
-            for c in base_card.components
-            if c.component_name != "transformer"
-        )
-    )
-
-    def with_transformer_size(new_size: Memory) -> list[ComponentInfo]:
-        assert base_card.components is not None
-        return [
-            ComponentInfo(
-                component_name=c.component_name,
-                component_path=c.component_path,
-                storage_size=new_size
-                if c.component_name == "transformer"
-                else c.storage_size,
-                n_layers=c.n_layers,
-                can_shard=c.can_shard,
-                safetensors_index_filename=c.safetensors_index_filename,
-            )
-            for c in base_card.components
-        ]
-
-    variants = {
-        base_name: ModelCard(
-            model_id=base_card.model_id,
-            storage_size=transformer_bytes + remaining_bytes,
-            n_layers=base_card.n_layers,
-            hidden_size=base_card.hidden_size,
-            supports_tensor=base_card.supports_tensor,
-            tasks=base_card.tasks,
-            components=with_transformer_size(transformer_bytes),
-        )
-    }
-
-    for quant in quantizations:
-        quant_transformer_bytes = Memory.from_bytes(
-            (num_transformer_bytes * quant) // 16
-        )
-        total_bytes = remaining_bytes + quant_transformer_bytes
-
-        model_id = ModelId(base_card.model_id + f"-{quant}bit")
-
-        variants[f"{base_name}-{quant}bit"] = ModelCard(
-            model_id=model_id,
-            storage_size=total_bytes,
-            n_layers=base_card.n_layers,
-            hidden_size=base_card.hidden_size,
-            supports_tensor=base_card.supports_tensor,
-            tasks=base_card.tasks,
-            components=with_transformer_size(quant_transformer_bytes),
-        )
-
-    return variants
-
-
-_image_model_cards: dict[str, ModelCard] = {}
-for _base_name, _base_card in _IMAGE_BASE_MODEL_CARDS.items():
-    _image_model_cards |= _generate_image_model_quant_variants(_base_name, _base_card)
-_IMAGE_MODEL_CARDS = _image_model_cards
-
 if EXO_ENABLE_IMAGE_MODELS:
     MODEL_CARDS.update(_IMAGE_MODEL_CARDS)
 
@@ -749,7 +621,7 @@ class ConfigData(BaseModel):
 
 async def get_config_data(model_id: ModelId) -> ConfigData:
     """Downloads and parses config.json for a model."""
-    from exo.download.download_utils import (
+    from exo.worker.download.download_utils import (
         download_file_with_retry,
         ensure_models_dir,
     )
@@ -771,11 +643,11 @@ async def get_config_data(model_id: ModelId) -> ConfigData:
 
 async def get_safetensors_size(model_id: ModelId) -> Memory:
     """Gets model size from safetensors index or falls back to HF API."""
-    from exo.download.download_utils import (
+    from exo.shared.types.worker.downloads import ModelSafetensorsIndex
+    from exo.worker.download.download_utils import (
         download_file_with_retry,
         ensure_models_dir,
     )
-    from exo.shared.types.worker.downloads import ModelSafetensorsIndex
 
     target_dir = (await ensure_models_dir()) / model_id.normalize()
     await aios.makedirs(target_dir, exist_ok=True)

src/exo/shared/topology.py

@@ -248,8 +248,8 @@ class Topology:
     ) -> list[list[NodeId]]:
         """
         Find cycles in the Thunderbolt topology where all nodes have TB bridge enabled.
-        Only returns cycles with >=2 nodes (2+ machines in a loop), as
-        1 node doesn't cause the broadcast storm problem.
+        Only returns cycles with >2 nodes (3+ machines in a loop), as cycles with
+        2 or fewer nodes don't cause the broadcast storm problem.
         """
         enabled_nodes = {
             node_id
@@ -257,7 +257,7 @@ class Topology:
             if status.enabled
         }
 
-        if len(enabled_nodes) < 2:
+        if len(enabled_nodes) < 3:
             return []
 
         thunderbolt_ips = _get_ips_with_interface_type(
@@ -288,7 +288,7 @@ class Topology:
         return [
             [graph[idx] for idx in cycle]
             for cycle in rx.simple_cycles(graph)
-            if len(cycle) >= 2
+            if len(cycle) > 2
         ]
 
 

src/exo/shared/types/api.py

@@ -7,11 +7,10 @@ from pydantic import BaseModel, Field, field_validator
 from pydantic_core import PydanticUseDefault
 
 from exo.shared.models.model_cards import ModelCard, ModelId
-from exo.shared.types.common import CommandId, NodeId
+from exo.shared.types.common import CommandId
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding, ShardMetadata
-from exo.utils.pydantic_ext import CamelCaseModel
+from exo.shared.types.worker.shards import Sharding
 
 FinishReason = Literal[
     "stop", "length", "tool_calls", "content_filter", "function_call", "error"
@@ -353,16 +352,3 @@ class ImageListItem(BaseModel, frozen=True):
 
 class ImageListResponse(BaseModel, frozen=True):
     data: list[ImageListItem]
-
-
-class StartDownloadParams(CamelCaseModel):
-    target_node_id: NodeId
-    shard_metadata: ShardMetadata
-
-
-class StartDownloadResponse(CamelCaseModel):
-    command_id: CommandId
-
-
-class DeleteDownloadResponse(CamelCaseModel):
-    command_id: CommandId

src/exo/shared/types/commands.py

@@ -1,6 +1,6 @@
 from pydantic import Field
 
-from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.models.model_cards import ModelCard
 from exo.shared.types.api import (
     ChatCompletionTaskParams,
     ImageEditsInternalParams,
@@ -9,7 +9,7 @@ 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.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding, ShardMetadata
+from exo.shared.types.worker.shards import Sharding
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
 
@@ -62,19 +62,6 @@ class RequestEventLog(BaseCommand):
     since_idx: int
 
 
-class StartDownload(BaseCommand):
-    target_node_id: NodeId
-    shard_metadata: ShardMetadata
-
-
-class DeleteDownload(BaseCommand):
-    target_node_id: NodeId
-    model_id: ModelId
-
-
-DownloadCommand = StartDownload | DeleteDownload
-
-
 Command = (
     TestCommand
     | RequestEventLog
@@ -92,8 +79,3 @@ Command = (
 class ForwarderCommand(CamelCaseModel):
     origin: NodeId
     command: Command
-
-
-class ForwarderDownloadCommand(CamelCaseModel):
-    origin: NodeId
-    command: DownloadCommand

src/exo/shared/types/mlx.py

@@ -1,12 +0,0 @@
-"""Shared types for MLX-related functionality."""
-
-from collections.abc import Sequence
-
-from mlx_lm.models.cache import (
-    KVCache,
-    QuantizedKVCache,
-    RotatingKVCache,
-)
-
-# This list contains one cache entry per transformer layer
-KVCacheType = Sequence[KVCache | RotatingKVCache | QuantizedKVCache]

src/exo/shared/types/profiling.py

@@ -48,7 +48,7 @@ class SystemPerformanceProfile(CamelCaseModel):
     ecpu_usage: float = 0.0
 
 
-InterfaceType = Literal["wifi", "ethernet", "maybe_ethernet", "thunderbolt", "unknown"]
+InterfaceType = Literal["wifi", "ethernet", "thunderbolt", "unknown"]
 
 
 class NetworkInterfaceInfo(CamelCaseModel):

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

@@ -30,7 +30,6 @@ class ImageGenerationResponse(BaseRunnerResponse):
     image_data: bytes
     format: Literal["png", "jpeg", "webp"] = "png"
     stats: ImageGenerationStats | None = None
-    image_index: int = 0
 
     def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
         for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
@@ -45,7 +44,6 @@ class PartialImageResponse(BaseRunnerResponse):
     format: Literal["png", "jpeg", "webp"] = "png"
     partial_index: int
     total_partials: int
-    image_index: int = 0
 
     def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
         for name, value in super().__repr_args__():  # pyright: ignore[reportAny]

src/exo/utils/info_gatherer/info_gatherer.py

@@ -349,8 +349,13 @@ class InfoGatherer:
     async def _monitor_misc(self):
         if self.misc_poll_interval is None:
             return
+        prev = await MiscData.gather()
+        await self.info_sender.send(prev)
         while True:
-            await self.info_sender.send(await MiscData.gather())
+            curr = await MiscData.gather()
+            if prev != curr:
+                prev = curr
+                await self.info_sender.send(curr)
             await anyio.sleep(self.misc_poll_interval)
 
     async def _monitor_system_profiler_thunderbolt_data(self):
@@ -360,12 +365,15 @@ class InfoGatherer:
         if iface_map is None:
             return
 
+        old_idents = []
         while True:
             data = await ThunderboltConnectivity.gather()
             assert data is not None
 
             idents = [it for i in data if (it := i.ident(iface_map)) is not None]
-            await self.info_sender.send(MacThunderboltIdentifiers(idents=idents))
+            if idents != old_idents:
+                await self.info_sender.send(MacThunderboltIdentifiers(idents=idents))
+            old_idents = idents
 
             conns = [it for i in data if (it := i.conn()) is not None]
             await self.info_sender.send(MacThunderboltConnections(conns=conns))
@@ -390,17 +398,22 @@ class InfoGatherer:
     async def _watch_system_info(self):
         if self.interface_watcher_interval is None:
             return
+        old_nics = []
         while True:
-            nics = await get_network_interfaces()
-            await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
+            nics = get_network_interfaces()
+            if nics != old_nics:
+                old_nics = nics
+                await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
             await anyio.sleep(self.interface_watcher_interval)
 
     async def _monitor_thunderbolt_bridge_status(self):
         if self.thunderbolt_bridge_poll_interval is None:
             return
+        prev: ThunderboltBridgeInfo | None = None
         while True:
             curr = await ThunderboltBridgeInfo.gather()
-            if curr is not None:
+            if curr is not None and prev != curr:
+                prev = curr
                 await self.info_sender.send(curr)
             await anyio.sleep(self.thunderbolt_bridge_poll_interval)
 

src/exo/utils/info_gatherer/system_info.py

@@ -1,6 +1,6 @@
 import socket
 import sys
-from subprocess import CalledProcessError
+from subprocess import CalledProcessError, run
 
 import psutil
 from anyio import run_process
@@ -16,7 +16,8 @@ async def get_friendly_name() -> str:
     """
     hostname = socket.gethostname()
 
-    if sys.platform != "darwin":
+    # TODO: better non mac support
+    if sys.platform != "darwin":  # 'darwin' is the platform name for macOS
         return hostname
 
     try:
@@ -27,20 +28,21 @@ async def get_friendly_name() -> str:
     return process.stdout.decode("utf-8", errors="replace").strip() or hostname
 
 
-async def _get_interface_types_from_networksetup() -> dict[str, InterfaceType]:
+def _get_interface_types_from_networksetup() -> dict[str, InterfaceType]:
     """Parse networksetup -listallhardwareports to get interface types."""
     if sys.platform != "darwin":
         return {}
-
     try:
-        result = await run_process(["networksetup", "-listallhardwareports"])
-    except CalledProcessError:
+        result = run(
+            ["networksetup", "-listallhardwareports"], capture_output=True, text=True
+        )
+    except Exception:
         return {}
 
     types: dict[str, InterfaceType] = {}
     current_type: InterfaceType = "unknown"
 
-    for line in result.stdout.decode().splitlines():
+    for line in result.stdout.splitlines():
         if line.startswith("Hardware Port:"):
             port_name = line.split(":", 1)[1].strip()
             if "Wi-Fi" in port_name:
@@ -53,15 +55,12 @@ async def _get_interface_types_from_networksetup() -> dict[str, InterfaceType]:
                 current_type = "unknown"
         elif line.startswith("Device:"):
             device = line.split(":", 1)[1].strip()
-            # enX is ethernet adapters or thunderbolt - these must be deprioritised
-            if device.startswith("en") and device not in ["en0", "en1"]:
-                current_type = "maybe_ethernet"
             types[device] = current_type
 
     return types
 
 
-async def get_network_interfaces() -> list[NetworkInterfaceInfo]:
+def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     """
     Retrieves detailed network interface information on macOS.
     Parses output from 'networksetup -listallhardwareports' and 'ifconfig'
@@ -69,7 +68,7 @@ async def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     Returns a list of NetworkInterfaceInfo objects.
     """
     interfaces_info: list[NetworkInterfaceInfo] = []
-    interface_types = await _get_interface_types_from_networksetup()
+    interface_types = _get_interface_types_from_networksetup()
 
     for iface, services in psutil.net_if_addrs().items():
         for service in services:

src/exo/utils/keyed_backoff.py

@@ -1,32 +0,0 @@
-import time
-from typing import Generic, TypeVar
-
-K = TypeVar("K")
-
-
-class KeyedBackoff(Generic[K]):
-    """Tracks exponential backoff state per key."""
-
-    def __init__(self, base: float = 0.5, cap: float = 10.0):
-        self._base = base
-        self._cap = cap
-        self._attempts: dict[K, int] = {}
-        self._last_time: dict[K, float] = {}
-
-    def should_proceed(self, key: K) -> bool:
-        """Returns True if enough time has elapsed since last attempt."""
-        now = time.monotonic()
-        last = self._last_time.get(key, 0.0)
-        attempts = self._attempts.get(key, 0)
-        delay = min(self._cap, self._base * (2.0**attempts))
-        return now - last >= delay
-
-    def record_attempt(self, key: K) -> None:
-        """Record that an attempt was made for this key."""
-        self._last_time[key] = time.monotonic()
-        self._attempts[key] = self._attempts.get(key, 0) + 1
-
-    def reset(self, key: K) -> None:
-        """Reset backoff state for a key (e.g., on success)."""
-        self._attempts.pop(key, None)
-        self._last_time.pop(key, None)

src/exo/worker/download/download_utils.py

@@ -24,15 +24,7 @@ from pydantic import (
     TypeAdapter,
 )
 
-from exo.download.huggingface_utils import (
-    filter_repo_objects,
-    get_allow_patterns,
-    get_auth_headers,
-    get_hf_endpoint,
-    get_hf_token,
-)
 from exo.shared.constants import EXO_MODELS_DIR
-from exo.shared.models.model_cards import ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.downloads import (
@@ -43,6 +35,13 @@ from exo.shared.types.worker.downloads import (
     RepoFileDownloadProgress,
 )
 from exo.shared.types.worker.shards import ShardMetadata
+from exo.worker.download.huggingface_utils import (
+    filter_repo_objects,
+    get_allow_patterns,
+    get_auth_headers,
+    get_hf_endpoint,
+    get_hf_token,
+)
 
 
 class HuggingFaceAuthenticationError(Exception):
@@ -121,20 +120,11 @@ async def ensure_models_dir() -> Path:
 
 
 async def delete_model(model_id: ModelId) -> bool:
-    models_dir = await ensure_models_dir()
-    model_dir = models_dir / model_id.normalize()
-    cache_dir = models_dir / "caches" / model_id.normalize()
-
-    deleted = False
-    if await aios.path.exists(model_dir):
-        await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
-        deleted = True
-
-    # Also clear cache
-    if await aios.path.exists(cache_dir):
-        await asyncio.to_thread(shutil.rmtree, cache_dir, ignore_errors=False)
-
-    return deleted
+    model_dir = await ensure_models_dir() / model_id.normalize()
+    if not await aios.path.exists(model_dir):
+        return False
+    await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
+    return True
 
 
 async def seed_models(seed_dir: str | Path):
@@ -160,28 +150,16 @@ async def fetch_file_list_with_cache(
     target_dir = (await ensure_models_dir()) / "caches" / model_id.normalize()
     await aios.makedirs(target_dir, exist_ok=True)
     cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json"
-
-    # Always try fresh first
-    try:
-        file_list = await fetch_file_list_with_retry(
-            model_id, revision, recursive=recursive
-        )
-        # Update cache with fresh data
-        async with aiofiles.open(cache_file, "w") as f:
-            await f.write(
-                TypeAdapter(list[FileListEntry]).dump_json(file_list).decode()
-            )
-        return file_list
-    except Exception as e:
-        # Fetch failed - try cache fallback
-        if await aios.path.exists(cache_file):
-            logger.warning(
-                f"Failed to fetch file list for {model_id}, using cached data: {e}"
-            )
-            async with aiofiles.open(cache_file, "r") as f:
-                return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
-        # No cache available, propagate the error
-        raise
+    if await aios.path.exists(cache_file):
+        async with aiofiles.open(cache_file, "r") as f:
+            return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
+    file_list = await fetch_file_list_with_retry(
+        model_id, revision, recursive=recursive
+    )
+    await aios.makedirs(cache_file.parent, exist_ok=True)
+    async with aiofiles.open(cache_file, "w") as f:
+        await f.write(TypeAdapter(list[FileListEntry]).dump_json(file_list).decode())
+    return file_list
 
 
 async def fetch_file_list_with_retry(
@@ -353,28 +331,8 @@ async def _download_file(
     target_dir: Path,
     on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None,
 ) -> Path:
-    target_path = target_dir / path
-
-    if await aios.path.exists(target_path):
-        local_size = (await aios.stat(target_path)).st_size
-
-        # Try to verify against remote, but allow offline operation
-        try:
-            remote_size, _ = await file_meta(model_id, revision, path)
-            if local_size != remote_size:
-                logger.info(
-                    f"File {path} size mismatch (local={local_size}, remote={remote_size}), re-downloading"
-                )
-                await aios.remove(target_path)
-            else:
-                return target_path
-        except Exception as e:
-            # Offline or network error - trust local file
-            logger.debug(
-                f"Could not verify {path} against remote (offline?): {e}, using local file"
-            )
-            return target_path
-
+    if await aios.path.exists(target_dir / path):
+        return target_dir / path
     await aios.makedirs((target_dir / path).parent, exist_ok=True)
     length, etag = await file_meta(model_id, revision, path)
     remote_hash = etag[:-5] if etag.endswith("-gzip") else etag
@@ -523,11 +481,6 @@ async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
         return ["*"]
 
 
-def is_image_model(shard: ShardMetadata) -> bool:
-    tasks = shard.model_card.tasks
-    return ModelTask.TextToImage in tasks or ModelTask.ImageToImage in tasks
-
-
 async def get_downloaded_size(path: Path) -> int:
     partial_path = path.with_suffix(path.suffix + ".partial")
     if await aios.path.exists(path):
@@ -569,40 +522,22 @@ async def download_shard(
             file_list, allow_patterns=allow_patterns, key=lambda x: x.path
         )
     )
-
-    # For image models, skip root-level safetensors files since weights
-    # are stored in component subdirectories (e.g., transformer/, vae/)
-    if is_image_model(shard):
-        filtered_file_list = [
-            f
-            for f in filtered_file_list
-            if "/" in f.path or not f.path.endswith(".safetensors")
-        ]
     file_progress: dict[str, RepoFileDownloadProgress] = {}
 
     async def on_progress_wrapper(
         file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
     ) -> None:
-        previous_progress = file_progress.get(file.path)
-
-        # Detect re-download: curr_bytes < previous downloaded means file was deleted and restarted
-        is_redownload = (
-            previous_progress is not None
-            and curr_bytes < previous_progress.downloaded.in_bytes
+        start_time = (
+            file_progress[file.path].start_time
+            if file.path in file_progress
+            else time.time()
+        )
+        downloaded_this_session = (
+            file_progress[file.path].downloaded_this_session.in_bytes
+            + (curr_bytes - file_progress[file.path].downloaded.in_bytes)
+            if file.path in file_progress
+            else curr_bytes
         )
-
-        if is_redownload or previous_progress is None:
-            # Fresh download or re-download: reset tracking
-            start_time = time.time()
-            downloaded_this_session = curr_bytes
-        else:
-            # Continuing download: accumulate
-            start_time = previous_progress.start_time
-            downloaded_this_session = (
-                previous_progress.downloaded_this_session.in_bytes
-                + (curr_bytes - previous_progress.downloaded.in_bytes)
-            )
-
         speed = (
             downloaded_this_session / (time.time() - start_time)
             if time.time() - start_time > 0

src/exo/worker/download/impl_shard_downloader.py

@@ -5,13 +5,13 @@ from typing import AsyncIterator, Callable
 
 from loguru import logger
 
-from exo.download.download_utils import RepoDownloadProgress, download_shard
-from exo.download.shard_downloader import ShardDownloader
 from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
     ShardMetadata,
 )
+from exo.worker.download.download_utils import RepoDownloadProgress, download_shard
+from exo.worker.download.shard_downloader import ShardDownloader
 
 
 def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:

src/exo/worker/download/shard_downloader.py

@@ -5,13 +5,13 @@ from datetime import timedelta
 from pathlib import Path
 from typing import AsyncIterator, Callable
 
-from exo.download.download_utils import RepoDownloadProgress
 from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
     ShardMetadata,
 )
+from exo.worker.download.download_utils import RepoDownloadProgress
 
 
 # TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?

src/exo/worker/engines/image/distributed_model.py

@@ -6,10 +6,10 @@ import mlx.core as mx
 from mflux.models.common.config.config import Config
 from PIL import Image
 
-from exo.download.download_utils import build_model_path
 from exo.shared.types.api import AdvancedImageParams
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.shards import PipelineShardMetadata
+from exo.worker.download.download_utils import build_model_path
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models import (
     create_adapter_for_model,
@@ -140,7 +140,6 @@ class DistributedImageModel:
             width=width,
             image_path=image_path,
             model_config=self._adapter.model.model_config,  # pyright: ignore[reportAny]
-            guidance=guidance_override if guidance_override is not None else 4.0,
         )
 
         num_sync_steps = self._config.get_num_sync_steps(steps)

src/exo/worker/engines/image/generate.py

@@ -75,20 +75,19 @@ def generate_image(
     intermediate images, then ImageGenerationResponse for the final image.
 
     Yields:
-        PartialImageResponse for intermediate images (if partial_images > 0, first image only)
-        ImageGenerationResponse for final complete images
+        PartialImageResponse for intermediate images (if partial_images > 0)
+        ImageGenerationResponse for the final complete image
     """
     width, height = parse_size(task.size)
     quality: Literal["low", "medium", "high"] = task.quality or "medium"
 
     advanced_params = task.advanced_params
     if advanced_params is not None and advanced_params.seed is not None:
-        base_seed = advanced_params.seed
+        seed = advanced_params.seed
     else:
-        base_seed = random.randint(0, 2**32 - 1)
+        seed = random.randint(0, 2**32 - 1)
 
     is_bench = getattr(task, "bench", False)
-    num_images = task.n or 1
 
     generation_start_time: float = 0.0
 
@@ -96,11 +95,7 @@ def generate_image(
         mx.reset_peak_memory()
         generation_start_time = time.perf_counter()
 
-    partial_images = (
-        task.partial_images
-        if task.partial_images is not None
-        else (3 if task.stream else 0)
-    )
+    partial_images = task.partial_images or (3 if task.stream else 0)
 
     image_path: Path | None = None
 
@@ -110,81 +105,72 @@ def generate_image(
             image_path = Path(tmpdir) / "input.png"
             image_path.write_bytes(base64.b64decode(task.image_data))
 
-        for image_num in range(num_images):
-            # Increment seed for each image to ensure unique results
-            current_seed = base_seed + image_num
+        # Iterate over generator results
+        for result in model.generate(
+            prompt=task.prompt,
+            height=height,
+            width=width,
+            quality=quality,
+            seed=seed,
+            image_path=image_path,
+            partial_images=partial_images,
+            advanced_params=advanced_params,
+        ):
+            if isinstance(result, tuple):
+                # Partial image: (Image, partial_index, total_partials)
+                image, partial_idx, total_partials = result
+                buffer = io.BytesIO()
+                image_format = task.output_format.upper()
+                if image_format == "JPG":
+                    image_format = "JPEG"
+                if image_format == "JPEG" and image.mode == "RGBA":
+                    image = image.convert("RGB")
+                image.save(buffer, format=image_format)
 
-            for result in model.generate(
-                prompt=task.prompt,
-                height=height,
-                width=width,
-                quality=quality,
-                seed=current_seed,
-                image_path=image_path,
-                partial_images=partial_images,
-                advanced_params=advanced_params,
-            ):
-                if isinstance(result, tuple):
-                    # Partial image: (Image, partial_index, total_partials)
-                    image, partial_idx, total_partials = result
-                    buffer = io.BytesIO()
-                    image_format = task.output_format.upper()
-                    if image_format == "JPG":
-                        image_format = "JPEG"
-                    if image_format == "JPEG" and image.mode == "RGBA":
-                        image = image.convert("RGB")
-                    image.save(buffer, format=image_format)
+                yield PartialImageResponse(
+                    image_data=buffer.getvalue(),
+                    format=task.output_format,
+                    partial_index=partial_idx,
+                    total_partials=total_partials,
+                )
+            else:
+                image = result
 
-                    yield PartialImageResponse(
-                        image_data=buffer.getvalue(),
-                        format=task.output_format,
-                        partial_index=partial_idx,
-                        total_partials=total_partials,
-                        image_index=image_num,
+                stats: ImageGenerationStats | None = None
+                if is_bench:
+                    generation_end_time = time.perf_counter()
+                    total_generation_time = generation_end_time - generation_start_time
+
+                    num_inference_steps = model.get_steps_for_quality(quality)
+
+                    seconds_per_step = (
+                        total_generation_time / num_inference_steps
+                        if num_inference_steps > 0
+                        else 0.0
                     )
-                else:
-                    image = result
 
-                    # Only include stats on the final image
-                    stats: ImageGenerationStats | None = None
-                    if is_bench and image_num == num_images - 1:
-                        generation_end_time = time.perf_counter()
-                        total_generation_time = (
-                            generation_end_time - generation_start_time
-                        )
+                    peak_memory_gb = mx.get_peak_memory() / (1024**3)
 
-                        num_inference_steps = model.get_steps_for_quality(quality)
-                        total_steps = num_inference_steps * num_images
-
-                        seconds_per_step = (
-                            total_generation_time / total_steps
-                            if total_steps > 0
-                            else 0.0
-                        )
-
-                        peak_memory_gb = mx.get_peak_memory() / (1024**3)
-
-                        stats = ImageGenerationStats(
-                            seconds_per_step=seconds_per_step,
-                            total_generation_time=total_generation_time,
-                            num_inference_steps=num_inference_steps,
-                            num_images=num_images,
-                            image_width=width,
-                            image_height=height,
-                            peak_memory_usage=Memory.from_gb(peak_memory_gb),
-                        )
-
-                    buffer = io.BytesIO()
-                    image_format = task.output_format.upper()
-                    if image_format == "JPG":
-                        image_format = "JPEG"
-                    if image_format == "JPEG" and image.mode == "RGBA":
-                        image = image.convert("RGB")
-                    image.save(buffer, format=image_format)
-
-                    yield ImageGenerationResponse(
-                        image_data=buffer.getvalue(),
-                        format=task.output_format,
-                        stats=stats,
-                        image_index=image_num,
+                    stats = ImageGenerationStats(
+                        seconds_per_step=seconds_per_step,
+                        total_generation_time=total_generation_time,
+                        num_inference_steps=num_inference_steps,
+                        num_images=task.n or 1,
+                        image_width=width,
+                        image_height=height,
+                        peak_memory_usage=Memory.from_gb(peak_memory_gb),
                     )
+
+                buffer = io.BytesIO()
+                image_format = task.output_format.upper()
+                if image_format == "JPG":
+                    image_format = "JPEG"
+                if image_format == "JPEG" and image.mode == "RGBA":
+                    image = image.convert("RGB")
+                image.save(buffer, format=image_format)
+
+                yield ImageGenerationResponse(
+                    image_data=buffer.getvalue(),
+                    format=task.output_format,
+                    stats=stats,
+                )

src/exo/worker/engines/image/models/__init__.py

@@ -33,7 +33,6 @@ _ADAPTER_REGISTRY: dict[str, AdapterFactory] = {
 # Config registry: maps model ID patterns to configs
 _CONFIG_REGISTRY: dict[str, ImageModelConfig] = {
     "flux.1-schnell": FLUX_SCHNELL_CONFIG,
-    "flux.1-krea-dev": FLUX_DEV_CONFIG,  # Must come before "flux.1-dev" for pattern matching
     "flux.1-dev": FLUX_DEV_CONFIG,
     "qwen-image-edit": QWEN_IMAGE_EDIT_CONFIG,  # Must come before "qwen-image" for pattern matching
     "qwen-image": QWEN_IMAGE_CONFIG,

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

@@ -19,8 +19,6 @@ from mlx_lm.models.deepseek_v32 import DeepseekV32MLP
 from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
 from mlx_lm.models.glm4_moe import Model as Glm4MoeModel
 from mlx_lm.models.glm4_moe import MoE
-from mlx_lm.models.glm4_moe_lite import Glm4MoeLiteDecoderLayer, Glm4MoeLiteMLP
-from mlx_lm.models.glm4_moe_lite import Model as GLM4MoeLiteModel
 from mlx_lm.models.gpt_oss import GptOssMoeModel
 from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.models.llama import Model as LlamaModel
@@ -147,10 +145,6 @@ class PipelineLastLayer(CustomMlxLayer):
             if cache is not None:
                 cache.keys = mx.depends(cache.keys, output)  # type: ignore[reportUnknownMemberType]
 
-        output = mx.distributed.all_gather(output, group=self.group)[
-            -output.shape[0] :
-        ]  # type :ignore
-
         return output
 
 
@@ -258,6 +252,10 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
         if cache is not None:
             cache[-1].state = mx.depends(cache[-1].state, logits)  # type: ignore
 
+        logits = mx.distributed.all_gather(logits, group=group)[
+            -logits.shape[0] :
+        ]  # type :ignore
+
         return logits
 
     cls.__call__ = patched_call
@@ -336,7 +334,15 @@ def tensor_auto_parallel(
         group=group,
     )
 
+    if hasattr(model, "shard") and not isinstance(model, GptOssModel):
+        try:
+            model.shard(group)  # type: ignore
+            return patch_tensor_model(model)
+        except (AttributeError, TypeError, NameError):
+            pass
+
     if isinstance(model, (LlamaModel, Ministral3Model)):
+        logger.warning("shouldn't be hit - upstream sharding exists")
         tensor_parallel_sharding_strategy = LlamaShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -345,6 +351,7 @@ def tensor_auto_parallel(
             sharded_to_all_linear_in_place,
         )
     elif isinstance(model, (DeepseekV3Model, DeepseekV32Model)):
+        logger.warning("shouldn't be hit - upstream sharding exists")
         tensor_parallel_sharding_strategy = DeepSeekShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -360,14 +367,6 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(model, GLM4MoeLiteModel):
-        tensor_parallel_sharding_strategy = GLM4MoeLiteShardingStrategy(
-            group,
-            all_to_sharded_linear,
-            sharded_to_all_linear,
-            all_to_sharded_linear_in_place,
-            sharded_to_all_linear_in_place,
-        )
     elif isinstance(model, (Qwen3MoeModel, Glm4MoeModel, Qwen3NextModel)):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
@@ -442,7 +441,7 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
             layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
             layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
             layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
-            mx.eval(layer)
+
         return model
 
 
@@ -517,8 +516,6 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
                 layer.mlp = ShardedDeepseekV3MoE(layer.mlp)  # type: ignore
                 layer.mlp.sharding_group = self.group
 
-            mx.eval(layer)
-
         return model
 
 
@@ -536,84 +533,6 @@ class ShardedDeepseekV3MoE(CustomMlxLayer):
         return y
 
 
-class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
-    def shard_model(
-        self,
-        model: nn.Module,
-        timeout_seconds: float,
-        on_timeout: TimeoutCallback | None,
-    ) -> nn.Module:
-        model = cast(GLM4MoeLiteModel, model)
-        for layer in model.layers:  # type: ignore
-            layer = cast(Glm4MoeLiteDecoderLayer, layer)
-            eval_with_timeout(
-                layer.parameters(),
-                timeout_seconds / len(model.layers),  # type: ignore
-                on_timeout,
-            )
-            if layer.self_attn.q_lora_rank is None:  # type: ignore
-                layer.self_attn.q_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_proj
-                )
-            else:
-                layer.self_attn.q_b_proj = self.all_to_sharded_linear(
-                    layer.self_attn.q_b_proj
-                )
-
-            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
-            layer.self_attn.num_heads //= self.N
-
-            # Logic from upstream mlx
-            num_heads = layer.self_attn.num_heads
-            sh = self.group.rank() * num_heads
-            eh = sh + num_heads
-
-            def shard_heads(w: mx.array, sh: int = sh, eh: int = eh) -> mx.array:
-                return w[sh:eh]
-
-            layer.self_attn.embed_q.apply(shard_heads)
-            layer.self_attn.unembed_out.apply(shard_heads)
-
-            if isinstance(layer.mlp, Glm4MoeLiteMLP):
-                layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
-                layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
-                layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
-
-            else:
-                if getattr(layer.mlp, "shared_experts", None) is not None:
-                    self.all_to_sharded_linear_in_place(
-                        layer.mlp.shared_experts.gate_proj
-                    )
-                    self.sharded_to_all_linear_in_place(
-                        layer.mlp.shared_experts.down_proj
-                    )
-                    self.all_to_sharded_linear_in_place(
-                        layer.mlp.shared_experts.up_proj
-                    )
-                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
-                self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
-                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
-                layer.mlp = ShardedGLM4MoeLiteMoE(layer.mlp)  # type: ignore
-                layer.mlp.sharding_group = self.group  # type: ignore
-            mx.eval(layer)
-
-        return model
-
-
-class ShardedGLM4MoeLiteMoE(CustomMlxLayer):
-    def __init__(self, layer: _LayerCallable):
-        super().__init__(layer)
-        self.sharding_group: mx.distributed.Group | None = None
-
-    def __call__(self, x: mx.array) -> mx.array:
-        if self.sharding_group is not None:
-            x = sum_gradients(self.sharding_group)(x)
-        y = self.original_layer.__call__(x)
-        if self.sharding_group is not None:
-            y = mx.distributed.all_sum(y, group=self.sharding_group)
-        return y
-
-
 class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
     def shard_model(
         self,
@@ -647,7 +566,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             )
             layer.block_sparse_moe = ShardedQwenMoE(layer.block_sparse_moe)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
             layer.block_sparse_moe.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
-            mx.eval(layer)
+
         return model
 
 
@@ -688,7 +607,6 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
-            mx.eval(layer)
         return model
 
 
@@ -743,7 +661,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
 
             layer.mlp = ShardedGptOssMoE(layer.mlp)  # type: ignore
             layer.mlp.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
-            mx.eval(layer)
+
         return model
 
 

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

@@ -1,81 +1,39 @@
-import os
+# type: ignore
+# TODO: Fix this file, including types!
 from copy import deepcopy
-from typing import Any, cast
+from typing import Callable
 
 import mlx.core as mx
-from mlx_lm.models.cache import (
-    KVCache,
-    QuantizedKVCache,
-    RotatingKVCache,
-    trim_prompt_cache,
-)
-from mlx_lm.models.gpt_oss import Model as GptOssModel
+from mlx_lm import stream_generate
+from mlx_lm.models.cache import _BaseCache, trim_prompt_cache
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
-from exo.shared.types.mlx import KVCacheType
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
-from exo.worker.runner.bootstrap import logger
-
-# Fraction of device memory above which LRU eviction kicks in
-_DEFAULT_MEMORY_THRESHOLD = 0.85
-_MEMORY_THRESHOLD = float(
-    os.environ.get("EXO_MEMORY_THRESHOLD", _DEFAULT_MEMORY_THRESHOLD)
-)
+from exo.worker.engines.mlx.constants import KEEP_KV_SIZE, KV_BITS, KV_GROUP_SIZE
+from exo.worker.engines.mlx.utils_mlx import make_kv_cache
 
 
 class KVPrefixCache:
-    def __init__(self, tokenizer: TokenizerWrapper):
+    def __init__(self):
+        # Only one prefix cache per runner.
         self.prompts: list[mx.array] = []  # mx array of tokens (ints)
-        self.caches: list[KVCacheType] = []
-        self._last_used: list[int] = []  # monotonic counter of last access per entry
-        self._access_counter: int = 0
-        self._tokenizer: TokenizerWrapper = tokenizer
+        self.caches: list[list[_BaseCache]] = []
 
-    def clear(self):
-        """Clear all cached prompts and caches."""
-        self.prompts.clear()
-        self.caches.clear()
-        self._last_used.clear()
-
-    def add_kv_cache(self, prompt: str, cache: KVCacheType):
-        """Add a new cache entry. Evicts LRU entries if memory is high."""
-        self._evict_if_needed()
-        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
+    def add_kv_cache(
+        self, tokenizer: TokenizerWrapper, prompt: str, cache: list[_BaseCache]
+    ):
+        tokenized_prompt = self.encode_prompt(tokenizer, prompt)
         self.prompts.append(tokenized_prompt)
         self.caches.append(deepcopy(cache))
-        self._access_counter += 1
-        self._last_used.append(self._access_counter)
-        logger.info(f"KV cache added: {len(tokenized_prompt)} tokens")
-
-    def update_kv_cache(
-        self,
-        index: int,
-        prompt: str,
-        cache: KVCacheType,
-    ):
-        """Update an existing cache entry in-place."""
-        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
-        self.prompts[index] = tokenized_prompt
-        self.caches[index] = deepcopy(cache)
-        self._access_counter += 1
-        self._last_used[index] = self._access_counter
-        logger.info(f"KV cache updated (index {index}): {len(tokenized_prompt)} tokens")
 
     def get_kv_cache(
         self,
         model: Model,
+        tokenizer: TokenizerWrapper,
+        sampler: Callable[[mx.array], mx.array],
         prompt: str,
-    ) -> tuple[KVCacheType, mx.array, int | None]:
-        """Get KV cache for prompt, returning remaining tokens to prefill.
-
-        Returns:
-            Tuple of (cache, remaining_tokens, matched_index) where:
-            - cache: KV cache to use for generation
-            - remaining_tokens: tokens that still need prefilling
-            - matched_index: index of the matched entry (None if no match)
-        """
-        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
+    ) -> list[_BaseCache]:
+        tokenized_prompt = self.encode_prompt(tokenizer, prompt)
         max_length = len(tokenized_prompt)
 
         best_snapshot_index, best_snapshot_length = None, 0
@@ -84,127 +42,63 @@ class KVPrefixCache:
             length = _get_prefix_length(tokenized_prompt, cached_prompt)
 
             if length == max_length:
-                # Exact match - cached prompt starts with our entire prompt
-                # Trim cache to prompt length - 1, return last token for stream_generate
-                prompt_cache = deepcopy(self.caches[i])
-                cached_length = _cache_length(self.caches[i])
-                tokens_to_trim = cached_length - (max_length - 1)
-                if tokens_to_trim > 0:
-                    trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
-                self._access_counter += 1
-                self._last_used[i] = self._access_counter
-                logger.info(f"KV cache exact match: {max_length} tokens (instant)")
-                return prompt_cache, tokenized_prompt[-1:], i
+                return self.caches[i]
 
             if length > best_snapshot_length:
                 best_snapshot_index, best_snapshot_length = i, length
 
         if best_snapshot_index is not None:
-            new_tokens = max_length - best_snapshot_length
-            logger.info(
-                f"KV cache prefix match: {best_snapshot_length}/{max_length} tokens "
-                f"(reusing {best_snapshot_length}, need to prefill {new_tokens})"
-            )
-
             prompt_cache = deepcopy(self.caches[best_snapshot_index])
-
-            # Trim removes tokens from the end, so we trim (cached_length - prefix_length) to keep the prefix
-            cached_length = _cache_length(self.caches[best_snapshot_index])
-            tokens_to_trim = cached_length - best_snapshot_length
-            if tokens_to_trim > 0:
-                trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
-
-            self._access_counter += 1
-            self._last_used[best_snapshot_index] = self._access_counter
-            remaining_tokens = tokenized_prompt[best_snapshot_length:]
-            return prompt_cache, remaining_tokens, best_snapshot_index
+            trim_prompt_cache(prompt_cache, max_length - best_snapshot_length)
+            tokenized_prompt = tokenized_prompt[best_snapshot_index:]
 
         else:
-            prompt_cache = make_kv_cache(model)
-            if len(self.prompts) == 0:
-                logger.info(f"KV cache empty, need to prefill {max_length} tokens")
-            else:
-                logger.info(
-                    f"KV cache no prefix match, need to prefill {max_length} tokens"
-                )
-
-            return prompt_cache, tokenized_prompt, None
-
-    def _evict_if_needed(self):
-        """Evict least recently used entries while memory pressure is high."""
-        if len(self.caches) == 0:
-            return
-
-        active: int = mx.metal.get_active_memory()
-        limit = int(mx.metal.device_info()["max_recommended_working_set_size"])
-        if active < limit * _MEMORY_THRESHOLD:
-            return
-
-        # Evict LRU entries until below threshold or only one entry left
-        while len(self.caches) > 0:
-            lru_index = self._last_used.index(min(self._last_used))
-            evicted_tokens = len(self.prompts[lru_index])
-            self.prompts.pop(lru_index)
-            self.caches.pop(lru_index)
-            self._last_used.pop(lru_index)
-            logger.info(
-                f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory pressure"
+            prompt_cache = make_kv_cache(
+                model,
+                # max_kv_size=MAX_KV_SIZE,
+                # keep=KEEP_KV_SIZE
             )
 
-            active = mx.metal.get_active_memory()
-            if active < limit * _MEMORY_THRESHOLD:
-                break
+        prefill(model, tokenizer, sampler, tokenized_prompt, prompt_cache)
 
+        return prompt_cache
 
-def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
-    """Encode a prompt string to token array.
-
-    For chat-templated prompts (which have their own structure markers like
-    <|im_user|>, <|im_middle|>, etc.), we should NOT add BOS/EOS tokens as
-    that would corrupt the prompt structure.
-    """
-    # Chat templates define their own structure - don't add BOS/EOS
-    tokenized_prompt = tokenizer.encode(prompt, add_special_tokens=False)
-    return mx.array(tokenized_prompt)
-
-
-def _cache_length(cache: KVCacheType) -> int:
-    """Get the number of tokens in a KV cache."""
-    # Use .offset attribute which all cache types have (len() not implemented in older QuantizedKVCache)
-    return max(c.offset for c in cache)  # type: ignore
+    def encode_prompt(self, tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
+        add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
+            tokenizer.bos_token
+        )
+        tokenized_prompt = tokenizer.encode(
+            prompt, add_special_tokens=add_special_tokens
+        )
+        return mx.array(tokenized_prompt)
 
 
 def _get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
-    """Find the length of the common prefix between two token arrays."""
-    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]))
+    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]), KEEP_KV_SIZE)
     if n == 0:
         return 0
 
-    equal = mx.equal(prompt[:n], cached_prompt[:n]).astype(mx.int32)
+    equal = (prompt[:n] == cached_prompt[:n]).astype(mx.int32)
     prefix_mask = mx.cumprod(equal)  # stays 1 until first mismatch, then 0 forever
     return int(mx.sum(prefix_mask).item())
 
 
-def make_kv_cache(
-    model: Model, max_kv_size: int | None = None, keep: int = 0
-) -> KVCacheType:
-    assert hasattr(model, "layers")
-
-    # TODO: Do this for all models
-    if hasattr(model, "make_cache") and isinstance(model, GptOssModel):
-        logger.info("Using MLX LM's make cache")
-        return model.make_cache()  # type: ignore
-
-    if max_kv_size is None:
-        if KV_CACHE_BITS is None:
-            logger.info("Using default KV cache")
-            return [KVCache() for _ in model.layers]
-        else:
-            logger.info("Using quantized KV cache")
-            return [
-                QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=KV_CACHE_BITS)
-                for _ in model.layers
-            ]
-    else:
-        logger.info(f"Using rotating KV cache with {max_kv_size=} with {keep=}")
-        return [RotatingKVCache(max_size=max_kv_size, keep=keep) for _ in model.layers]
+def prefill(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    sampler: Callable[[mx.array], mx.array],
+    prompt: mx.array,
+    cache: list[_BaseCache],
+) -> None:
+    for _ in stream_generate(
+        model=model,
+        tokenizer=tokenizer,
+        prompt=prompt,
+        max_tokens=0,
+        sampler=sampler,
+        prompt_cache=cache,
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
+    ):
+        pass

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

@@ -4,7 +4,7 @@
 KV_GROUP_SIZE: int | None = 32
 KV_BITS: int | None = None
 ATTENTION_KV_BITS: int | None = 4
-MAX_TOKENS: int = 32168
+MAX_TOKENS: int = 8192
 MAX_KV_SIZE: int | None = 3200
 KEEP_KV_SIZE: int | None = 1600
 QUANTIZE_MODEL_MODE: str | None = "affine"

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

@@ -1,12 +1,12 @@
-import time
 from typing import Any, Callable, Generator, cast, get_args
 
 import mlx.core as mx
 from mlx_lm.generate import stream_generate
-from mlx_lm.models.cache import trim_prompt_cache
+from mlx_lm.models.cache import KVCache
 from mlx_lm.sample_utils import make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
+# from exo.engines.mlx.cache import KVPrefixCache
 from exo.shared.types.api import (
     BenchChatCompletionTaskParams,
     ChatCompletionMessage,
@@ -14,78 +14,35 @@ from exo.shared.types.api import (
     GenerationStats,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType
 from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.cache import KVPrefixCache, encode_prompt, make_kv_cache
 from exo.worker.engines.mlx.constants import KV_BITS, KV_GROUP_SIZE, MAX_TOKENS
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
+    make_kv_cache,
     mx_barrier,
 )
 from exo.worker.runner.bootstrap import logger
 
 generation_stream = mx.new_stream(mx.default_device())
 
-_MIN_PREFIX_HIT_TO_UPDATE = 1000
 
-
-def prefill(
-    model: Model,
-    tokenizer: TokenizerWrapper,
-    sampler: Callable[[mx.array], mx.array],
-    prompt_tokens: mx.array,
-    cache: KVCacheType,
-) -> float:
-    """Prefill the KV cache with prompt tokens.
-
-    This runs the model over the prompt tokens to populate the cache,
-    then trims off the extra generated token.
-
-    Returns:
-        tokens_per_sec
-    """
-    num_tokens = len(prompt_tokens)
-    if num_tokens == 0:
-        return 0.0
-
-    logger.debug(f"Prefilling {num_tokens} tokens...")
-    start_time = time.perf_counter()
-
-    def progress_callback(processed: int, total: int) -> None:
-        elapsed = time.time() - start_time
-        tok_per_sec = processed / elapsed if elapsed > 0 else 0
-        logger.debug(
-            f"Prefill progress: {processed}/{total} tokens ({tok_per_sec:.1f} tok/s)"
-        )
-
-    # Use max_tokens=1 because max_tokens=0 does not work.
-    # We just throw away the generated token - we only care about filling the cache
-    for _ in stream_generate(
-        model=model,
-        tokenizer=tokenizer,
-        prompt=prompt_tokens,
-        max_tokens=1,
-        sampler=sampler,
-        prompt_cache=cache,
-        prefill_step_size=2048,
-        kv_group_size=KV_GROUP_SIZE,
-        kv_bits=KV_BITS,
-        prompt_progress_callback=progress_callback,
-    ):
-        break  # Stop after first iteration - cache is now filled
-    trim_prompt_cache(cast(list[Any], cache), 1)
-
-    elapsed = time.perf_counter() - start_time
-    tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0.0
-    logger.debug(
-        f"Prefill complete: {num_tokens} tokens in {elapsed:.2f}s "
-        f"({tokens_per_sec:.1f} tok/s)"
-    )
-    return tokens_per_sec
+def maybe_quantize_kv_cache(
+    prompt_cache: list[KVCache | Any],
+    quantized_kv_start: int,
+    kv_group_size: int,
+    kv_bits: int | None,
+) -> None:
+    if kv_bits is None:
+        return
+    for e, c in enumerate(prompt_cache):
+        if (
+            hasattr(c, "to_quantized") and c.offset >= quantized_kv_start  # type: ignore
+        ):
+            prompt_cache[e] = c.to_quantized(group_size=kv_group_size, bits=kv_bits)
 
 
 def warmup_inference(
@@ -163,7 +120,6 @@ def mlx_generate(
     tokenizer: TokenizerWrapper,
     task: ChatCompletionTaskParams,
     prompt: str,
-    kv_prefix_cache: KVPrefixCache | None = None,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
@@ -175,22 +131,7 @@ def mlx_generate(
     if task.seed is not None:
         mx.random.seed(task.seed)
 
-    # Do not use the prefix cache if we are trying to do benchmarks.
-    if is_bench:
-        kv_prefix_cache = None
-
-    # Use prefix cache if available, otherwise create fresh cache
-    prefix_hit_length = 0
-    matched_index: int | None = None
-    if kv_prefix_cache is None:
-        caches = make_kv_cache(model=model)
-        prompt_tokens = encode_prompt(tokenizer, prompt)
-    else:
-        caches, prompt_tokens, matched_index = kv_prefix_cache.get_kv_cache(
-            model, prompt
-        )
-        all_prompt_tokens = encode_prompt(tokenizer, prompt)
-        prefix_hit_length = len(all_prompt_tokens) - len(prompt_tokens)
+    caches = make_kv_cache(model=model)
 
     logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
     if is_bench:
@@ -203,19 +144,11 @@ def mlx_generate(
         top_p=task.top_p if task.top_p is not None else 1.0,
     )
 
-    # Prefill cache with all tokens except the last one
-    prefill_tps = prefill(model, tokenizer, sampler, prompt_tokens[:-1], caches)
-
-    # stream_generate starts from the last token
-    last_token = prompt_tokens[-1:]
-
     max_tokens = task.max_tokens or MAX_TOKENS
-    generated_text_parts: list[str] = []
-    generation_start_time = time.perf_counter()
     for out in stream_generate(
         model=model,
         tokenizer=tokenizer,
-        prompt=last_token,
+        prompt=prompt,
         max_tokens=max_tokens,
         sampler=sampler,
         logits_processors=logits_processors,
@@ -225,13 +158,12 @@ def mlx_generate(
         kv_group_size=KV_GROUP_SIZE,
         kv_bits=KV_BITS,
     ):
-        generated_text_parts.append(out.text)
         logger.info(out.text)
 
         stats: GenerationStats | None = None
         if out.finish_reason is not None:
             stats = GenerationStats(
-                prompt_tps=float(prefill_tps or out.prompt_tps),
+                prompt_tps=float(out.prompt_tps),
                 generation_tps=float(out.generation_tps),
                 prompt_tokens=int(out.prompt_tokens),
                 generation_tokens=int(out.generation_tokens),
@@ -253,26 +185,6 @@ def mlx_generate(
         )
 
         if out.finish_reason is not None:
-            # Log generation stats
-            generation_elapsed = time.perf_counter() - generation_start_time
-            generated_tokens = len(generated_text_parts)
-            generation_tps = (
-                generated_tokens / generation_elapsed if generation_elapsed > 0 else 0.0
-            )
-            logger.debug(
-                f"Generation complete: prefill {prompt_tokens} tokens @ "
-                f"{prefill_tps:.1f} tok/s, generated {generated_tokens} tokens @ "
-                f"{generation_tps:.1f} tok/s"
-            )
-            if kv_prefix_cache is not None:
-                full_prompt = prompt + "".join(generated_text_parts)
-                if (
-                    matched_index is not None
-                    and prefix_hit_length >= _MIN_PREFIX_HIT_TO_UPDATE
-                ):
-                    kv_prefix_cache.update_kv_cache(matched_index, full_prompt, caches)
-                else:
-                    kv_prefix_cache.add_kv_cache(full_prompt, caches)
             break
 
         # TODO: Do we want an mx_barrier?

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

@@ -18,12 +18,15 @@ try:
 except ImportError:
     pass  # transformers < 5.0 or bytes_to_unicode not available
 
-from mlx_lm.models.cache import KVCache
+from mlx_lm.models.cache import KVCache, QuantizedKVCache, RotatingKVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3Model
+from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.models.model_cards import ModelId
 from exo.worker.engines.mlx.constants import (
+    CACHE_GROUP_SIZE,
+    KV_CACHE_BITS,
     TRUST_REMOTE_CODE,
 )
 
@@ -38,7 +41,6 @@ import mlx.nn as nn
 from mlx_lm.utils import load_model
 from pydantic import RootModel
 
-from exo.download.download_utils import build_model_path
 from exo.shared.types.api import ChatCompletionMessageText
 from exo.shared.types.common import Host
 from exo.shared.types.memory import Memory
@@ -53,6 +55,7 @@ from exo.shared.types.worker.shards import (
     ShardMetadata,
     TensorShardMetadata,
 )
+from exo.worker.download.download_utils import build_model_path
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.auto_parallel import (
     TimeoutCallback,
@@ -402,11 +405,7 @@ def apply_chat_template(
                 continue
 
             message.content = "\n".join(c.text for c in message.content).strip()
-        if (
-            message.content is None
-            and message.thinking is None
-            and message.tool_calls is None
-        ):
+        if message.content is None and message.thinking is None:
             continue
 
         # Null values are not valid when applying templates in tokenizer
@@ -463,6 +462,31 @@ class NullKVCache(KVCache):
         raise NotImplementedError("We should not be setting a NullKVCache.")
 
 
+def make_kv_cache(
+    model: Model, max_kv_size: int | None = None, keep: int = 0
+) -> list[KVCache | RotatingKVCache | QuantizedKVCache]:
+    assert hasattr(model, "layers")
+
+    # TODO: Do this for all models
+    if hasattr(model, "make_cache") and isinstance(model, GptOssModel):
+        logger.info("Using MLX LM's make cache")
+        return model.make_cache()  # type: ignore
+
+    if max_kv_size is None:
+        if KV_CACHE_BITS is None:
+            logger.info("Using default KV cache")
+            return [KVCache() for _ in model.layers]
+        else:
+            logger.info("Using quantized KV cache")
+            return [
+                QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=KV_CACHE_BITS)
+                for _ in model.layers
+            ]
+    else:
+        logger.info(f"Using rotating KV cache with {max_kv_size=} with {keep=}")
+        return [RotatingKVCache(max_size=max_kv_size, keep=keep) for _ in model.layers]
+
+
 def mlx_force_oom(size: int = 40000) -> None:
     """
     Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.

src/exo/worker/main.py

@@ -1,9 +1,8 @@
 from datetime import datetime, timezone
 from random import random
-from typing import Iterator
 
 import anyio
-from anyio import CancelScope, create_task_group, fail_after
+from anyio import CancelScope, create_task_group, current_time, fail_after
 from anyio.abc import TaskGroup
 from loguru import logger
 
@@ -11,12 +10,7 @@ from exo.routing.connection_message import ConnectionMessage, ConnectionMessageT
 from exo.shared.apply import apply
 from exo.shared.models.model_cards import ModelId
 from exo.shared.types.api import ImageEditsInternalParams
-from exo.shared.types.commands import (
-    ForwarderCommand,
-    ForwarderDownloadCommand,
-    RequestEventLog,
-    StartDownload,
-)
+from exo.shared.types.commands import ForwarderCommand, RequestEventLog
 from exo.shared.types.common import CommandId, NodeId, SessionId
 from exo.shared.types.events import (
     Event,
@@ -24,6 +18,7 @@ from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
     InputChunkReceived,
+    NodeDownloadProgress,
     NodeGatheredInfo,
     TaskCreated,
     TaskStatusUpdated,
@@ -41,12 +36,23 @@ from exo.shared.types.tasks import (
     TaskStatus,
 )
 from exo.shared.types.topology import Connection, SocketConnection
+from exo.shared.types.worker.downloads import (
+    DownloadCompleted,
+    DownloadFailed,
+    DownloadOngoing,
+    DownloadPending,
+    DownloadProgress,
+)
 from exo.shared.types.worker.runners import RunnerId
+from exo.shared.types.worker.shards import ShardMetadata
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.event_buffer import OrderedBuffer
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
 from exo.utils.info_gatherer.net_profile import check_reachable
-from exo.utils.keyed_backoff import KeyedBackoff
+from exo.worker.download.download_utils import (
+    map_repo_download_progress_to_download_progress_data,
+)
+from exo.worker.download.shard_downloader import RepoDownloadProgress, ShardDownloader
 from exo.worker.plan import plan
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
 
@@ -56,6 +62,7 @@ class Worker:
         self,
         node_id: NodeId,
         session_id: SessionId,
+        shard_downloader: ShardDownloader,
         *,
         connection_message_receiver: Receiver[ConnectionMessage],
         global_event_receiver: Receiver[ForwarderEvent],
@@ -63,22 +70,23 @@ class Worker:
         # This is for requesting updates. It doesn't need to be a general command sender right now,
         # but I think it's the correct way to be thinking about commands
         command_sender: Sender[ForwarderCommand],
-        download_command_sender: Sender[ForwarderDownloadCommand],
-        event_index_counter: Iterator[int],
     ):
         self.node_id: NodeId = node_id
         self.session_id: SessionId = session_id
 
+        self.shard_downloader: ShardDownloader = shard_downloader
+        self._pending_downloads: dict[RunnerId, ShardMetadata] = {}
+
         self.global_event_receiver = global_event_receiver
         self.local_event_sender = local_event_sender
-        self.event_index_counter = event_index_counter
+        self.local_event_index = 0
         self.command_sender = command_sender
-        self.download_command_sender = download_command_sender
         self.connection_message_receiver = connection_message_receiver
         self.event_buffer = OrderedBuffer[Event]()
         self.out_for_delivery: dict[EventId, ForwarderEvent] = {}
 
         self.state: State = State()
+        self.download_status: dict[ModelId, DownloadProgress] = {}
         self.runners: dict[RunnerId, RunnerSupervisor] = {}
         self._tg: TaskGroup = create_task_group()
 
@@ -93,8 +101,6 @@ class Worker:
         self.input_chunk_buffer: dict[CommandId, dict[int, str]] = {}
         self.input_chunk_counts: dict[CommandId, int] = {}
 
-        self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0)
-
     async def run(self):
         logger.info("Starting Worker")
 
@@ -105,6 +111,7 @@ class Worker:
             tg.start_soon(info_gatherer.run)
             tg.start_soon(self._forward_info, info_recv)
             tg.start_soon(self.plan_step)
+            tg.start_soon(self._emit_existing_download_progress)
             tg.start_soon(self._connection_message_event_writer)
             tg.start_soon(self._resend_out_for_delivery)
             tg.start_soon(self._event_applier)
@@ -114,7 +121,6 @@ class Worker:
         # 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()
 
@@ -173,9 +179,11 @@ class Worker:
     async def plan_step(self):
         while True:
             await anyio.sleep(0.1)
+            # 3. based on the updated state, we plan & execute an operation.
             task: Task | None = plan(
                 self.node_id,
                 self.runners,
+                self.download_status,
                 self.state.downloads,
                 self.state.instances,
                 self.state.runners,
@@ -199,26 +207,42 @@ class Worker:
                         )
                     )
                 case DownloadModel(shard_metadata=shard):
-                    model_id = shard.model_card.model_id
-                    if not self._download_backoff.should_proceed(model_id):
-                        continue
-
-                    self._download_backoff.record_attempt(model_id)
-
-                    await self.download_command_sender.send(
-                        ForwarderDownloadCommand(
-                            origin=self.node_id,
-                            command=StartDownload(
-                                target_node_id=self.node_id,
-                                shard_metadata=shard,
-                            ),
+                    if shard.model_card.model_id not in self.download_status:
+                        progress = DownloadPending(
+                            shard_metadata=shard, node_id=self.node_id
+                        )
+                        self.download_status[shard.model_card.model_id] = progress
+                        await self.event_sender.send(
+                            NodeDownloadProgress(download_progress=progress)
+                        )
+                    initial_progress = (
+                        await self.shard_downloader.get_shard_download_status_for_shard(
+                            shard
                         )
                     )
-                    await self.event_sender.send(
-                        TaskStatusUpdated(
-                            task_id=task.task_id, task_status=TaskStatus.Running
+                    if initial_progress.status == "complete":
+                        progress = DownloadCompleted(
+                            shard_metadata=shard,
+                            node_id=self.node_id,
+                            total_bytes=initial_progress.total_bytes,
                         )
-                    )
+                        self.download_status[shard.model_card.model_id] = progress
+                        await self.event_sender.send(
+                            NodeDownloadProgress(download_progress=progress)
+                        )
+                        await self.event_sender.send(
+                            TaskStatusUpdated(
+                                task_id=task.task_id,
+                                task_status=TaskStatus.Complete,
+                            )
+                        )
+                    else:
+                        await self.event_sender.send(
+                            TaskStatusUpdated(
+                                task_id=task.task_id, task_status=TaskStatus.Running
+                            )
+                        )
+                        self._handle_shard_download_process(task, initial_progress)
                 case Shutdown(runner_id=runner_id):
                     try:
                         with fail_after(3):
@@ -363,17 +387,104 @@ class Worker:
         self._tg.start_soon(runner.run)
         return runner
 
+    def _handle_shard_download_process(
+        self,
+        task: DownloadModel,
+        initial_progress: RepoDownloadProgress,
+    ):
+        """Manages the shard download process with progress tracking."""
+        status = DownloadOngoing(
+            node_id=self.node_id,
+            shard_metadata=task.shard_metadata,
+            download_progress=map_repo_download_progress_to_download_progress_data(
+                initial_progress
+            ),
+        )
+        self.download_status[task.shard_metadata.model_card.model_id] = status
+        self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
+
+        last_progress_time = 0.0
+        throttle_interval_secs = 1.0
+
+        async def download_progress_callback(
+            shard: ShardMetadata, progress: RepoDownloadProgress
+        ) -> None:
+            nonlocal self
+            nonlocal last_progress_time
+            if progress.status == "complete":
+                status = DownloadCompleted(
+                    shard_metadata=shard,
+                    node_id=self.node_id,
+                    total_bytes=progress.total_bytes,
+                )
+                self.download_status[shard.model_card.model_id] = status
+                await self.event_sender.send(
+                    NodeDownloadProgress(download_progress=status)
+                )
+                await self.event_sender.send(
+                    TaskStatusUpdated(
+                        task_id=task.task_id, task_status=TaskStatus.Complete
+                    )
+                )
+            elif (
+                progress.status == "in_progress"
+                and current_time() - last_progress_time > throttle_interval_secs
+            ):
+                status = DownloadOngoing(
+                    node_id=self.node_id,
+                    shard_metadata=shard,
+                    download_progress=map_repo_download_progress_to_download_progress_data(
+                        progress
+                    ),
+                )
+                self.download_status[shard.model_card.model_id] = status
+                await self.event_sender.send(
+                    NodeDownloadProgress(download_progress=status)
+                )
+                last_progress_time = current_time()
+
+        self.shard_downloader.on_progress(download_progress_callback)
+
+        async def download_with_error_handling() -> None:
+            try:
+                await self.shard_downloader.ensure_shard(task.shard_metadata)
+            except Exception as e:
+                error_message = str(e)
+                logger.error(
+                    f"Download failed for {task.shard_metadata.model_card.model_id}: {error_message}"
+                )
+                failed_status = DownloadFailed(
+                    node_id=self.node_id,
+                    shard_metadata=task.shard_metadata,
+                    error_message=error_message,
+                )
+                self.download_status[task.shard_metadata.model_card.model_id] = (
+                    failed_status
+                )
+                await self.event_sender.send(
+                    NodeDownloadProgress(download_progress=failed_status)
+                )
+                await self.event_sender.send(
+                    TaskStatusUpdated(
+                        task_id=task.task_id, task_status=TaskStatus.Failed
+                    )
+                )
+
+        self._tg.start_soon(download_with_error_handling)
+
     async def _forward_events(self) -> None:
         with self.event_receiver as events:
             async for event in events:
-                idx = next(self.event_index_counter)
                 fe = ForwarderEvent(
-                    origin_idx=idx,
+                    origin_idx=self.local_event_index,
                     origin=self.node_id,
                     session=self.session_id,
                     event=event,
                 )
-                logger.debug(f"Worker published event {idx}: {str(event)[:100]}")
+                logger.debug(
+                    f"Worker published event {self.local_event_index}: {str(event)[:100]}"
+                )
+                self.local_event_index += 1
                 await self.local_event_sender.send(fe)
                 self.out_for_delivery[event.event_id] = fe
 
@@ -421,3 +532,42 @@ class Worker:
                     await self.event_sender.send(TopologyEdgeDeleted(conn=conn))
 
             await anyio.sleep(10)
+
+    async def _emit_existing_download_progress(self) -> None:
+        try:
+            while True:
+                logger.debug("Fetching and emitting existing download progress...")
+                async for (
+                    _,
+                    progress,
+                ) in self.shard_downloader.get_shard_download_status():
+                    if progress.status == "complete":
+                        status = DownloadCompleted(
+                            node_id=self.node_id,
+                            shard_metadata=progress.shard,
+                            total_bytes=progress.total_bytes,
+                        )
+                    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
+                            )
+                        else:
+                            status = DownloadOngoing(
+                                node_id=self.node_id,
+                                shard_metadata=progress.shard,
+                                download_progress=map_repo_download_progress_to_download_progress_data(
+                                    progress
+                                ),
+                            )
+                    else:
+                        continue
+
+                    self.download_status[progress.shard.model_card.model_id] = status
+                    await self.event_sender.send(
+                        NodeDownloadProgress(download_progress=status)
+                    )
+                logger.debug("Done emitting existing download progress.")
+                await anyio.sleep(5 * 60)  # 5 minutes
+        except Exception as e:
+            logger.error(f"Error emitting existing download progress: {e}")

src/exo/worker/plan.py

@@ -2,6 +2,7 @@
 
 from collections.abc import Mapping, Sequence
 
+from exo.shared.models.model_cards import ModelId
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.tasks import (
     ChatCompletion,
@@ -44,6 +45,9 @@ def plan(
     node_id: NodeId,
     # Runners is expected to be FRESH and so should not come from state
     runners: Mapping[RunnerId, RunnerSupervisor],
+    # DL_status is expected to be FRESH and so should not come from state
+    download_status: Mapping[ModelId, DownloadProgress],
+    # gdls is not expected to be fresh
     global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
     instances: Mapping[InstanceId, Instance],
     all_runners: Mapping[RunnerId, RunnerStatus],  # all global
@@ -55,7 +59,7 @@ def plan(
     return (
         _kill_runner(runners, all_runners, instances)
         or _create_runner(node_id, runners, instances)
-        or _model_needs_download(node_id, runners, global_download_status)
+        or _model_needs_download(runners, download_status)
         or _init_distributed_backend(runners, all_runners)
         or _load_model(runners, all_runners, global_download_status)
         or _ready_to_warmup(runners, all_runners)
@@ -111,15 +115,9 @@ def _create_runner(
 
 
 def _model_needs_download(
-    node_id: NodeId,
     runners: Mapping[RunnerId, RunnerSupervisor],
-    global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
+    download_status: Mapping[ModelId, DownloadProgress],
 ) -> DownloadModel | None:
-    local_downloads = global_download_status.get(node_id, [])
-    download_status = {
-        dp.shard_metadata.model_card.model_id: dp for dp in local_downloads
-    }
-
     for runner in runners.values():
         model_id = runner.bound_instance.bound_shard.model_card.model_id
         if isinstance(runner.status, RunnerIdle) and (

src/exo/worker/runner/bootstrap.py

@@ -7,7 +7,6 @@ from exo.shared.types.tasks import Task
 from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
-from exo.worker.tests.patches import load_null_model
 
 logger: "loguru.Logger" = loguru.logger
 
@@ -17,8 +16,6 @@ def entrypoint(
     event_sender: MpSender[Event],
     task_receiver: MpReceiver[Task],
     _logger: "loguru.Logger",
-    *,
-    _load_null_models: bool = False,
 ) -> None:
     fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
     if fast_synch_override == "on" or (
@@ -32,13 +29,6 @@ def entrypoint(
     else:
         os.environ["MLX_METAL_FAST_SYNCH"] = "0"
 
-    p = None
-    if _load_null_models:
-        from unittest.mock import patch
-
-        p = patch("mlx_lm.utils.load_model", new=load_null_model)
-        p.start()
-
     global logger
     logger = _logger
 
@@ -62,8 +52,6 @@ def entrypoint(
             )
         )
     finally:
-        if p is not None:
-            p.stop()
         try:
             event_sender.close()
             task_receiver.close()

src/exo/worker/runner/runner.py

@@ -70,7 +70,6 @@ from exo.worker.engines.image import (
     warmup_image_generator,
 )
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
@@ -104,7 +103,6 @@ def main(
     model: Model | DistributedImageModel | None = None
     tokenizer = None
     group = None
-    kv_prefix_cache: KVPrefixCache | None = None
 
     current_status: RunnerStatus = RunnerIdle()
     logger.info("runner created")
@@ -163,8 +161,6 @@ def main(
                         logger.info(
                             f"model has_tool_calling={tokenizer.has_tool_calling}"
                         )
-                        kv_prefix_cache = KVPrefixCache(tokenizer)
-
                     elif (
                         ModelTask.TextToImage in shard_metadata.model_card.tasks
                         or ModelTask.ImageToImage in shard_metadata.model_card.tasks
@@ -174,6 +170,7 @@ def main(
                         raise ValueError(
                             f"Unknown model task(s): {shard_metadata.model_card.tasks}"
                         )
+
                     current_status = RunnerLoaded()
                     logger.info("runner loaded")
                 case StartWarmup() if isinstance(current_status, RunnerLoaded):
@@ -241,9 +238,12 @@ def main(
                             tokenizer=tokenizer,
                             task=task_params,
                             prompt=prompt,
-                            kv_prefix_cache=kv_prefix_cache,
                         )
 
+                        # GPT-OSS specific parsing to match other model formats.
+                        if isinstance(model, GptOssModel):
+                            mlx_generator = parse_gpt_oss(mlx_generator)
+
                         # For other thinking models (GLM, etc.), check if we need to
                         # prepend the thinking tag that was consumed by the chat template
                         if detect_thinking_prompt_suffix(prompt, tokenizer):
@@ -257,16 +257,10 @@ def main(
                             patch_kimi_tokenizer(tokenizer)
 
                         # GLM models need patched parser (upstream has bug with None regex match)
-                        elif "glm" in shard_metadata.model_card.model_id.lower():
+                        if "glm" in shard_metadata.model_card.model_id.lower():
                             patch_glm_tokenizer(tokenizer)
 
-                        # GPT-OSS specific parsing to match other model formats.
-                        elif isinstance(model, GptOssModel):
-                            mlx_generator = parse_gpt_oss(mlx_generator)
-
-                        if tokenizer.has_tool_calling and not isinstance(
-                            model, GptOssModel
-                        ):
+                        if tokenizer.has_tool_calling:
                             assert tokenizer.tool_call_start
                             assert tokenizer.tool_call_end
                             assert tokenizer.tool_parser  # pyright: ignore[reportAny]
@@ -495,10 +489,9 @@ def get_gpt_oss_encoding():
 
 
 def filter_kimi_tokens(
-    responses: Generator[GenerationResponse | ToolCallResponse],
+    responses: Generator[GenerationResponse],
 ) -> Generator[GenerationResponse]:
     for resp in responses:
-        assert isinstance(resp, GenerationResponse)
         if (
             resp.text == "<|tool_calls_section_begin|>"
             or resp.text == "<|tool_calls_section_end|>"
@@ -508,44 +501,17 @@ def filter_kimi_tokens(
 
 
 def parse_gpt_oss(
-    responses: Generator[GenerationResponse | ToolCallResponse],
-) -> Generator[GenerationResponse | ToolCallResponse]:
+    responses: Generator[GenerationResponse],
+) -> Generator[GenerationResponse]:
     encoding = get_gpt_oss_encoding()
     stream = StreamableParser(encoding, role=Role.ASSISTANT)
     thinking = False
-    current_tool_name: str | None = None
-    tool_arg_parts: list[str] = []
 
     for response in responses:
-        assert isinstance(response, GenerationResponse)
         stream.process(response.token)
 
         delta = stream.last_content_delta
         ch = stream.current_channel
-        recipient = stream.current_recipient
-
-        if recipient != current_tool_name:
-            if current_tool_name is not None:
-                prefix = "functions."
-                if current_tool_name.startswith(prefix):
-                    current_tool_name = current_tool_name[len(prefix) :]
-                yield ToolCallResponse(
-                    tool_calls=[
-                        ToolCallItem(
-                            name=current_tool_name,
-                            arguments="".join(tool_arg_parts).strip(),
-                        )
-                    ]
-                )
-                tool_arg_parts = []
-                break
-            current_tool_name = recipient
-
-        # If inside a tool call, accumulate arguments
-        if current_tool_name is not None:
-            if delta:
-                tool_arg_parts.append(delta)
-            continue
 
         if ch == "analysis" and not thinking:
             thinking = True
@@ -562,12 +528,13 @@ def parse_gpt_oss(
             if thinking:
                 yield response.model_copy(update={"text": "</think>"})
             yield response
+            break
 
 
 def parse_thinking_models(
-    responses: Generator[GenerationResponse | ToolCallResponse],
+    responses: Generator[GenerationResponse],
     tokenizer: TokenizerWrapper,
-) -> Generator[GenerationResponse | ToolCallResponse]:
+) -> Generator[GenerationResponse]:
     """
     For models that inject thinking tags in the prompt (like GLM-4.7),
     prepend the thinking tag to the output stream so the frontend
@@ -575,9 +542,6 @@ def parse_thinking_models(
     """
     first = True
     for response in responses:
-        if isinstance(response, ToolCallResponse):
-            yield response
-            continue
         if first:
             first = False
             yield response.model_copy(
@@ -648,7 +612,7 @@ def _process_image_response(
         command_id=command_id,
         model_id=shard_metadata.model_card.model_id,
         event_sender=event_sender,
-        image_index=response.image_index,
+        image_index=response.partial_index if is_partial else image_index,
         is_partial=is_partial,
         partial_index=response.partial_index if is_partial else None,
         total_partials=response.total_partials if is_partial else None,
@@ -658,7 +622,7 @@ def _process_image_response(
 
 
 def parse_tool_calls(
-    responses: Generator[GenerationResponse | ToolCallResponse],
+    responses: Generator[GenerationResponse],
     tool_call_start: str,
     tool_call_end: str,
     tool_parser: Callable[[str], dict[str, Any] | list[dict[str, Any]]],
@@ -666,7 +630,6 @@ def parse_tool_calls(
     in_tool_call = False
     tool_call_text_parts: list[str] = []
     for response in responses:
-        assert isinstance(response, GenerationResponse)
         # assumption: the tool call start is one token
         if response.text == tool_call_start:
             in_tool_call = True

src/exo/worker/tests/patches.py

@@ -1,50 +0,0 @@
-# type: ignore
-
-import importlib
-import json
-from pathlib import Path
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    from exo.worker.engines.mlx import Model
-
-
-def load_null_model(path: Path, **_: object) -> "tuple[Model, dict[str, Any]]":
-    with open(path / "config.json", "r") as f:
-        cfg = json.load(f)
-        model, args = _get_classes(cfg)
-        model = model(args.from_dict(cfg))
-        return model, cfg
-
-
-def _get_classes(config: dict):
-    """
-    Retrieve the model and model args classes based on the configuration.
-
-    Args:
-        config (dict): The model configuration.
-
-    Returns:
-        A tuple containing the Model class and the ModelArgs class.
-    """
-    model_type = config["model_type"]
-    model_type = MODEL_REMAPPING.get(model_type, model_type)
-    try:
-        arch = importlib.import_module(f"mlx_lm.models.{model_type}")
-    except ImportError:
-        msg = f"Model type {model_type} not supported."
-        raise ValueError(msg) from None
-
-    return arch.Model, arch.ModelArgs
-
-
-MODEL_REMAPPING = {
-    "mistral": "llama",
-    "llava": "mistral3",
-    "phi-msft": "phixtral",
-    "falcon_mamba": "mamba",
-    "kimi_k2": "deepseek_v3",
-    "qwen2_5_vl": "qwen2_vl",
-    "minimax_m2": "minimax",
-    "iquestcoder": "llama",
-}

src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py

@@ -1,545 +0,0 @@
-# type: ignore
-import time
-from typing import cast
-from unittest.mock import patch
-
-import mlx.core as mx
-import pytest
-from mlx_lm.models.cache import KVCache
-from mlx_lm.sample_utils import make_sampler
-
-from exo.shared.types.api import ChatCompletionMessage
-from exo.shared.types.common import ModelId
-from exo.shared.types.tasks import ChatCompletionTaskParams
-from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.cache import (
-    KVPrefixCache,
-    _cache_length,
-    _get_prefix_length,
-    encode_prompt,
-    make_kv_cache,
-)
-from exo.worker.engines.mlx.generator.generate import mlx_generate, prefill
-from exo.worker.engines.mlx.utils_mlx import apply_chat_template
-from exo.worker.tests.unittests.test_mlx.conftest import (
-    DEFAULT_GPT_OSS_CONFIG,
-    DEFAULT_GPT_OSS_MODEL_ID,
-)
-
-
-def _check_model_exists() -> bool:
-    return DEFAULT_GPT_OSS_CONFIG.model_path.exists()
-
-
-class TestGetPrefixLength:
-    def test_identical_arrays(self):
-        a = mx.array([1, 2, 3, 4, 5])
-        b = mx.array([1, 2, 3, 4, 5])
-        assert _get_prefix_length(a, b) == 5
-
-    def test_no_common_prefix(self):
-        a = mx.array([1, 2, 3])
-        b = mx.array([4, 5, 6])
-        assert _get_prefix_length(a, b) == 0
-
-    def test_partial_prefix(self):
-        a = mx.array([1, 2, 3, 4, 5])
-        b = mx.array([1, 2, 3, 7, 8])
-        assert _get_prefix_length(a, b) == 3
-
-    def test_prompt_longer_than_cached(self):
-        a = mx.array([1, 2, 3, 4, 5])
-        b = mx.array([1, 2, 3])
-        assert _get_prefix_length(a, b) == 3
-
-    def test_cached_longer_than_prompt(self):
-        a = mx.array([1, 2, 3])
-        b = mx.array([1, 2, 3, 4, 5])
-        assert _get_prefix_length(a, b) == 3
-
-    def test_single_token_match(self):
-        a = mx.array([1, 2, 3])
-        b = mx.array([1, 5, 6])
-        assert _get_prefix_length(a, b) == 1
-
-    def test_empty_prompt(self):
-        a = mx.array([]).astype(mx.int32)
-        b = mx.array([1, 2, 3])
-        assert _get_prefix_length(a, b) == 0
-
-    def test_empty_cached(self):
-        a = mx.array([1, 2, 3])
-        b = mx.array([]).astype(mx.int32)
-        assert _get_prefix_length(a, b) == 0
-
-    def test_both_empty(self):
-        a = mx.array([]).astype(mx.int32)
-        b = mx.array([]).astype(mx.int32)
-        assert _get_prefix_length(a, b) == 0
-
-
-class TestKVPrefix:
-    @pytest.fixture
-    def mock_tokenizer(self):
-        """Create a minimal mock tokenizer for tests that don't need real tokenization."""
-        from unittest.mock import MagicMock
-
-        tokenizer = MagicMock()
-        tokenizer.encode.return_value = [1, 2, 3]
-        return tokenizer
-
-    def test_starts_empty(self, mock_tokenizer):
-        cache = KVPrefixCache(mock_tokenizer)
-        assert len(cache.prompts) == 0
-        assert len(cache.caches) == 0
-
-    def test_clear_empties_cache(self, mock_tokenizer):
-        cache = KVPrefixCache(mock_tokenizer)
-        cache.prompts.append(mx.array([1, 2, 3]))
-        cache.caches.append([KVCache()])
-        cache.clear()
-        assert len(cache.prompts) == 0
-        assert len(cache.caches) == 0
-
-    def test_clear_on_empty_cache(self, mock_tokenizer):
-        cache = KVPrefixCache(mock_tokenizer)
-        cache.clear()
-        assert len(cache.prompts) == 0
-
-
-def _load_gpt_oss() -> tuple[Model, object]:
-    from mlx_lm.utils import load_model
-
-    from exo.worker.engines.mlx.utils_mlx import load_tokenizer_for_model_id
-
-    model_path = DEFAULT_GPT_OSS_CONFIG.model_path
-    model_id = ModelId(DEFAULT_GPT_OSS_MODEL_ID)
-
-    model, _ = load_model(model_path, lazy=False)
-    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
-    return cast(Model, model), tokenizer
-
-
-@pytest.mark.slow
-@pytest.mark.skipif(
-    not _check_model_exists(),
-    reason=f"GPT-OSS model not found at {DEFAULT_GPT_OSS_CONFIG.model_path}",
-)
-class TestKVPrefixCacheWithModel:
-    @pytest.fixture(scope="class")
-    def model_and_tokenizer(self):
-        model, tokenizer = _load_gpt_oss()
-        return model, tokenizer
-
-    def test_prefill_populates_cache(self, model_and_tokenizer):
-        model, tokenizer = model_and_tokenizer
-
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Hello!!")],
-            max_tokens=1,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        tokens = encode_prompt(tokenizer, prompt)
-        cache = make_kv_cache(model)
-
-        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-
-        # Cache should now hold the prompt tokens
-        assert _cache_length(cache) == len(tokens)
-
-    def test_add_and_get_exact_match(self, model_and_tokenizer):
-        model, tokenizer = model_and_tokenizer
-
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Test exact")],
-            max_tokens=1,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        tokens = encode_prompt(tokenizer, prompt)
-        cache = make_kv_cache(model)
-
-        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        kv_prefix_cache.add_kv_cache(prompt, cache)
-
-        assert len(kv_prefix_cache.prompts) == 1
-        stored_length = _cache_length(kv_prefix_cache.caches[0])
-        assert stored_length > 0
-
-        # Retrieve with same prompt: exact match
-        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
-            model, prompt
-        )
-        assert matched_index == 0
-
-        # Exact match returns only last token
-        assert len(remaining_tokens) == 1
-        assert mx.array_equal(remaining_tokens, tokens[-1:])
-
-    def test_add_and_get_prefix_match(self, model_and_tokenizer):
-        """get_kv_cache with a longer prompt sharing prefix should return partial match."""
-        model, tokenizer = model_and_tokenizer
-
-        short_task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Hi")],
-            max_tokens=1,
-        )
-        short_prompt = apply_chat_template(tokenizer, short_task)
-        short_tokens = encode_prompt(tokenizer, short_prompt)
-        cache = make_kv_cache(model)
-
-        prefill(model, tokenizer, make_sampler(0.0), short_tokens, cache)
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        kv_prefix_cache.add_kv_cache(short_prompt, cache)
-
-        # Query with longer prompt that shares the chat template prefix
-        long_task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[
-                ChatCompletionMessage(role="user", content="Hi there, how are you?")
-            ],
-            max_tokens=1,
-        )
-        long_prompt = apply_chat_template(tokenizer, long_task)
-        long_tokens = encode_prompt(tokenizer, long_prompt)
-
-        # The prompts share a prefix (chat template preamble + "Hi")
-        expected_prefix = _get_prefix_length(long_tokens, short_tokens)
-        assert expected_prefix > 0, (
-            "Prompts should share a prefix from the chat template"
-        )
-
-        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
-            model, long_prompt
-        )
-        assert matched_index == 0
-
-        # remaining_tokens should be the suffix after the shared prefix
-        assert len(remaining_tokens) == len(long_tokens) - expected_prefix
-        assert mx.array_equal(remaining_tokens, long_tokens[expected_prefix:])
-
-    def test_stored_cache_not_mutated_after_get_and_generation(
-        self, model_and_tokenizer
-    ):
-        """Getting a cache and then mutating it (as generation does) must not corrupt stored cache."""
-        model, tokenizer = model_and_tokenizer
-
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Mutation test")],
-            max_tokens=1,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        tokens = encode_prompt(tokenizer, prompt)
-        cache = make_kv_cache(model)
-
-        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        kv_prefix_cache.add_kv_cache(prompt, cache)
-
-        stored_length = _cache_length(kv_prefix_cache.caches[0])
-
-        # Get cache and mutate it (simulating what generation does)
-        result_cache, _, matched_index = kv_prefix_cache.get_kv_cache(model, prompt)
-        assert matched_index == 0
-
-        # Simulate generation: feed many additional tokens through the cache
-        head_dim = result_cache[0].keys.shape[-1]
-        num_heads = result_cache[0].keys.shape[1]
-        extra_keys = mx.random.normal((1, num_heads, 50, head_dim))
-        extra_values = mx.random.normal((1, num_heads, 50, head_dim))
-        for layer_cache in result_cache:
-            layer_cache.update_and_fetch(extra_keys, extra_values)
-        mx.eval([c.keys for c in result_cache])
-
-        # Stored cache must be unchanged
-        assert _cache_length(kv_prefix_cache.caches[0]) == stored_length
-
-    def test_stored_cache_survives_repeated_get_mutate_cycles(
-        self, model_and_tokenizer
-    ):
-        """Multiple get+mutate cycles (like repeated user requests) must not corrupt cache."""
-        model, tokenizer = model_and_tokenizer
-
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Repeat test")],
-            max_tokens=1,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        tokens = encode_prompt(tokenizer, prompt)
-        cache = make_kv_cache(model)
-
-        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        kv_prefix_cache.add_kv_cache(prompt, cache)
-
-        stored_length = _cache_length(kv_prefix_cache.caches[0])
-
-        for i in range(3):
-            result_cache, _, _ = kv_prefix_cache.get_kv_cache(model, prompt)
-
-            head_dim = result_cache[0].keys.shape[-1]
-            num_heads = result_cache[0].keys.shape[1]
-            extra = mx.random.normal((1, num_heads, 30, head_dim))
-            for layer_cache in result_cache:
-                layer_cache.update_and_fetch(extra, extra)
-            mx.eval([c.keys for c in result_cache])
-
-            assert _cache_length(kv_prefix_cache.caches[0]) == stored_length, (
-                f"Failed on loop {i}"
-            )
-
-    def test_mlx_generate_populates_cache(self, model_and_tokenizer):
-        """mlx_generate should save the cache after generation completes."""
-        model, tokenizer = model_and_tokenizer
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Hello")],
-            max_tokens=5,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        prompt_tokens = encode_prompt(tokenizer, prompt)
-
-        # Consume the entire generator so the cache-saving code after yield runs
-        generated_tokens = 0
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
-            prompt=prompt,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            generated_tokens += 1
-
-        assert len(kv_prefix_cache.prompts) == 1
-        assert len(kv_prefix_cache.caches) == 1
-        # Cache should contain prompt + generated tokens
-        expected_length = len(prompt_tokens) + generated_tokens
-        assert _cache_length(kv_prefix_cache.caches[0]) == expected_length
-
-    def test_mlx_generate_second_call_gets_prefix_hit(self, model_and_tokenizer):
-        """Second mlx_generate call with same prompt should get a prefix hit from stored cache."""
-        model, tokenizer = model_and_tokenizer
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Reuse test")],
-            max_tokens=5,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        prompt_tokens = encode_prompt(tokenizer, prompt)
-
-        # First generation populates cache
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
-            prompt=prompt,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            pass
-
-        assert len(kv_prefix_cache.prompts) == 1
-
-        # Second call should find a prefix match (the stored cache contains
-        # prompt + generated tokens, which shares the prompt prefix)
-        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
-            model, prompt
-        )
-        # The stored cache is longer than the prompt (it includes generated tokens),
-        # so this is a prefix match where our prompt is fully contained
-        assert matched_index == 0
-        # Exact match: remaining_tokens is just the last token
-        assert len(remaining_tokens) == 1
-        assert mx.array_equal(remaining_tokens, prompt_tokens[-1:])
-
-    def test_mlx_generate_long_prompt_updates_cache_in_place(self, model_and_tokenizer):
-        """With a prompt > 1000 tokens, second generation should update the cache entry in-place."""
-        model, tokenizer = model_and_tokenizer
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-
-        # Build a long user message (> 1000 tokens) to exceed _MIN_PREFIX_HIT_TO_UPDATE
-        base_text = "The quick brown fox jumps over the lazy dog. "
-        base_tokens = tokenizer.encode(base_text)
-        repeats = (1200 // len(base_tokens)) + 2
-        long_content = base_text * repeats
-
-        task1 = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content=long_content)],
-            max_tokens=5,
-        )
-        prompt1 = apply_chat_template(tokenizer, task1)
-        prompt1_tokens = encode_prompt(tokenizer, prompt1)
-        assert len(prompt1_tokens) > 1000, (
-            "Prompt must exceed _MIN_PREFIX_HIT_TO_UPDATE"
-        )
-
-        # First generation populates the cache (must prefill all tokens)
-        t0 = time.perf_counter()
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task1,
-            prompt=prompt1,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            pass
-        first_gen_time = time.perf_counter() - t0
-
-        assert len(kv_prefix_cache.prompts) == 1
-        first_cache_length = _cache_length(kv_prefix_cache.caches[0])
-
-        # Second generation: same long prompt + extra content (simulating multi-turn)
-        task2 = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[
-                ChatCompletionMessage(role="user", content=long_content),
-                ChatCompletionMessage(role="assistant", content="Sure, I can help."),
-                ChatCompletionMessage(role="user", content="Tell me more."),
-            ],
-            max_tokens=5,
-        )
-        prompt2 = apply_chat_template(tokenizer, task2)
-        prompt2_tokens = encode_prompt(tokenizer, prompt2)
-
-        # Verify the prompts share a long prefix
-        prefix_len = _get_prefix_length(prompt2_tokens, prompt1_tokens)
-        assert prefix_len > 1000, "Prompts must share > 1000 token prefix"
-
-        # Second generation should reuse the cached prefix (only prefill new tokens)
-        t0 = time.perf_counter()
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task2,
-            prompt=prompt2,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            pass
-        second_gen_time = time.perf_counter() - t0
-
-        # Second generation should be significantly faster due to prefix cache hit - hopefully not flaky
-        assert second_gen_time < first_gen_time * 0.5, (
-            f"Expected prefix cache speedup: "
-            f"first={first_gen_time:.2f}s, second={second_gen_time:.2f}s"
-        )
-
-        # With prefix_hit > 1000, should update in-place (not add a second entry)
-        assert len(kv_prefix_cache.prompts) == 1
-        # Updated cache should be longer (prompt2 + generated > prompt1 + generated)
-        updated_cache_length = _cache_length(kv_prefix_cache.caches[0])
-        assert updated_cache_length > first_cache_length
-
-    def test_mlx_generate_stored_cache_not_mutated(self, model_and_tokenizer):
-        """After mlx_generate saves a cache, a second generation must not corrupt the stored copy."""
-        model, tokenizer = model_and_tokenizer
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Immutable test")],
-            max_tokens=5,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-
-        # First generation populates cache
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
-            prompt=prompt,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            pass
-
-        first_cache_length = _cache_length(kv_prefix_cache.caches[0])
-
-        # Second generation gets the cache and mutates it during generation
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
-            prompt=prompt,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            pass
-
-        # The first stored cache must not have been mutated by the second generation
-        assert _cache_length(kv_prefix_cache.caches[0]) == first_cache_length
-
-    def test_evicts_lru_entry_under_memory_pressure(self, model_and_tokenizer):
-        """Under memory pressure, adding a new cache entry evicts the least recently used one."""
-        model, tokenizer = model_and_tokenizer
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-
-        # Add three cache entries with different prompts
-        prompts = ["First entry", "Second entry", "Third entry"]
-        for i, content in enumerate(prompts):
-            task = ChatCompletionTaskParams(
-                model=DEFAULT_GPT_OSS_MODEL_ID,
-                messages=[ChatCompletionMessage(role="user", content=content)],
-                max_tokens=1,
-            )
-            prompt = apply_chat_template(tokenizer, task)
-            tokens = encode_prompt(tokenizer, prompt)
-            cache = make_kv_cache(model)
-            prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-            kv_prefix_cache.add_kv_cache(prompt, cache)
-            # Stagger _last_used so LRU order is deterministic
-            kv_prefix_cache._last_used[i] = float(i)
-
-        assert len(kv_prefix_cache.prompts) == 3
-
-        # Access the third entry to make it most recently used
-        kv_prefix_cache._last_used[2] = 100.0
-        # Entry 0 (_last_used=0.0) is LRU, entry 1 (_last_used=1.0) is next
-
-        # Simulate memory pressure: active memory exceeds threshold
-        fake_limit = 1000
-        fake_active = int(fake_limit * 0.90)  # Above _MEMORY_THRESHOLD (0.85)
-
-        with (
-            patch(
-                "exo.worker.engines.mlx.cache.mx.metal.get_active_memory",
-                return_value=fake_active,
-            ),
-            patch(
-                "exo.worker.engines.mlx.cache.mx.metal.device_info",
-                return_value={"max_recommended_working_set_size": fake_limit},
-            ),
-        ):
-            # Trigger eviction by adding a new entry
-            task = ChatCompletionTaskParams(
-                model=DEFAULT_GPT_OSS_MODEL_ID,
-                messages=[ChatCompletionMessage(role="user", content="New entry")],
-                max_tokens=1,
-            )
-            prompt = apply_chat_template(tokenizer, task)
-            tokens = encode_prompt(tokenizer, prompt)
-            cache = make_kv_cache(model)
-            prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-            kv_prefix_cache.add_kv_cache(prompt, cache)
-
-        # LRU entries should have been evicted (entries 0, 1, 2 in order of _last_used)
-        # Since fake_active stays above threshold after each eviction (we don't change it),
-        # all old entries get evicted, leaving only the newly added one
-        assert len(kv_prefix_cache.prompts) == 1
-        # The surviving entry should be the newly added one
-        new_tokens = encode_prompt(tokenizer, prompt)
-        assert _get_prefix_length(kv_prefix_cache.prompts[0], new_tokens) == len(
-            new_tokens
-        )

src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py

@@ -11,12 +11,12 @@ from pathlib import Path
 
 import pytest
 
-from exo.download.download_utils import (
+from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
+from exo.worker.download.download_utils import (
     download_file_with_retry,
     ensure_models_dir,
     fetch_file_list_with_cache,
 )
-from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
 from exo.worker.engines.mlx.utils_mlx import (
     get_eos_token_ids_for_model,
     load_tokenizer_for_model_id,

src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py

@@ -1,5 +1,5 @@
 import exo.worker.plan as plan_mod
-from exo.shared.types.common import NodeId
+from exo.shared.types.common import ModelId, NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.tasks import LoadModel
 from exo.shared.types.worker.downloads import DownloadCompleted, DownloadProgress
@@ -45,9 +45,13 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
     instances = {INSTANCE_1_ID: instance}
     all_runners = {RUNNER_1_ID: RunnerIdle()}
 
+    # No entry for this shard -> should trigger DownloadModel
+    download_status: dict[ModelId, DownloadProgress] = {}
+
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status=download_status,
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -88,6 +92,14 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
         RUNNER_2_ID: RunnerConnected(),
     }
 
+    # Local node has already marked its shard as downloaded (not actually used by _load_model)
+    local_download_status = {
+        MODEL_A_ID: DownloadCompleted(
+            shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
+        )
+    }
+
+    # Global view has completed downloads for both nodes
     global_download_status = {
         NODE_A: [
             DownloadCompleted(
@@ -104,6 +116,7 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,
@@ -135,26 +148,30 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
     instances = {INSTANCE_1_ID: instance}
     all_runners = {RUNNER_1_ID: RunnerIdle()}
 
-    # Global state shows shard is downloaded for NODE_A
+    # Local status claims the shard is downloaded already
+    local_download_status = {
+        MODEL_A_ID: DownloadCompleted(
+            shard_metadata=shard, node_id=NODE_A, total_bytes=Memory()
+        )
+    }
+
+    # Global view hasn't caught up yet (no completed shards recorded for NODE_A)
     global_download_status: dict[NodeId, list[DownloadProgress]] = {
-        NODE_A: [
-            DownloadCompleted(
-                shard_metadata=shard, node_id=NODE_A, total_bytes=Memory()
-            )
-        ],
+        NODE_A: [],
         NODE_B: [],
     }
 
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,
         tasks={},
     )
 
-    assert not isinstance(result, plan_mod.DownloadModel)
+    assert result is None
 
 
 def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
@@ -185,6 +202,12 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
         RUNNER_2_ID: RunnerConnected(),
     }
 
+    # Only NODE_A's shard is recorded as downloaded globally
+    local_download_status = {
+        MODEL_A_ID: DownloadCompleted(
+            shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
+        )
+    }
     global_download_status = {
         NODE_A: [
             DownloadCompleted(
@@ -197,6 +220,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,
@@ -221,6 +245,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,

src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py

@@ -47,7 +47,8 @@ def test_plan_kills_runner_when_instance_missing():
 
     result = plan_mod.plan(
         node_id=NODE_A,
-        runners=runners,  # type: ignore[arg-type]
+        runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -86,7 +87,8 @@ def test_plan_kills_runner_when_sibling_failed():
 
     result = plan_mod.plan(
         node_id=NODE_A,
-        runners=runners,  # type: ignore[arg-type]
+        runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -118,6 +120,7 @@ def test_plan_creates_runner_when_missing_for_node():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -155,7 +158,8 @@ def test_plan_does_not_create_runner_when_supervisor_already_present():
 
     result = plan_mod.plan(
         node_id=NODE_A,
-        runners=runners,  # type: ignore[arg-type]
+        runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -185,6 +189,7 @@ def test_plan_does_not_create_runner_for_unassigned_node():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,

src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py

@@ -65,6 +65,7 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -112,6 +113,7 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: [], NODE_B: []},
         instances=instances,
         all_runners=all_runners,
@@ -156,6 +158,7 @@ def test_plan_does_not_forward_tasks_for_other_instances():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -218,6 +221,7 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: [], NODE_B: []},
         instances=instances,
         all_runners=all_runners,
@@ -257,6 +261,7 @@ def test_plan_returns_none_when_nothing_to_do():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: [], NODE_B: []},
         instances=instances,
         all_runners=all_runners,

src/exo/worker/tests/unittests/test_plan/test_warmup.py

@@ -57,6 +57,7 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
     result = plan_mod.plan(
         node_id=NODE_B,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -98,6 +99,7 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -138,6 +140,7 @@ def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warmin
     result = plan_mod.plan(
         node_id=NODE_B,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: [], NODE_B: []},
         instances=instances,
         all_runners=all_runners,
@@ -182,6 +185,7 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -198,6 +202,7 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -241,6 +246,7 @@ def test_plan_starts_warmup_for_connecting_rank_after_others_warming():
     result = plan_mod.plan(
         node_id=NODE_B,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_B: []},
         instances=instances,
         all_runners=all_runners,
@@ -283,6 +289,7 @@ def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warmi
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: [], NODE_B: []},
         instances=instances,
         all_runners=all_runners,
@@ -324,6 +331,7 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
+        download_status={},
         global_download_status={NODE_A: [], NODE_B: []},
         instances=instances,
         all_runners=all_runners,

tests/headless_runner.py

@@ -1,6 +1,7 @@
 import multiprocessing as mp
 import socket
 import time
+import typing
 
 import anyio
 from fastapi import FastAPI
@@ -15,7 +16,7 @@ from exo.shared.models.model_cards import MODEL_CARDS, ModelId
 from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
 from exo.shared.types.commands import CommandId
 from exo.shared.types.common import Host, NodeId
-from exo.shared.types.events import Event, RunnerStatusUpdated
+from exo.shared.types.events import Event
 from exo.shared.types.tasks import (
     ChatCompletion,
     ConnectToGroup,
@@ -31,17 +32,22 @@ from exo.shared.types.worker.instances import (
     MlxJacclInstance,
     MlxRingInstance,
 )
-from exo.shared.types.worker.runners import RunnerFailed, RunnerId, ShardAssignments
+from exo.shared.types.worker.runners import RunnerId, ShardAssignments
 from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
 from exo.utils.channels import MpReceiver, MpSender, channel, mp_channel
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
+from exo.worker.download.impl_shard_downloader import (
+    build_full_shard,
+    exo_shard_downloader,
+)
 from exo.worker.runner.bootstrap import entrypoint
 
-MODEL_CARDS = {"haha": MODEL_CARDS["qwen3-coder-480b-a35b-8bit"]}
 
 class Tests(BaseModel):
     # list[hostname, ip addr]
     devs: list[list[str]]
+    model_id: str
+    kind: typing.Literal["init", "warmup", "inference"]
 
 
 mp.set_start_method("spawn", force=True)
@@ -50,14 +56,16 @@ logger_setup(None)
 
 async def main():
     logger.info("starting cool server majig")
+    await assert_downloads()
     cfg = Config()
-    cfg.bind = "0.0.0.0:8000"
+    cfg.bind = "0.0.0.0:52415"
     # nb: shared.logging needs updating if any of this changes
     cfg.accesslog = "-"
     cfg.errorlog = "-"
     cfg.logger_class = InterceptLogger
     app = FastAPI()
-    app.post("/run_test")(run_test)
+    app.post("/ring")(ring_backend)
+    app.post("/jaccl")(jaccl_backend)
     app.post("/tb_detection")(tb_detection)
     shutdown = anyio.Event()
     await serve(
@@ -79,7 +87,28 @@ async def tb_detection():
         return recv.collect()
 
 
-async def run_test(test: Tests):
+async def assert_downloads():
+    sd = exo_shard_downloader()
+    # await sd.ensure_shard(await build_full_shard(MODEL_CARDS["qwen3-0.6b"].model_id))
+    await sd.ensure_shard(
+        await build_full_shard(MODEL_CARDS["llama-3.1-8b-bf16"].model_id)
+    )
+    await sd.ensure_shard(await build_full_shard(MODEL_CARDS["qwen3-30b"].model_id))
+    await sd.ensure_shard(
+        await build_full_shard(MODEL_CARDS["gpt-oss-120b-MXFP4-Q8"].model_id)
+    )
+    await sd.ensure_shard(
+        await build_full_shard(MODEL_CARDS["gpt-oss-20b-4bit"].model_id)
+    )
+    await sd.ensure_shard(
+        await build_full_shard(MODEL_CARDS["glm-4.7-8bit-gs32"].model_id)
+    )
+    await sd.ensure_shard(
+        await build_full_shard(MODEL_CARDS["minimax-m2.1-8bit"].model_id)
+    )
+
+
+async def ring_backend(test: Tests):
     iid = InstanceId(str(hash(str(test.devs))))
     weird_hn = socket.gethostname()
     for dev in test.devs:
@@ -88,30 +117,10 @@ async def run_test(test: Tests):
             break
     else:
         raise ValueError(f"{weird_hn} not in {test.devs}")
-
-    async def run():
-        for card in MODEL_CARDS.values():
-            for instance in (
-                ring_instance(test, card.model_id, iid, hn),
-                jaccl_instance(test, card.model_id, iid),
-            ):
-                recv = await execute_test(test, instance, hn)
-
-                with recv:
-                    try:
-                        async for item in recv:
-                            yield item.model_dump_json() + "\n"
-                            if isinstance(item, RunnerStatusUpdated) and isinstance(
-                                item.runner_status, RunnerFailed
-                            ):
-                                return
-                    except anyio.ClosedResourceError:
-                        pass
-
-    return StreamingResponse(run())
+    return await execute_test(test, ring_instance(test, iid, hn), hn)
 
 
-def ring_instance(test: Tests, model_id: ModelId, iid: InstanceId, hn: str) -> Instance:
+def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
     hbn = [Host(ip="i dont care", port=52416) for _ in test.devs]
     world_size = len(test.devs)
     for i in range(world_size):
@@ -126,13 +135,13 @@ def ring_instance(test: Tests, model_id: ModelId, iid: InstanceId, hn: str) -> I
     else:
         raise ValueError(f"{hn} not in {test.devs}")
 
-    card = next(card for card in MODEL_CARDS.values() if card.model_id == model_id)
+    card = MODEL_CARDS[test.model_id]
     instance = MlxRingInstance(
         instance_id=iid,
         ephemeral_port=52416,
         hosts_by_node={NodeId(hn): hbn},
         shard_assignments=ShardAssignments(
-            model_id=model_id,
+            model_id=ModelId(test.model_id),
             node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs},
             runner_to_shard={
                 RunnerId(test.devs[i][0]): PipelineShardMetadata(
@@ -154,7 +163,7 @@ def ring_instance(test: Tests, model_id: ModelId, iid: InstanceId, hn: str) -> I
     return instance
 
 
-async def execute_test(test: Tests, instance: Instance, hn: str) -> MpReceiver[Event]:
+async def execute_test(test: Tests, instance: Instance, hn: str):
     world_size = len(test.devs)
     iid = InstanceId(str(hash(str(test.devs))))
     _handle, recv, send = new_runner(instance, hn)
@@ -162,33 +171,60 @@ async def execute_test(test: Tests, instance: Instance, hn: str) -> MpReceiver[E
         send.send(ConnectToGroup(instance_id=iid))
     send.send(LoadModel(instance_id=iid))
 
-    for card in MODEL_CARDS.values():
-        send.send(StartWarmup(instance_id=iid))
-        send.send(
-            ChatCompletion(
-                task_params=ChatCompletionTaskParams(
-                    model=card.model_id,
-                    messages=[
-                        ChatCompletionMessage(
-                            role="system", content="You are a helpful assistant"
-                        ),
-                        ChatCompletionMessage(
-                            role="user", content="What is the capital of France?"
-                        ),
-                    ],
-                ),
-                command_id=CommandId("yo"),
-                instance_id=iid,
+    match test.kind:
+        case "init":
+            pass
+        case "warmup":
+            send.send(StartWarmup(instance_id=iid))
+        case "inference":
+            send.send(StartWarmup(instance_id=iid))
+            send.send(
+                ChatCompletion(
+                    task_params=ChatCompletionTaskParams(
+                        model=test.model_id,
+                        messages=[
+                            ChatCompletionMessage(
+                                role="system", content="You are a helpful assistant"
+                            ),
+                            ChatCompletionMessage(
+                                role="user", content="What is the capital of France?"
+                            ),
+                        ],
+                    ),
+                    command_id=CommandId("yo"),
+                    instance_id=iid,
+                )
             )
-        )
 
     send.send(Shutdown(runner_id=RunnerId(hn), instance_id=iid))
 
-    return recv
+    async def map_recv():
+        with recv:
+            try:
+                async for item in recv:
+                    yield item.model_dump_json() + "\n"
+            except anyio.ClosedResourceError:
+                pass
+
+    ret = StreamingResponse(map_recv())
+    ret._pls_dont_gc = _handle  # type: ignore
+    return ret
 
 
-def jaccl_instance(test: Tests, model_id: ModelId, iid: InstanceId):
-    card = next(card for card in MODEL_CARDS.values() if card.model_id == model_id)
+async def jaccl_backend(test: Tests):
+    iid = InstanceId(str(hash(str(test.devs))))
+    weird_hn = socket.gethostname()
+    for dev in test.devs:
+        if weird_hn.startswith(dev[0]) or dev[0].startswith(weird_hn):
+            hn = dev[0]
+            break
+    else:
+        raise ValueError(f"{weird_hn} not in {test.devs}")
+    return await execute_test(test, jaccl_instance(test, iid), hn)
+
+
+def jaccl_instance(test: Tests, iid: InstanceId):
+    card = MODEL_CARDS[test.model_id]
     world_size = len(test.devs)
 
     return MlxJacclInstance(
@@ -199,7 +235,7 @@ def jaccl_instance(test: Tests, model_id: ModelId, iid: InstanceId):
             NodeId(host[0]): test.devs[0][1] + ":52416" for host in test.devs
         },
         shard_assignments=ShardAssignments(
-            model_id=model_id,
+            model_id=ModelId(test.model_id),
             node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs},
             runner_to_shard={
                 RunnerId(test.devs[i][0]): TensorShardMetadata(
@@ -234,7 +270,6 @@ def new_runner(
             task_recv,
             logger,
         ),
-        kwargs={"_load_null_models": True},
     )
     runner_process._pls_dont_gc = (ev_send, task_recv)  # type: ignore
     runner_process.start()

tests/start_distributed_test.sh

@@ -6,8 +6,19 @@ query() {
   tailscale status | awk -v find="$1" '$2 == find { print $1 }'
 }
 
-if [[ $# -lt 1 ]]; then
-  echo "USAGE: $0 [host1] [host2] ..."
+if [[ $# -lt 2 ]]; then
+  echo "USAGE: $0 <test kind> [host1] [host2] ..."
+  exit 1
+fi
+
+
+kind=$1
+shift
+
+test_kinds="ring jaccl"
+
+if ! echo "$test_kinds" | grep -q "$kind"; then
+  printf "%s is not a known test kind.\nCurrent test kinds are %s" "$kind" "$test_kinds"
   exit 1
 fi
 
@@ -23,12 +34,23 @@ done
 devs_raw=$(printf "[\"%s\", \"%s\"], " "${weaved[@]}")
 devs="[${devs_raw%, }]"
 
-for i in "${!ips[@]}"; do  
-  { 
-    curl -sN \
-      -X POST "http://${ips[$i]}:8000/run_test" \
-      -H "Content-Type: application/json" -d "{\"devs\": ${devs}}" \
-    2>&1 | sed "s/^/\n${hostnames[$i]}@${ips[$i]}: /" || echo "curl to ${hostnames[$i]} failed" && exit 1
-  } &
+model_ids=("qwen3-30b" "gpt-oss-120b-MXFP4-Q8" "kimi-k2-thinking")
+
+for model_id in "${model_ids[@]}"; do
+  for i in "${!ips[@]}"; do  
+    { 
+      req="{
+        \"model_id\": \"${model_id}\",
+        \"devs\": ${devs},
+        \"kind\": \"inference\"
+       }"
+      echo "req $req"
+      curl -sN \
+        -X POST "http://${ips[$i]}:52415/${kind}" \
+        -H "Content-Type: application/json" -d "$req" \
+      2>&1 | sed "s/^/\n${hostnames[$i]}@${ips[$i]}: /" || echo "curl to ${hostnames[$i]} failed" && exit 1
+    } &
+  done
+  wait
 done
-wait
+

uv.lock

@@ -412,10 +412,10 @@ requires-dist = [
     { name = "huggingface-hub", specifier = ">=0.33.4" },
     { name = "hypercorn", specifier = ">=0.18.0" },
     { name = "loguru", specifier = ">=0.7.3" },
-    { name = "mflux", specifier = "==0.15.4" },
+    { name = "mflux", specifier = ">=0.14.2" },
     { name = "mlx", marker = "sys_platform == 'darwin'", specifier = "==0.30.3" },
     { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.3" },
-    { name = "mlx-lm", specifier = "==0.30.5" },
+    { name = "mlx-lm", git = "https://github.com/AlexCheema/mlx-lm.git?rev=fix-transformers-5.0.0rc2" },
     { name = "openai-harmony", specifier = ">=0.0.8" },
     { name = "pillow", specifier = ">=11.0,<12.0" },
     { name = "psutil", specifier = ">=7.0.0" },
@@ -458,6 +458,16 @@ dev = [
     { name = "pytest-asyncio", specifier = ">=1.0.0" },
 ]
 
+[[package]]
+name = "tomlkit"
+version = "0.14.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c3/af/14b24e41977adb296d6bd1fb59402cf7d60ce364f90c890bd2ec65c43b5a/tomlkit-0.14.0.tar.gz", hash = "sha256:cf00efca415dbd57575befb1f6634c4f42d2d87dbba376128adb42c121b87064", size = 187167 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b5/11/87d6d29fb5d237229d67973a6c9e06e048f01cf4994dee194ab0ea841814/tomlkit-0.14.0-py3-none-any.whl", hash = "sha256:592064ed85b40fa213469f81ac584f67a4f2992509a7c3ea2d632208623a3680", size = 39310 },
+]
+
+
 [[package]]
 name = "fastapi"
 version = "0.128.0"
@@ -987,7 +997,7 @@ wheels = [
 
 [[package]]
 name = "mflux"
-version = "0.15.4"
+version = "0.15.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1013,9 +1023,9 @@ dependencies = [
     { name = "twine", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "urllib3", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/a6/f8/95322db7a865e4df6bad108b1c99aa7fbe211aac3f298f3ad696c2744a39/mflux-0.15.4.tar.gz", hash = "sha256:138e1aedae86e13eafeb8faec017945fcdcca42c3234daabcd81a83c9a202ace", size = 741228, upload-time = "2026-01-20T15:39:26.807Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/23/c5/dd12e16714702255d89b7ccc6f217c405a9fdcf2af950a2236892c50a219/mflux-0.15.3.tar.gz", hash = "sha256:e32ea66a81aad4f77eea2415b17c27fc3d9ce662a842565c62871ff570f4ef2f", size = 740701, upload-time = "2026-01-19T22:54:59.066Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/8e/be/81cf4ce2d1933b9b210c028a05ac95e958008c0d43e377a5f2757b7f2d4d/mflux-0.15.4-py3-none-any.whl", hash = "sha256:f04d9b1d7c5cd67880f483ab29fb2097648a25459eef9c5ee6480fad46de5e82", size = 987644, upload-time = "2026-01-20T15:39:24.817Z" },
+    { url = "https://files.pythonhosted.org/packages/cf/9f/a673ee12877a0943a4059c51b5beb6cf909c92f25384365cf8beeb475159/mflux-0.15.3-py3-none-any.whl", hash = "sha256:631cfcc038f27e9bd0ff76c25c2bc7373562b8f64cf0ce961fc268a246fa699e", size = 987270, upload-time = "2026-01-19T22:54:57.155Z" },
 ]
 
 [[package]]
@@ -1072,8 +1082,8 @@ wheels = [
 
 [[package]]
 name = "mlx-lm"
-version = "0.30.5"
-source = { registry = "https://pypi.org/simple" }
+version = "0.30.4"
+source = { git = "https://github.com/AlexCheema/mlx-lm.git?rev=fix-transformers-5.0.0rc2#a5daf2b894f31793dfaef0fdf9bc3ed683176ad6" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", marker = "sys_platform == 'darwin'" },
@@ -1083,10 +1093,6 @@ dependencies = [
     { name = "sentencepiece", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/0b/90/4469d9f75f196e6255f59a89441abe0079925d30a001462e1c1c4bc4e6a1/mlx_lm-0.30.5.tar.gz", hash = "sha256:9e6cb258c65b766c6af25cb90958aef40acab67139f05839eef19864cb3154f6", size = 262367, upload-time = "2026-01-25T15:29:30.125Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/89/ba/66db6e1e5f1ef506655b562932f6bd8f72600116d5f31f92d71c1f200b3f/mlx_lm-0.30.5-py3-none-any.whl", hash = "sha256:a80bc8e3efdebe81813b0f6eb403fb66a7a15071e256f4e7102ada986acb75bb", size = 366716, upload-time = "2026-01-25T15:29:28.29Z" },
-]
 
 [[package]]
 name = "mlx-metal"
@@ -2221,15 +2227,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/44/6f/7120676b6d73228c96e17f1f794d8ab046fc910d781c8d151120c3f1569e/toml-0.10.2-py2.py3-none-any.whl", hash = "sha256:806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b", size = 16588, upload-time = "2020-11-01T01:40:20.672Z" },
 ]
 
-[[package]]
-name = "tomlkit"
-version = "0.14.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/c3/af/14b24e41977adb296d6bd1fb59402cf7d60ce364f90c890bd2ec65c43b5a/tomlkit-0.14.0.tar.gz", hash = "sha256:cf00efca415dbd57575befb1f6634c4f42d2d87dbba376128adb42c121b87064", size = 187167, upload-time = "2026-01-13T01:14:53.304Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/b5/11/87d6d29fb5d237229d67973a6c9e06e048f01cf4994dee194ab0ea841814/tomlkit-0.14.0-py3-none-any.whl", hash = "sha256:592064ed85b40fa213469f81ac584f67a4f2992509a7c3ea2d632208623a3680", size = 39310, upload-time = "2026-01-13T01:14:51.965Z" },
-]
-
 [[package]]
 name = "torch"
 version = "2.9.1"
@@ -2285,7 +2282,7 @@ wheels = [
 
 [[package]]
 name = "transformers"
-version = "5.0.0rc3"
+version = "5.0.0rc2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -2300,9 +2297,9 @@ dependencies = [
     { name = "tqdm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "typer-slim", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/3f/a3/7c116a8d85f69ea7749cf4c2df79e64c35d028e5fc7ea0168f299d03b8c7/transformers-5.0.0rc3.tar.gz", hash = "sha256:a0315b92b7e087617ade42ec9e6e92ee7620541cc5d6a3331886c52cbe306f5c", size = 8388520, upload-time = "2026-01-14T16:49:02.952Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/94/e2/86b1bd5264272953370a5e50a91da38d7a53a87c5faf3fd3ff62d7353879/transformers-5.0.0rc2.tar.gz", hash = "sha256:9f2fa5e132433dd7eb910dc224b32de0baf758f3b6ffc918dbb632e0af85c07a", size = 8362532, upload-time = "2026-01-07T16:58:02.603Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/1e/f2/ae2b8968764253bdf38a48dee3c299b8d0bedf7c8ffbe3449fca9bd95338/transformers-5.0.0rc3-py3-none-any.whl", hash = "sha256:383fad27f4f73092d330e45fae384681e5c8521e1dc1cf6cb1a297780e68bf2d", size = 10107087, upload-time = "2026-01-14T16:48:59.393Z" },
+    { url = "https://files.pythonhosted.org/packages/b4/eb/9526a77354a2126f5b220f4792dc8494d573773c098dac6a5ad1fc7a5f17/transformers-5.0.0rc2-py3-none-any.whl", hash = "sha256:f8f2a14060ab11f20a0eec39d827af54c1589c327c5799d82808ae3f4167418a", size = 10067329, upload-time = "2026-01-07T16:57:59.617Z" },
 ]
 
 [[package]]