lets try the other thing

nix mlx compilation for better portability
Update README.md with some changes from release 1.0.61 (#1157 )
2026-01-21 12:30:22 -05:00 · 2026-01-19 16:55:54 +00:00 · 2026-01-19 16:51:29 +00:00 · 2026-01-19 16:43:18 +00:00 · 2026-01-19 16:19:37 +00:00
12 changed files with 229 additions and 260 deletions
--- a/.github/actions/typecheck/action.yml
+++ b/.github/actions/typecheck/action.yml
@@ -0,0 +1,12 @@
+name: Type Check
+
+description: "Run type checker"
+
+runs:
+  using: "composite"
+  steps:
+    - name: Run type checker
+      run: |
+        nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just sync
+        nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just check
+      shell: bash
--- a/.github/workflows/build-app.yml
+++ b/.github/workflows/build-app.yml
@@ -161,6 +161,17 @@ jobs:
      - name: Install Homebrew packages
        run: brew install just awscli macmon

+      - name: Install UV
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
+          cache-dependency-glob: uv.lock
+
+      - name: Setup Python
+        run: |
+          uv python install
+          uv sync --locked
+
      - name: Install Nix
        uses: cachix/install-nix-action@v31
        with:
@@ -172,6 +183,12 @@ jobs:
          name: exo
          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"

+      - name: Build dashboard
+        run: |
+          DASHBOARD_OUT=$(nix build .#dashboard --print-build-logs --no-link --print-out-paths)
+          mkdir -p dashboard/build
+          cp -r "$DASHBOARD_OUT"/* dashboard/build/
+
      - name: Install Sparkle CLI
        run: |
          CLI_URL="${SPARKLE_CLI_URL:-https://github.com/sparkle-project/Sparkle/releases/download/${SPARKLE_VERSION}/Sparkle-${SPARKLE_VERSION}.tar.xz}"
@@ -227,10 +244,7 @@ jobs:
      # ============================================================

      - name: Build PyInstaller bundle
-        run: |
-          PYINSTALLER_OUT=$(nix build .#exo-pyinstaller --print-build-logs --no-link --print-out-paths)
-          mkdir -p dist
-          cp -r "$PYINSTALLER_OUT" dist/exo
+        run: uv run pyinstaller packaging/pyinstaller/exo.spec

      - name: Build Swift app
        env:
--- a/.github/workflows/pipeline.yml
+++ b/.github/workflows/pipeline.yml
@@ -26,14 +26,73 @@ jobs:
          name: exo
          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"

-      - name: Load nix develop environment
-        run: nix run github:nicknovitski/nix-develop/v1
+      - name: Configure git user
+        run: |
+          git config --local user.email "github-actions@users.noreply.github.com"
+          git config --local user.name  "github-actions bot"
+        shell: bash

-      - name: Sync dependencies
-        run: uv sync --all-packages
+      - name: Pull LFS files
+        run: |
+          echo "Pulling Git LFS files..."
+          git lfs pull
+        shell: bash

-      - name: Run type checker
-        run: uv run basedpyright --project pyproject.toml
+      - name: Setup Nix Environment
+        run: |
+          echo "Checking for nix installation..."
+          
+          # Check if nix binary exists directly
+          if [ -f /nix/var/nix/profiles/default/bin/nix ]; then
+            echo "Found nix binary at /nix/var/nix/profiles/default/bin/nix"
+            export PATH="/nix/var/nix/profiles/default/bin:$PATH"
+            echo "PATH=$PATH" >> $GITHUB_ENV
+            nix --version
+          elif [ -f /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh ]; then
+            echo "Found nix profile script, sourcing..."
+            source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh
+            nix --version
+          elif command -v nix >/dev/null 2>&1; then
+            echo "Nix already in PATH"
+            nix --version
+          else
+            echo "Nix not found. Debugging info:"
+            echo "Contents of /nix/var/nix/profiles/default/:"
+            ls -la /nix/var/nix/profiles/default/ 2>/dev/null || echo "Directory not found"
+            echo "Contents of /nix/var/nix/profiles/default/bin/:"
+            ls -la /nix/var/nix/profiles/default/bin/ 2>/dev/null || echo "Directory not found"
+            exit 1
+          fi
+        shell: bash
+
+      - name: Configure basedpyright include for local MLX
+        run: |
+          RUNNER_LABELS='${{ toJSON(runner.labels) }}'
+          if echo "$RUNNER_LABELS" | grep -q "local_mlx"; then
+            if [ -d "/Users/Shared/mlx" ]; then
+              echo "Updating [tool.basedpyright].include to use /Users/Shared/mlx"
+              awk '
+                BEGIN { in=0 }
+                /^\[tool\.basedpyright\]/ { in=1; print; next }
+                in && /^\[/ { in=0 }  # next section
+                in && /^[ \t]*include[ \t]*=/ {
+                  print "include = [\"/Users/Shared/mlx\"]"
+                  next
+                }
+                { print }
+              ' pyproject.toml > pyproject.toml.tmp && mv pyproject.toml.tmp pyproject.toml
+
+              echo "New [tool.basedpyright] section:"
+              sed -n '/^\[tool\.basedpyright\]/,/^\[/p' pyproject.toml | sed '$d' || true
+            else
+              echo "local_mlx tag present but /Users/Shared/mlx not found; leaving pyproject unchanged."
+            fi
+          else
+            echo "Runner does not have 'local_mlx' tag; leaving pyproject unchanged."
+          fi
+        shell: bash
+
+      - uses: ./.github/actions/typecheck

  nix:
    name: Build and check (${{ matrix.system }})
--- a/README.md
+++ b/README.md
@@ -42,7 +42,7 @@ exo includes a built-in dashboard for managing your cluster and chatting with mo
  <summary>Qwen3-235B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA</summary>
  <img src="docs/benchmarks/jeffgeerling/mac-studio-cluster-ai-full-1-qwen3-235b.jpeg" alt="Benchmark - Qwen3-235B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA" width="80%" />
  <p>
-    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
+    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
  </p>
 </details>

@@ -50,7 +50,7 @@ exo includes a built-in dashboard for managing your cluster and chatting with mo
  <summary>DeepSeek v3.1 671B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA</summary>
  <img src="docs/benchmarks/jeffgeerling/mac-studio-cluster-ai-full-2-deepseek-3.1-671b.jpeg" alt="Benchmark - DeepSeek v3.1 671B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA" width="80%" />
  <p>
-    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
+    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
  </p>
 </details>

@@ -58,7 +58,7 @@ exo includes a built-in dashboard for managing your cluster and chatting with mo
  <summary>Kimi K2 Thinking (native 4-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA</summary>
  <img src="docs/benchmarks/jeffgeerling/mac-studio-cluster-ai-full-3-kimi-k2-thinking.jpeg" alt="Benchmark - Kimi K2 Thinking (native 4-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA" width="80%" />
  <p>
-    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
+    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
  </p>
 </details>

@@ -163,6 +163,24 @@ This starts the exo dashboard and API at http://localhost:52415/

 **Important note for Linux users:** Currently, exo runs on CPU on Linux. GPU support for Linux platforms is under development. If you'd like to see support for your specific Linux hardware, please [search for existing feature requests](https://github.com/exo-explore/exo/issues) or create a new one.

+**Configuration Options:**
+
+- `--no-worker`: Run exo without the worker component. Useful for coordinator-only nodes that handle networking and orchestration but don't execute inference tasks. This is helpful for machines without sufficient GPU resources but with good network connectivity.
+
+  ```bash
+  uv run exo --no-worker
+  ```
+
+**File Locations (Linux):**
+
+exo follows the [XDG Base Directory Specification](https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html) on Linux:
+
+- **Configuration files**: `~/.config/exo/` (or `$XDG_CONFIG_HOME/exo/`)
+- **Data files**: `~/.local/share/exo/` (or `$XDG_DATA_HOME/exo/`)
+- **Cache files**: `~/.cache/exo/` (or `$XDG_CACHE_HOME/exo/`)
+
+You can override these locations by setting the corresponding XDG environment variables.
+
 ### macOS App

 exo ships a macOS app that runs in the background on your Mac.
@@ -175,6 +193,19 @@ Download the latest build here: [EXO-latest.dmg](https://assets.exolabs.net/EXO-

 The app will ask for permission to modify system settings and install a new Network profile. Improvements to this are being worked on.

+**Custom Namespace for Cluster Isolation:**
+
+The macOS app includes a custom namespace feature that allows you to isolate your exo cluster from others on the same network. This is configured through the `EXO_LIBP2P_NAMESPACE` setting:
+
+- **Use cases**:
+  - Running multiple separate exo clusters on the same network
+  - Isolating development/testing clusters from production clusters
+  - Preventing accidental cluster joining
+
+- **Configuration**: Access this setting in the app's Advanced settings (or set the `EXO_LIBP2P_NAMESPACE` environment variable when running from source)
+
+The namespace is logged on startup for debugging purposes.
+
 #### Uninstalling the macOS App

 The recommended way to uninstall is through the app itself: click the menu bar icon → Advanced → Uninstall. This cleanly removes all system components.
@@ -321,6 +352,52 @@ For further details, see:

 ---

+## Benchmarking
+
+The `exo-bench` tool measures model prefill and token generation speed across different placement configurations. This helps you optimize model performance and validate improvements.
+
+**Prerequisites:**
+- Nodes should be running with `uv run exo` before benchmarking
+- The tool uses the `/bench/chat/completions` endpoint
+
+**Basic usage:**
+
+```bash
+uv run bench/exo_bench.py \
+  --model llama-3.2-1b \
+  --pp 128,256,512 \
+  --tg 128,256
+```
+
+**Key parameters:**
+
+- `--model`: Model to benchmark (short ID or HuggingFace ID)
+- `--pp`: Prompt size hints (comma-separated integers)
+- `--tg`: Generation lengths (comma-separated integers)
+- `--max-nodes`: Limit placements to N nodes (default: 4)
+- `--instance-meta`: Filter by `ring`, `jaccl`, or `both` (default: both)
+- `--sharding`: Filter by `pipeline`, `tensor`, or `both` (default: both)
+- `--repeat`: Number of repetitions per configuration (default: 1)
+- `--warmup`: Warmup runs per placement (default: 0)
+- `--json-out`: Output file for results (default: bench/results.json)
+
+**Example with filters:**
+
+```bash
+uv run bench/exo_bench.py \
+  --model llama-3.2-1b \
+  --pp 128,512 \
+  --tg 128 \
+  --max-nodes 2 \
+  --sharding tensor \
+  --repeat 3 \
+  --json-out my-results.json
+```
+
+The tool outputs performance metrics including prompt tokens per second (prompt_tps), generation tokens per second (generation_tps), and peak memory usage for each configuration.
+
+---
+
 ## Hardware Accelerator Support

 On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working on extending hardware accelerator support. If you'd like support for a new hardware platform, please [search for an existing feature request](https://github.com/exo-explore/exo/issues) and add a thumbs up so we know what hardware is important to the community.
@@ -329,4 +406,4 @@ On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working

 ## Contributing

-See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
+See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
--- a/flake.lock
+++ b/flake.lock
@@ -21,9 +21,7 @@
          "nixpkgs"
        ],
        "purescript-overlay": "purescript-overlay",
-        "pyproject-nix": [
-          "pyproject-nix"
-        ]
+        "pyproject-nix": "pyproject-nix"
      },
      "locked": {
        "lastModified": 1765953015,
@@ -151,44 +149,19 @@
        "type": "github"
      }
    },
-    "pyproject-build-systems": {
-      "inputs": {
-        "nixpkgs": [
-          "nixpkgs"
-        ],
-        "pyproject-nix": [
-          "pyproject-nix"
-        ],
-        "uv2nix": [
-          "uv2nix"
-        ]
-      },
-      "locked": {
-        "lastModified": 1763662255,
-        "narHash": "sha256-4bocaOyLa3AfiS8KrWjZQYu+IAta05u3gYZzZ6zXbT0=",
-        "owner": "pyproject-nix",
-        "repo": "build-system-pkgs",
-        "rev": "042904167604c681a090c07eb6967b4dd4dae88c",
-        "type": "github"
-      },
-      "original": {
-        "owner": "pyproject-nix",
-        "repo": "build-system-pkgs",
-        "type": "github"
-      }
-    },
    "pyproject-nix": {
      "inputs": {
        "nixpkgs": [
+          "dream2nix",
          "nixpkgs"
        ]
      },
      "locked": {
-        "lastModified": 1764134915,
-        "narHash": "sha256-xaKvtPx6YAnA3HQVp5LwyYG1MaN4LLehpQI8xEdBvBY=",
+        "lastModified": 1763017646,
+        "narHash": "sha256-Z+R2lveIp6Skn1VPH3taQIuMhABg1IizJd8oVdmdHsQ=",
        "owner": "pyproject-nix",
        "repo": "pyproject.nix",
-        "rev": "2c8df1383b32e5443c921f61224b198a2282a657",
+        "rev": "47bd6f296502842643078d66128f7b5e5370790c",
        "type": "github"
      },
      "original": {
@@ -205,10 +178,7 @@
        "flake-parts": "flake-parts",
        "nixpkgs": "nixpkgs",
        "nixpkgs-swift": "nixpkgs-swift",
-        "pyproject-build-systems": "pyproject-build-systems",
-        "pyproject-nix": "pyproject-nix",
-        "treefmt-nix": "treefmt-nix",
-        "uv2nix": "uv2nix"
+        "treefmt-nix": "treefmt-nix"
      }
    },
    "rust-analyzer-src": {
@@ -269,29 +239,6 @@
        "repo": "treefmt-nix",
        "type": "github"
      }
-    },
-    "uv2nix": {
-      "inputs": {
-        "nixpkgs": [
-          "nixpkgs"
-        ],
-        "pyproject-nix": [
-          "pyproject-nix"
-        ]
-      },
-      "locked": {
-        "lastModified": 1767701098,
-        "narHash": "sha256-CJhKZnWb3gumR9oTRjFvCg/6lYTGbZRU7xtvcyWIRwU=",
-        "owner": "pyproject-nix",
-        "repo": "uv2nix",
-        "rev": "9d357f0d2ce6f5f35ec7959d7e704452352eb4da",
-        "type": "github"
-      },
-      "original": {
-        "owner": "pyproject-nix",
-        "repo": "uv2nix",
-        "type": "github"
-      }
    }
  },
  "root": "root",
--- a/flake.nix
+++ b/flake.nix
@@ -24,26 +24,6 @@
    dream2nix = {
      url = "github:nix-community/dream2nix";
      inputs.nixpkgs.follows = "nixpkgs";
-      inputs.pyproject-nix.follows = "pyproject-nix";
-    };
-
-    # Python packaging with uv2nix
-    pyproject-nix = {
-      url = "github:pyproject-nix/pyproject.nix";
-      inputs.nixpkgs.follows = "nixpkgs";
-    };
-
-    uv2nix = {
-      url = "github:pyproject-nix/uv2nix";
-      inputs.pyproject-nix.follows = "pyproject-nix";
-      inputs.nixpkgs.follows = "nixpkgs";
-    };
-
-    pyproject-build-systems = {
-      url = "github:pyproject-nix/build-system-pkgs";
-      inputs.pyproject-nix.follows = "pyproject-nix";
-      inputs.uv2nix.follows = "uv2nix";
-      inputs.nixpkgs.follows = "nixpkgs";
    };

    # Pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
@@ -68,7 +48,6 @@
        inputs.treefmt-nix.flakeModule
        ./dashboard/parts.nix
        ./rust/parts.nix
-        ./python/parts.nix
      ];

      perSystem =
@@ -102,9 +81,16 @@
            };
          };

+          checks.lint = pkgs.runCommand "lint-check" { } ''
+            export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
+            ${pkgs.ruff}/bin/ruff check ${inputs.self}/
+            touch $out
+          '';
+
          packages =
            if pkgs.stdenv.isDarwin then {
              metal = pkgs.callPackage ./nix/metalWrapper.nix { metalVersion = "310"; };
+              mlx = pkgs.callPackage ./nix/mlx.nix {};
            } else { };

          devShells.default = with pkgs; pkgs.mkShell {
--- a/nix/mlx.nix
+++ b/nix/mlx.nix
@@ -1,7 +1,9 @@
 { stdenv
 , lib
-, buildPythonPackage
 , fetchFromGitHub
+, fetchPypi
+, pyprojectHook
+, pypaInstallHook
 , replaceVars
 , fetchzip
 , setuptools
@@ -32,7 +34,7 @@ let
    hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
  };

-  mlx = buildPythonPackage rec {
+  mlx = stdenv.mkDerivation rec {
    pname = "mlx";
    version = "0.30.1";
    pyproject = true;
@@ -98,6 +100,8 @@ let
    nativeBuildInputs = [
      cmake
      metal
+      pyprojectHook
+      pypaInstallHook
    ];

    buildInputs = [
--- a/python/parts.nix
+++ b/python/parts.nix
@@ -1,143 +0,0 @@
-{ inputs, ... }:
-{
-  perSystem =
-    { self', pkgs, lib, system, ... }:
-    let
-      # Load workspace from uv.lock
-      workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
-        workspaceRoot = inputs.self;
-      };
-
-      # Create overlay from workspace
-      overlay = workspace.mkPyprojectOverlay { };
-
-      # Override overlay to inject Nix-built components
-      exoOverlay = final: _: {
-        # Replace workspace exo_pyo3_bindings with Nix-built wheel
-        exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
-          pname = "exo-pyo3-bindings";
-          version = "0.1.0";
-          src = self'.packages.exo_pyo3_bindings;
-          # Install from pre-built wheel
-          nativeBuildInputs = [ final.pyprojectWheelHook ];
-          dontStrip = true;
-        };
-      };
-
-      python = pkgs.python313;
-
-      # Overlay to provide build systems for source builds
-      buildSystemsOverlay = final: prev: {
-        # mlx-lm is a git dependency that needs setuptools
-        mlx-lm = prev.mlx-lm.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-
-        # Build MLX from source with proper dependencies
-        mlx = pkgs.callPythonPackage ./nix/mlx.nix;
-      };
-
-      pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
-        inherit python;
-      }).overrideScope (
-        lib.composeManyExtensions [
-          inputs.pyproject-build-systems.overlays.default
-          overlay
-          exoOverlay
-          buildSystemsOverlay
-        ]
-      );
-      exoVenv = pythonSet.mkVirtualEnv "exo-env" workspace.deps.default;
-
-      # Virtual environment with dev dependencies for testing
-      testVenv = pythonSet.mkVirtualEnv "exo-test-env" (
-        workspace.deps.default // {
-          exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
-        }
-      );
-
-      exoPackage = pkgs.runCommand "exo"
-        {
-          nativeBuildInputs = [ pkgs.makeWrapper ];
-        }
-        ''
-          mkdir -p $out/bin
-
-          # Create wrapper scripts
-          for script in exo exo-master exo-worker; do
-            makeWrapper ${exoVenv}/bin/$script $out/bin/$script \
-              --set DASHBOARD_DIR ${self'.packages.dashboard}
-          done
-        '';
-
-      pyinstallerPackage =
-        let
-          venv = pythonSet.mkVirtualEnv "exo-pyinstaller-env" (
-            workspace.deps.default
-            // {
-              # Include pyinstaller in the environment
-              exo = [ "dev" ];
-            }
-          );
-        in
-        pkgs.stdenv.mkDerivation {
-          pname = "exo-pyinstaller";
-          version = "0.3.0";
-
-          src = inputs.self;
-
-          nativeBuildInputs = [ venv pkgs.makeWrapper pkgs.macmon pkgs.darwin.system_cmds ];
-
-          buildPhase = ''
-            # macmon must be in PATH for PyInstaller to bundle it
-            export PATH="${pkgs.macmon}/bin:$PATH"
-            # HOME must be writable for PyInstaller's cache
-            export HOME="$TMPDIR"
-
-            # Copy dashboard to expected location
-            mkdir -p dashboard/build
-            cp -r ${self'.packages.dashboard}/* dashboard/build/
-
-            # Run PyInstaller
-            ${venv}/bin/python -m PyInstaller packaging/pyinstaller/exo.spec
-          '';
-
-          installPhase = ''
-            cp -r dist/exo $out
-          '';
-        };
-    in
-    {
-      # Python package only available on macOS for now due to the dependency on
-      # mlx/mlx-cpu being tricky to build on Linux. We can either remove this
-      # dependency in the PyProject or build it with Nix.
-      packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
-        exo = exoPackage;
-        exo-pyinstaller = pyinstallerPackage;
-      };
-
-      checks = {
-        # Ruff linting (works on all platforms)
-        lint = pkgs.runCommand "ruff-lint" { } ''
-          export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
-          ${pkgs.ruff}/bin/ruff check ${inputs.self}/
-          touch $out
-        '';
-      }
-      # Pytest only on macOS (requires MLX)
-      // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
-        pytest = pkgs.runCommand "pytest"
-          {
-            nativeBuildInputs = [ testVenv ];
-          } ''
-          export HOME="$TMPDIR"
-          export EXO_TESTS=1
-          cd ${inputs.self}
-          ${testVenv}/bin/python -m pytest src -m "not slow" --import-mode=importlib
-          touch $out
-        '';
-      };
-    };
-}
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/worker/download/download_utils.py
@@ -5,6 +5,7 @@ import shutil
 import ssl
 import time
 import traceback
+from collections.abc import Awaitable
 from datetime import timedelta
 from pathlib import Path
 from typing import Callable, Literal
@@ -525,7 +526,7 @@ async def download_progress_for_local_path(

 async def download_shard(
    shard: ShardMetadata,
-    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], None],
+    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    max_parallel_downloads: int = 8,
    skip_download: bool = False,
    allow_patterns: list[str] | None = None,
@@ -566,9 +567,9 @@ async def download_shard(
    )
    file_progress: dict[str, RepoFileDownloadProgress] = {}

-    def on_progress_wrapper(
+    async def on_progress_wrapper(
        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
-    ):
+    ) -> None:
        start_time = (
            file_progress[file.path].start_time
            if file.path in file_progress
@@ -604,7 +605,7 @@ async def download_shard(
            else "in_progress",
            start_time=start_time,
        )
-        on_progress(
+        await on_progress(
            shard,
            calculate_repo_progress(
                shard,
@@ -632,14 +633,21 @@ async def download_shard(

    semaphore = asyncio.Semaphore(max_parallel_downloads)

-    async def download_with_semaphore(file: FileListEntry):
+    def schedule_progress(
+        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
+    ) -> None:
+        asyncio.create_task(
+            on_progress_wrapper(file, curr_bytes, total_bytes, is_renamed)
+        )
+
+    async def download_with_semaphore(file: FileListEntry) -> None:
        async with semaphore:
            await download_file_with_retry(
                str(shard.model_meta.model_id),
                revision,
                file.path,
                target_dir,
-                lambda curr_bytes, total_bytes, is_renamed: on_progress_wrapper(
+                lambda curr_bytes, total_bytes, is_renamed: schedule_progress(
                    file, curr_bytes, total_bytes, is_renamed
                ),
            )
@@ -651,7 +659,7 @@ async def download_shard(
    final_repo_progress = calculate_repo_progress(
        shard, str(shard.model_meta.model_id), revision, file_progress, all_start_time
    )
-    on_progress(shard, final_repo_progress)
+    await on_progress(shard, final_repo_progress)
    if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):
        return target_dir / gguf.path, final_repo_progress
    else:
--- a/src/exo/worker/download/impl_shard_downloader.py
+++ b/src/exo/worker/download/impl_shard_downloader.py
@@ -1,4 +1,5 @@
 import asyncio
+from collections.abc import Awaitable
 from pathlib import Path
 from typing import AsyncIterator, Callable

@@ -48,7 +49,8 @@ class SingletonShardDownloader(ShardDownloader):
        self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}

    def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    ) -> None:
        self.shard_downloader.on_progress(callback)

@@ -83,7 +85,8 @@ class CachedShardDownloader(ShardDownloader):
        self.cache: dict[tuple[str, ShardMetadata], Path] = {}

    def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    ) -> None:
        self.shard_downloader.on_progress(callback)

@@ -113,17 +116,18 @@ class ResumableShardDownloader(ShardDownloader):
    def __init__(self, max_parallel_downloads: int = 8):
        self.max_parallel_downloads = max_parallel_downloads
        self.on_progress_callbacks: list[
-            Callable[[ShardMetadata, RepoDownloadProgress], None]
+            Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
        ] = []

-    def on_progress_wrapper(
+    async def on_progress_wrapper(
        self, shard: ShardMetadata, progress: RepoDownloadProgress
    ) -> None:
        for callback in self.on_progress_callbacks:
-            callback(shard, progress)
+            await callback(shard, progress)

    def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    ) -> None:
        self.on_progress_callbacks.append(callback)

--- a/src/exo/worker/download/shard_downloader.py
+++ b/src/exo/worker/download/shard_downloader.py
@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from collections.abc import Awaitable
 from copy import copy
 from datetime import timedelta
 from pathlib import Path
@@ -31,7 +32,8 @@ class ShardDownloader(ABC):

    @abstractmethod
    def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    ) -> None:
        pass

@@ -59,7 +61,8 @@ class NoopShardDownloader(ShardDownloader):
        return Path("/tmp/noop_shard")

    def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
    ) -> None:
        pass

--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -359,8 +359,7 @@ class Worker:
        last_progress_time = 0.0
        throttle_interval_secs = 1.0

-        # TODO: i hate callbacks
-        def download_progress_callback(
+        async def download_progress_callback(
            shard: ShardMetadata, progress: RepoDownloadProgress
        ) -> None:
            nonlocal self
@@ -372,11 +371,10 @@ class Worker:
                    total_bytes=progress.total_bytes,
                )
                self.download_status[shard.model_meta.model_id] = status
-                # Footgun!
-                self.event_sender.send_nowait(
+                await self.event_sender.send(
                    NodeDownloadProgress(download_progress=status)
                )
-                self.event_sender.send_nowait(
+                await self.event_sender.send(
                    TaskStatusUpdated(
                        task_id=task.task_id, task_status=TaskStatus.Complete
                    )
@@ -393,7 +391,7 @@ class Worker:
                    ),
                )
                self.download_status[shard.model_meta.model_id] = status
-                self.event_sender.send_nowait(
+                await self.event_sender.send(
                    NodeDownloadProgress(download_progress=status)
                )
                last_progress_time = current_time()
Author	SHA1	Message	Date
Evan	c378b1f138	lets try the other thing	2026-01-19 16:55:54 +00:00
Evan	1c06c1da8b	nix mlx compilation for better portability	2026-01-19 16:51:29 +00:00
Andrei Onel	ce3ad391b1	Update README.md with some changes from release 1.0.61 (#1157 ) Updated README.md with documentation for four new features: - added a "Benchmarking" section documenting the exo-bench tool for measuring model performance across different placement configurations - documented the custom namespace feature for cluster isolation in the macOS app section - added a "Configuration Options" subsection explaining the --no-worker CLI flag for coordinator-only nodes - added a "File Locations (Linux)" subsection documenting XDG Base Directory Specification compliance on Linux systems Issue #930	2026-01-19 16:43:18 +00:00
Jake Hillion	fb0151630d	shard_downloader: make on_progress callback async The on_progress callback was synchronous but always invoked from async contexts, forcing the use of send_nowait() which could raise WouldBlock if the channel buffer was full, potentially dropping progress updates. Changed the callback type from `Callable[[ShardMetadata, RepoDownloadProgress], None]` to return a coroutine, updated all implementations to be async, and replaced send_nowait() with await send() in the worker's download progress handler. This allows proper backpressure handling when sending download progress events through the channel, eliminating the "Footgun!" that was previously documented in the code. Test plan: - Built a DMG and ran it on one node. All existing models showed as downloaded. - Downloaded a new model. The progress bar on the download page worked. - Downloaded another new model. The progress bar on the home page worked.	2026-01-19 16:19:37 +00:00